# DiffPO: A causal diffusion model for learning distributions of potential outcomes

arXiv: https://arxiv.org/abs/2410.08924 · HTML: https://mlsanity.com/paper/2410.08924/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 42 of its 67 references that are in the map.
- First posted on arXiv 11 Oct 2024.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 42)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation](https://mlsanity.com/paper/2107.03502/) (2021)
- [Fast Sampling of Diffusion Models with Exponential Integrator](https://mlsanity.com/paper/2204.13902/) (2022)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Learning Representations for Counterfactual Inference](https://mlsanity.com/paper/1605.03661/) (2016)
- [Diffusion Models are Minimax Optimal Distribution Estimators](https://mlsanity.com/paper/2303.01861/) (2023)
- [Bayesian Inference of Individualized Treatment Effects using Multi-task Gaussian Processes](https://mlsanity.com/paper/1704.02801/) (2017)
- [Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations](https://mlsanity.com/paper/2002.04083/) (2020)
- [Orthogonal Statistical Learning](https://mlsanity.com/paper/1901.09036/) (2019)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [GENIE: Higher-Order Denoising Diffusion Solvers](https://mlsanity.com/paper/2210.05475/) (2022)
- [Causal Autoregressive Flows](https://mlsanity.com/paper/2011.02268/) (2020)
- [Diffusion Causal Models for Counterfactual Estimation](https://mlsanity.com/paper/2202.10166/) (2022)
- [Estimating the Effects of Continuous-valued Interventions using Generative Adversarial Networks](https://mlsanity.com/paper/2002.12326/) (2020)
- [Double/Debiased/Neyman Machine Learning of Treatment Effects](https://mlsanity.com/paper/1701.08687/) (2017)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Generalization Bounds and Representation Learning for Estimation of Potential Outcomes and Causal Effects](https://mlsanity.com/paper/2001.07426/) (2020)
- [Diffusion Models for Causal Discovery via Topological Ordering](https://mlsanity.com/paper/2210.06201/) (2022)
- [Treatment effect estimation with disentangled latent factors](https://mlsanity.com/paper/2001.10652/) (2020)
- [Counterfactual Representation Learning with Balancing Weights](https://mlsanity.com/paper/2010.12618/) (2020)
- [Learning Weighted Representations for Generalization Across Designs](https://mlsanity.com/paper/1802.08598/) (2018)

## Built on by (most cited first, 17 of 17)

- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [Foundation Models for Causal Inference via Prior-Data Fitted Networks](https://mlsanity.com/paper/2506.10914/) (2025)
- [LightningRL: Breaking the Accuracy-Parallelism Trade-off of Block-wise dLLMs via Reinforcement Learning](https://mlsanity.com/paper/2603.13319/) (2026)
- [Counterfactual Probabilistic Diffusion with Expert Models](https://mlsanity.com/paper/2508.13355/) (2025)
- [Improving the Generation and Evaluation of Synthetic Data for Downstream Medical Causal Inference](https://mlsanity.com/paper/2510.18768/) (2025)
- [Flow-based Generative Modeling of Potential Outcomes and Counterfactuals](https://mlsanity.com/paper/2505.16051/) (2025)
- [A Diffusion-Based Method for Learning the Multi-Outcome Distribution of Medical Treatments](https://mlsanity.com/paper/2506.01533/) (2025)
- [CausalDiffTab: Mixed-Type Causal-Aware Diffusion for Tabular Data Generation](https://mlsanity.com/paper/2506.14206/) (2025)
- [PrivATE: Differentially Private Confidence Intervals for Average Treatment Effects](https://mlsanity.com/paper/2505.21641/) (2025)
- [LLM-Driven Treatment Effect Estimation Under Inference Time Text Confounding](https://mlsanity.com/paper/2507.02843/) (2025)
- [SurvDiff: A Diffusion Model for Generating Synthetic Data in Survival Analysis](https://mlsanity.com/paper/2509.22352/) (2025)
- [GDR-learners: Orthogonal Learning of Generative Models for Potential Outcomes](https://mlsanity.com/paper/2509.22953/) (2025)
- [Retrospective Counterfactual Prediction by Conditioning on the Factual Outcome: A Cross-World Approach](https://mlsanity.com/paper/2603.27320/) (2026)
- [A Generative Framework for Causal Estimation via Importance-Weighted Diffusion Distillation](https://mlsanity.com/paper/2505.11444/) (2025)
- [Interpreting the Synchronization Gap: The Hidden Mechanism Inside Diffusion Transformers](https://mlsanity.com/paper/2603.20987/) (2026)
- [RepFlow: Representation Enhanced Flow Matching for Causal Effect Estimation](https://mlsanity.com/paper/2605.05890/) (2026)
- [Clinically-Grounded Counterfactual Reasoning for Medical Video Diagnosis](https://mlsanity.com/paper/2605.26483/) (2026)

## Abstract (arXiv)

Predicting potential outcomes of interventions from observational data is crucial for decision-making in medicine, but the task is challenging due to the fundamental problem of causal inference. Existing methods are largely limited to point estimates of potential outcomes with no uncertain quantification; thus, the full information about the distributions of potential outcomes is typically ignored. In this paper, we propose a novel causal diffusion model called DiffPO, which is carefully designed for reliable inferences in medicine by learning the distribution of potential outcomes. In our DiffPO, we leverage a tailored conditional denoising diffusion model to learn complex distributions, where we address the selection bias through a novel orthogonal diffusion loss. Another strength of our DiffPO method is that it is highly flexible (e.g., it can also be used to estimate different causal quantities such as CATE). Across a wide range of experiments, we show that our method achieves state-of-the-art performance.