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

要約

観察データから個別化された治療効果を推定することは、主に非ランダム化治療の割り当てからの共変量の不均衡と交絡バイアスによる因果推論の中心的な課題です。
逆確率重み付け(IPW)はこの問題に対する確立されたソリューションですが、最新の深い学習フレームワークへの統合は限られたままです。
この作業では、拡散モデルの事前抑制と重要な加重スコア蒸留を組み合わせて、潜在的な迅速な因果推定を潜在的な結果の予測と治療効果の推定を含む可能性を可能にする、重要な加重拡散蒸留(IWDD)を提案します。
IPWを自然に前処理された拡散モデルの蒸留に組み込む方法を示し、IPWを明示的に計算する必要性を排除するランダム化ベースの調整をさらに導入し、さらに重要なことに、勾配推定値の分散を確実に削減することを実証します。
経験的な結果は、IWDDが他のベースラインと比較して最高の勝利率で最先端のサンプル外予測パフォーマンスを達成し、因果推定を大幅に改善し、個別の治療戦略の開発をサポートすることを示しています。
再現性と将来の研究のためにPytorchコードをリリースします。

要約(オリジナル)

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.

arxiv情報

著者 Xinran Song,Tianyu Chen,Mingyuan Zhou
発行日 2025-05-16 17:00:52+00:00
arxivサイト arxiv_id(pdf)

提供元, 利用サービス

arxiv.jp, Google

カテゴリー: cs.LG, stat.AP, stat.ME, stat.ML パーマリンク