[論文レビュー] Optimal Transport for Treatment Effect Estimation

Hao Wang, Zhichao Chen+8 more|arXiv (Cornell University)|Oct 27, 2023

Advanced Causal Inference Techniques被引用数 5

ひとこと要約

本稿では、ミニバッチサンプリング効果（MSE）と観察されない交絡要因効果（UCE）を緩和するための緩和された質量保存正則化子と近接的事実的結果正則化子を用いた、最適輸送に基づく新しい条件付き平均処置効果（CATE）推定手法であるESCFRを提案する。本手法は、確率的最適輸送を用いて全データ空間における処置群の整合を図ることで、観測データにおいて最先端の性能を達成する。

ABSTRACT

Estimating conditional average treatment effect from observational data is highly challenging due to the existence of treatment selection bias. Prevalent methods mitigate this issue by aligning distributions of different treatment groups in the latent space. However, there are two critical problems that these methods fail to address: (1) mini-batch sampling effects (MSE), which causes misalignment in non-ideal mini-batches with outcome imbalance and outliers; (2) unobserved confounder effects (UCE), which results in inaccurate discrepancy calculation due to the neglect of unobserved confounders. To tackle these problems, we propose a principled approach named Entire Space CounterFactual Regression (ESCFR), which is a new take on optimal transport in the context of causality. Specifically, based on the framework of stochastic optimal transport, we propose a relaxed mass-preserving regularizer to address the MSE issue and design a proximal factual outcome regularizer to handle the UCE issue. Extensive experiments demonstrate that our proposed ESCFR can successfully tackle the treatment selection bias and achieve significantly better performance than state-of-the-art methods.

研究の動機と目的

• 処置選択バイアスを是正すること。これは、処置群と非処置群の間で予測変数の分布シフトが生じるためである。
• ミニバッチサンプリング効果（MSE）を克服すること。これは、外れ値や不均衡なバッチが、既存手法における分布差の推定を歪めるためである。
• 観察されない交絡要因効果（UCE）を軽減すること。これは、標準的な交絡なし仮定のもとで、差の計算がバイアスを受けるためである。
• 局所的なミニバッチではなく、全データ空間上で動作する、整合的な因果推論フレームワークを構築すること。
• 観察されない交絡要因を考慮する近接正則化子を統合することで、反実仮想結果の推定を改善すること。

提案手法

• 観測データからのCATE推定を目的とした、確率的最適輸送に基づく全空間反実仮想回帰（ESCFR）フレームワークを提案する。
• 外れ値や不均衡なバッチに起因するミニバッチサンプリングのアーティファクトへの感受性を低減することで、分布整合性を安定化する緩和された質量保存正則化子を導入する。
• 観察されない交絡要因が結果に与える単調な影響をモデル化することで、それらの影響を考慮する近接的事実的結果正則化子を設計する。
• 全データ空間における分布差を推定するために、潜在的にバイアスを受けるミニバッチ推定に依存しない一般化されたSinkhorn差を採用する。
• 代表的マッピングと反実仮想結果を同時に学習するための双対最適化フレームワークを用い、分布整合性と結果の忠実性の両方を保証する。
• 観察されない交絡要因と結果の間の単調性仮定を適用することで、可能な輸送戦略の特定集合を縮小し、推定の正確性を向上させる。

実験結果

リサーチクエスチョン

• RQ1 ミニバッチサンプリング効果が分布差推定に与える影響を軽減することで、観測データにおけるCATE推定をどのように改善できるか？
• RQ2 観察されない交絡要因は、既存の最適輸送ベースのCATE推定器にどの程度バイアスをもたらすのか。そして、そのバイアスはどのように是正できるか？
• RQ3 局所的なミニバッチではなく、全データ空間上で動作する整合的な最適輸送フレームワークを設計することは可能か？
• RQ4 事実的結果に近接正則化子を統合することで、観察されない交絡要因下での反実仮想予測の正確性はどのように向上するか？
• RQ5 観察されない交絡要因と結果の間の単調性仮定を活用することで、輸送戦略選択の曖昧さを低減できるか？

主な発見

• ESCFRは、複数のベンチマークデータセットにおいて、SOTA手法を著しく上回るCATE推定性能を示し、分布シフトに対して優れたロバストネスを示した。
• 緩和された質量保存正則化子は、ミニバッチ内の外れ値や結果の不均衡に起因する影響を効果的に低減し、トレーニングプロセスの安定性と整合性の正確性を向上させた。
• 近接的事実的結果正則化子は、観察されない交絡要因に起因するバイアスを低減し、特に単調性仮定が成り立つ場合には、より正確な反実仮想予測を実現した。
• 実験結果から、ESCFRは特に高い観察されない交絡要因が存在する状況において、既存のベースラインと比較してより低い平均二乗誤差（MSE）を達成した。
• 本手法は、複雑で非線形な結果関係や隠れた交絡要因を有する多様なデータセットにおいて、一貫した性能向上を示した。
• アブレーションスタディの結果、質量保存正則化子と近接正則化子の両方が性能向上に顕著に寄与しており、それらの設計の妥当性が裏付けられた。