Research map / Probabilistic and scientific modeling / Treatment effects
DiffPO: A causal diffusion model for learning distributions of potential outcomes
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DiffPO, a conditional diffusion model with an inverse-propensity-weighted orthogonal loss that learns full distributions of potential outcomes from observational data, evaluated on synthetic, ACIC, and IHDP data.
Problem
Predicting potential outcome distributions (not just point estimates) from observational data under selection bias, where only one potential outcome is observed per individual.
Approach
Uses a tailored conditional denoising diffusion model to learn p(Y|X,A) and adds an inverse-propensity-weighted diffusion loss (with a learned propensity network) that is Neyman-orthogonal, so the loss approximates the potential-outcome distributions; the model can also yield CATE as the difference of predicted outcome means.
Key results
- On the synthetic dataset, DiffPO achieves out-of-sample empirical W1 of 0.125 (a=0) and 0.091 (a=1), versus 0.943 and 1.005 for the best and worst baselines respectively in the table.
- DiffPO's 95% predictive interval empirical coverage is 0.981 for a=0 and 0.908 for a=1, while MC-dropout baselines reach roughly 0.80-0.89.
- On ACIC 2018 out-of-sample, DiffPO is best on 36.86% of datasets versus 22.13% for the next-best DR-learner (CATE estimation).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DiffPO instead of the strongest reported baseline | Synthetic potential-outcome distribution prediction · Synthetic, a=0 and a=1, in- and out-of-sample | Wasserstein distance: 0.921 ± 0.112 -> 0.043 ± 0.021; 0.943 ± 0.145 -> 0.125 ± 0.051; 0.909 ± 0.112 -> 0.032 ± 0.037; 1.078 ± 0.192 -> 0.091 ± | Helped |
| Use DiffPO instead of the strongest reported baseline | Synthetic potential-outcome predictive intervals · Synthetic, a=0 and a=1, 95% and 99% PI | Empirical coverage: 0.892 -> 0.981; 0.954 -> 0.985; 0.873 -> 0.908; 0.885 -> 0.926 | Helped |
| Use DiffPO instead of the strongest reported baseline | Synthetic potential-outcome point prediction · Synthetic, a=0 and a=1, in- and out-of-sample | RMSE: 0.263 -> 0.143; 0.278 -> 0.156; 0.261 -> 0.162; 0.293 -> 0.187 | Helped |
| Use DiffPO instead of the strongest reported baseline | CATE estimation · ACIC 2016 and ACIC 2018, in- and out-of-sample | Runs with best performance (%): 15.71 -> 26.73; 15.67 -> 25.82; 23.64 -> 31.34; 22.13 -> 36.86 | Helped |
Surprising findings
- Point-estimate RMSE of the flexible-distribution approach is compared to MC-dropout baselines; the authors note MC dropout intervals are not faithful.
Limitations
- Sampling process is slow compared with one-pass baselines
- Relies on standard causal assumptions (consistency, unconfoundedness, overlap)
- Evaluation largely on synthetic and semi-synthetic data since counterfactuals are unobserved
Experimental setup
Datasets: synthetic dataset ACIC 2016 ACIC 2018 IHDP
Benchmarks: synthetic dataset ACIC 2016 ACIC 2018 IHDP
Metrics: empirical Wasserstein distance (k=1) predictive interval coverage (95%, 99%) RMSE PEHE (% best)
Baselines: S-learner T-learner TARNet CFR GANITE DR-learner RA-learner TEDVAE
Builds on: denoising diffusion probabilistic models CSDI DR-learner neyman-orthogonal learning
Closest papers by meaning
- Reducing Confounding Bias without Data Splitting for Causal Inference via Optimal Transport (ICML 2025 · similarity 0.87)
- DiffPuter: Empowering Diffusion Models for Missing Data Imputation (ICLR 2025 · similarity 0.86)
- Learning Mixtures of Unknown Causal Interventions (NeurIPS 2024 · similarity 0.86)
- When Additive Noise Meets Unobserved Mediators: Bivariate Denoising Diffusion for Causal Discovery (NeurIPS 2025 · similarity 0.86)
- Marginal Causal Flows for Validation and Inference (NeurIPS 2024 · similarity 0.86)
- Orthogonal Survival Learners for Estimating Heterogeneous Treatment Effects from Time-to-Event Data (NeurIPS 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.