# A Generative Framework for Causal Estimation via Importance-Weighted Diffusion Distillation

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 30 of its 50 references that are in the map.
- First posted on arXiv 16 May 2025.

## Placed in

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

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [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)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [ProlificDreamer: High-Fidelity and Diverse Text-to-3D Generation with Variational Score Distillation](https://mlsanity.com/paper/2305.16213/) (2023)
- [One-step Diffusion with Distribution Matching Distillation](https://mlsanity.com/paper/2311.18828/) (2023)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Diff-Instruct: A Universal Approach for Transferring Knowledge From Pre-trained Diffusion Models](https://mlsanity.com/paper/2305.18455/) (2023)
- [Score identity Distillation: Exponentially Fast Distillation of Pretrained Diffusion Models for One-Step Generation](https://mlsanity.com/paper/2404.04057/) (2024)
- [Diffusion-GAN: Training GANs with Diffusion](https://mlsanity.com/paper/2206.02262/) (2022)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [CARD: Classification and Regression Diffusion Models](https://mlsanity.com/paper/2206.07275/) (2022)
- [Diffusion Causal Models for Counterfactual Estimation](https://mlsanity.com/paper/2202.10166/) (2022)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [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)
- [Semi-Implicit Variational Inference](https://mlsanity.com/paper/1805.11183/) (2018)
- [DeepMatch: Balancing Deep Covariate Representations for Causal Inference Using Adversarial Training](https://mlsanity.com/paper/1802.05664/) (2018)
- [Truncated Diffusion Probabilistic Models and Diffusion-based Adversarial Auto-Encoders](https://mlsanity.com/paper/2202.09671/) (2022)
- [General Identifiability and Achievability for Causal Representation Learning](https://mlsanity.com/paper/2310.15450/) (2023)
- [Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation](https://mlsanity.com/paper/2211.01939/) (2022)
- [Causal Modeling with Stationary Diffusions](https://mlsanity.com/paper/2310.17405/) (2023)
- [Causal Diffusion Autoencoders: Toward Counterfactual Generation via Diffusion Probabilistic Models](https://mlsanity.com/paper/2404.17735/) (2024)
- [DiffPO: A causal diffusion model for learning distributions of potential outcomes](https://mlsanity.com/paper/2410.08924/) (2024)
- [Hierarchical Semi-Implicit Variational Inference with Application to Diffusion Model Acceleration](https://mlsanity.com/paper/2310.17153/) (2023)
- [Denoising Score Distillation: From Noisy Diffusion Pretraining to One-Step High-Quality Generation](https://mlsanity.com/paper/2503.07578/) (2025)
- [Diffusion Model in Causal Inference with Unmeasured Confounders](https://mlsanity.com/paper/2308.03669/) (2023)
- [Generalizing the intention-to-treat effect of an active control against placebo from historical placebo-controlled trials to an active-controlled trial: A case study of the efficacy of daily oral TDF/FTC in the HPTN 084 study](https://mlsanity.com/paper/2304.03476/) (2023)

## Abstract (arXiv)

Estimating individualized treatment effects from observational data is a central challenge in causal inference, largely due to covariate imbalance and confounding bias from non-randomized treatment assignment. While inverse probability weighting (IPW) is a well-established solution to this problem, its integration into modern deep learning frameworks remains limited. In this work, we propose Importance-Weighted Diffusion Distillation (IWDD), a novel generative framework that combines the pretraining of diffusion models with importance-weighted score distillation to enable accurate and fast causal estimation-including potential outcome prediction and treatment effect estimation. We demonstrate how IPW can be naturally incorporated into the distillation of pretrained diffusion models, and further introduce a randomization-based adjustment that eliminates the need to compute IPW explicitly-thereby simplifying computation and, more importantly, provably reducing the variance of gradient estimates. Empirical results show that IWDD achieves state-of-the-art out-of-sample prediction performance, with the highest win rates compared to other baselines, significantly improving causal estimation and supporting the development of individualized treatment strategies. We will release our PyTorch code for reproducibility and future research.