[論文レビュー] Exploring Transformer Backbones for Heterogeneous Treatment Effect Estimation

Yifan Zhang, Hanlin Zhang+3 more|arXiv (Cornell University)|Feb 2, 2022

Advanced Causal Inference Techniques被引用数 12

ひとこと要約

本稿では、自己注意機構を用いて治療と共変量の間の相互作用をモデル化するトランスフォーマーベースのバックボーンを用いた、新しい治療効果推定手法TransTEEを提案する。これにより、離散的・連続的・構造的治療のすべてにおいて、柔軟で頑健かつ解釈可能な異質的治療効果推定が可能になる。TransTEEは多様なベンチマークで強力なベースラインを上回り、事前学習済み言語モデルの共変量調整と監査が可能になる。

ABSTRACT

Previous works on Treatment Effect Estimation (TEE) are not in widespread use because they are predominantly theoretical, where strong parametric assumptions are made but untractable for practical application. Recent work uses multilayer perceptron (MLP) for modeling casual relationships, however, MLPs lag far behind recent advances in ML methodology, which limits their applicability and generalizability. To extend beyond the single domain formulation and towards more realistic learning scenarios, we explore model design spaces beyond MLPs, i.e., transformer backbones, which provide flexibility where attention layers govern interactions among treatments and covariates to exploit structural similarities of potential outcomes for confounding control. Through careful model design, Transformers as Treatment Effect Estimators (TransTEE) is proposed. We show empirically that TransTEE can: (1) serve as a general purpose treatment effect estimator that significantly outperforms competitive baselines in a variety of challenging TEE problems (e.g., discrete, continuous, structured, or dosage-associated treatments) and is applicable to both when covariates are tabular and when they consist of structural data (e.g., texts, graphs); (2) yield multiple advantages: compatibility with propensity score modeling, parameter efficiency, robustness to continuous treatment value distribution shifts, explainable in covariate adjustment, and real-world utility in auditing pre-trained language models

研究の動機と目的

• 既存のMLPベースの治療効果推定器の限界、すなわち一般化性能の低さ、パラメータの非効率性、分布シフトに伴う脆さを是正すること。
• 異質的治療効果推定における複雑な治療-共変量相互作用をモデル化するためのMLPの代替として、トランスフォーマーベースのバックボーンを検討すること。
• 傾向スコアモデリングと互換性があり、パラメータ効率的で、連続的治療の分布シフトに頑健な統一フレームワークの開発。
• 注目度に基づく特徴量重要度を用いた、因果推論における解釈可能な共変量調整の実現。
• 事前学習済み言語モデルの公平性および社会的影響を監査するための実用性の検証。

提案手法

• TransTEEは、すべての治療と共変量を一括して埋め込むための単一のトランスフォーマーベースのバックボーンを採用し、マルチタスクアーキテクチャを回避し、エンドツーエンド学習を可能にする。
• 自己注意機構により、治療と共変量の間の動的な相互作用をモデル化し、適応的共変量選択と交絡要因の制御を実現する。
• 注目度重みを活用して交絡要因および関連共変量を特定・優先順位付けし、条件付き無視可能性のための排他的原因基準と整合する。
• 傾向スコアモデリングと互換性があるように設計されており、観察研究における頑健性と解釈可能性を向上させる。
• 表形式および構造的入力（例：テキスト、グラフ）の両方をサポートしており、広範な適用可能性を有する。
• 合成ベンチマークを超えた実世界の実用性を評価するため、サーロゲートモデリングタスクを導入する。

実験結果

リサーチクエスチョン

• RQ1 トランスフォーマーベースのバックボーンは、離散的・連続的・構造的治療を含む多様な治療タイプにおいて、異質的治療効果を効果的にモデル化できるか？
• RQ2 TransTEEは、最先端のMLPベースのベースラインと比較して、性能、頑健性、パラメータ効率性の面でどのように優れているか？
• RQ3 TransTEEの注目度メカニズムは、因果推論における解釈可能で意味のある共変量調整をどの程度実現できるか？
• RQ4 TransTEEは、公平性および社会的影響を評価する目的で、事前学習済み言語モデルの監査に有効に利用できるか？
• RQ5 連続的治療値の分布シフトに対して、TransTEEはどの程度頑健であるか？

主な発見

• TransTEEは、IHDP、News、TCGAを含む4種類の治療タイプにわたる6つのベンチマークで、競合するベースラインを顕著に上回り、優れた一般化性能と推定性能を示した。
• 半合成および実世界のデータセットにおいて、平均治療効果（ATE）推定で最先端の性能を達成し、サーロゲートタスクでは最大ATEが0.9976に達した。
• 連続的治療値の分布シフトに対しても、TransTEEは頑健である。一方で、MLPベースのモデルは性能が著しく低下する。
• TransTEEの注目度重みは、交絡要因および関連共変量を効果的に特定・優先順位付けでき、因果同定の原則と整合する解釈可能な特徴量重要度スコアを提供する。
• TransTEEは、事前学習済み言語モデルの公平性評価に有効に利用でき、極端なATEが0.9976および0.0に達する摂動解析により、社会的バイアスを明らかにした。
• フレームワークはパラメータ効率的であり、傾向スコアモデリングと互換性があるため、実世界の因果推論の現場での実用的導入を促進する。