arXiv:2505.11444cs.LGstat.AP2025-05被引 1

用扩散模型生成因果估计,提升治疗效果预测精度。

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

  • 将逆概率加权融入扩散模型蒸馏,实现高效因果推断。
  • 在多个数据集上优于基线,出样本预测胜率最高。
  • 无需显式计算权重,降低梯度方差,适合个性化医疗研究。

从观测数据中估计个体化治疗效应是因果推断的核心挑战,主要源于协变量不平衡和非随机治疗分配带来的混淆偏差。尽管逆概率加权(IPW)是解决该问题的经典方法,但其在现代深度学习框架中的集成仍受限。本文提出重要性加权扩散蒸馏(IWDD),一种结合预训练扩散模型与重要性加权得分蒸馏的生成式框架,可实现快速、准确的因果估计——包括潜在结果预测与治疗效应估计。我们展示了如何将IPW自然融入预训练扩散模型的蒸馏过程,并引入基于随机化的调整机制,避免显式计算IPW,从而简化计算并严格降低梯度估计方差。实验表明,IWDD在多个数据集上实现最优的出样本预测性能,胜率显著高于其他基线,显著提升因果估计能力,支持个性化治疗策略发展。代码将开源以供复现与后续研究。

原文摘要 · Abstract (English)

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.

因果推断扩散模型治疗效应生成模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。