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

TL;DR

IWDD tackles covariate imbalance in observational causal inference by coupling diffusion pretraining with importance-weighted distillation and a randomization-based input adjustment that obviates explicit propensity estimation. The framework trains a conditional diffusion teacher and distills to a fast one-step generator using an IPW-free objective, with Fisher divergence and SiD for stable optimization and reduced gradient variance. Empirically, IWDD achieves state-of-the-art out-of-sample performance on standard causal benchmarks and delivers faster sampling compared with diffusion baselines, enabling robust individualized treatment effect estimation. The approach offers a practical, scalable path to accurate potential outcome prediction and heterogeneous treatment effect estimation in observational data, with code to support reproducibility.

Abstract

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.

Figures (3)

• Figure 1: Overview of IWDD. We first pretrain a conditional diffusion model $f_\phi(y \mid x, z)$ to approximate the true conditional distribution $p(y \mid x, z)$ over observational data. In the distillation stage, we train a generator $q_\theta(y \mid x, z)$ using marginal sampling, which implicitly applies importance weighting without requiring explicit propensity estimation. We apply a randomization-based sampling adjustment: covariates $x$ are shuffled and treatments $z$ are independently sampled. The distillation algorithm is detailed in Algorithm \ref{['alg:iwdd']}.
• Figure 2: Marginal distributions of $y$ in training and testing sets. Due to the shift in treatment assignment, the induced distribution of $y$ differs across domains.
• Figure 3: Synthetic data example: estimated potential outcome distributions $Y(0)$ and $Y(1)$ from different models. The pretrained diffusion model performs well in-sample and for $Y(0)$ out-of-sample, but struggles with $Y(1)$. IWDD improves estimation for out-of-sample $Y(1)$ while maintaining performance elsewhere.

Theorems & Definitions (4)

• Lemma 1
• proof
• Theorem 1
• proof : Proof sketch