Bayesian X-Learner: Calibrated Posterior Inference for Heterogeneous Treatment Effects under Heavy-Tailed Outcomes

Eichi Uehara による論文「Bayesian X-Learner: Calibrated Posterior Inference for Heterogeneous Treatment Effects under Heavy-Tailed Outcomes」の詳細な技術的サマリーを以下に示す。

1. 問題定義

現実世界の応用における条件付き平均処置効果（CATE）、すなわち $\tau(x)$ の推定は、既存のツールが同時に満たすことができない重要な 3 つの要件に直面している：

1. 異質性: 個人間で変化する処置効果（$\tau(x)$）を推定する能力。
2. 較正された不確実性: これらの異質な効果に対する統計的に妥当な不確実性の定量化（信頼区間）の提供。
3. 重尾頑健性: 臨床試験（稀な重篤な事象）、マーケティング（収益の「ホエール」）、経済学などの分野で一般的である「重尾」を持つ結果分布に対する頑健性。

既存手法の限界:

• メタラーナー（S/T/X-Learners）: 異質性を提供するが、通常は較正された不確実性を持たない点推定値のみを出力する。
• Causal Forests: 異質性を提供し、ブートストラップによる近似不確実性を提供するが、ブートストラップ区間は重尾汚染下でカバレッジを維持できないことが多い。
• 古典的ベイズモデル（例：Causal BART, BCF）: 完全なベイズ事後分布を提供するが、ガウス尤度仮定に依存している。単一の極端な外れ値（「ホエール」）は、二乗残差の無制限な影響により、事後平均を歪め、分散を過大評価する。
• 頑健な頻度論的手法: 点推定値または頻度論的区間を提供することが多いが、複雑な下流の意思決定（例：不確実性への統合）に必要な完全な事後オブジェクトを欠いている。

2. 手法：ベイズ X-Learner

本論文は、結果が重度に汚染されている場合でも $\tau(x)$ に対して較正された MCMC 事後分布を提供するように設計された 3 フェーズのフレームワークであるベイズ X-Learnerを提案する。

フェーズ 1：クロスフィットによるノイズ隔離

• 目的: 擬似アウトカムへの汚染の伝播を防ぐために、ノイズ関数（$\hat{\mu}_0, \hat{\mu}_1, \hat{\pi}$）を頑健に推定する。
• メカニズム: $K$ フォールド・クロスフィットを使用する。
• 頑健性: 標準的な二乗誤差損失の代わりに、Huber のミニマックス理論に基づき、contamination_severity 列挙値（例："severe" $\to$ $\delta=0.5$）からマッピングされたパラメータ $\delta$ を持つCatBoost と Huber 損失を採用する。
• 標準化: 極端なスケールのデータの場合、オプションの前標準化ステップ（$Y \to Y/\text{MAD}(Y)$）により、Huber の $\delta$ パラメータが単位スケールで動作し、その理論的最適性が維持されるようにする。

フェーズ 2：ダブルロバスト（DR）擬似アウトカム

• 目的: ノイズモデルが不完全であっても $\tau(x)$ に対して整合性を持つ目標変数を構築する。
• メカニズム: Kennedy (2020) に従い、各単位に対して DR 擬似アウトカム（$D_i$）を構築する：
  $$D_i = \hat{\mu}_1(X_i) - \hat{\mu}_0(X_i) + \frac{W_i(Y_i - \hat{\mu}_1(X_i))}{\hat{\pi}(X_i)} - \frac{(1-W_i)(Y_i - \hat{\mu}_0(X_i))}{1-\hat{\pi}(X_i)}$$
• これらの擬似アウトカムは、最終的なベイズ層の回帰ターゲットとして機能する。

フェーズ 3：Welsch 擬似尤度によるベイズ更新

• 目的: CATE パラメータ $\beta$（ここで $\tau(x) = \phi(x)^\top \beta$）の事後分布を推論する。
• 革新: 標準的なガウス尤度をWelsch 再降下擬似尤度に置き換える。
  • 損失関数: $\rho_W(r) = \frac{c^2}{2}(1 - \exp(-r^2/c^2))$。
  • 影響関数: $\psi_W(r) = r \exp(-r^2/c^2)$。
  • 主要な特性: 影響関数は有界であり、大きな残差に対してゼロに再降下する。これは、極端な外れ値が事後平均に与える影響が減衰することを意味し、影響が依然として高い、またはゆっくりと減衰するガウスまたは学生 t 尤度とは対照的である。
• 推論: NumPyro を介した NUTS（No-U-Turn Sampler）を使用して、一般化（Gibbs）事後分布からサンプリングする。
• 較正: Welsch 損失は真の対数尤度ではないため、事後分布は「一般化事後分布」である。論文は、信頼区間が正しい頻度論的カバレッジを持つようにするための学習率 $\eta$ の較正（トレースマッチングまたは機能固有のマッチングを介して）を導入する。
• モジュラー・ベイズ: ノイズ推定における不確実性を考慮するため、本手法はノイズパラメータをベイズ・ブートストラップで抽出し、その不確実性を最終的な事後分布に伝播させるモジュラー・ベイズ・アプローチをサポートする。

3. 主要な区別：汚染対シグナル

重要な理論的貢献は、2 つの重尾レジーム間の操作的な区別にある：

1. 汚染としてのテール: 外れ値はノイズ（例：測定誤差、収益のホエール）である。解決策は、Welsch 尤度と Huber ノイズによる重み付け低下である。
2. シグナルとしてのテール: 特定のサブグループが真に異なる処置効果を持つ場合。解決策は、テール感知基底関数（例：$\phi(x) = [1, \mathbb{I}(|x| > c)]$）による表現である。

• 警告: 本論文は、「シグナルとしてのテール」を処理するためにデータ層のリスケーリング（ヒル推定量正規化など）を使用することは、真の異質性をノイズとして扱うため、積極的に有害であることを実証している。

4. 主要な貢献

1. アルゴリズム: DR 擬似アウトカムと Welsch 再降下尤度を組み合わせた新規のベイズ X-Learner。これは、異質的で、較正され、かつ頑健なCATE 推定のための、初めて広く適用可能なツールを提供する。
2. 理論的枠組み: 汚染処理（尤度層）対シグナル処理（基底層）の分離を形式化し、一般化事後分布に対する漸近的な較正結果（命題 1）を提供する。
3. 実証的検証:
   • IHDP ベンチマーク: クリーンなデータにおいて、デフォルト設定は平均 $\sqrt{\epsilon_{PEHE}} = 0.56$ を達成し、S/T/X-ラーナー、Causal BART、Causal Forests を上回るか同等の性能を示し、最も狭い分散を有する。
   • ホエール DGP: 最大 20〜25% の汚染を持つ合成データにおいて、「severe」設定は RMSE $\approx 0.13$ を回復し、較正された 95% 信頼区間を維持する。一方、ガウスおよび学生 t ベースラインは壊滅的に失敗する。
   • 実世界データ: Hillstrom メールマーケティング RCT（$N=42,613$）で検証され、ナイーブな推定量は外れ値に駆動され（$0.77 のリフトを示唆）、頑健なベイズ X-Learner は顧客の大部分に対してゼロに近い効果を正しく特定し、狭く較正された区間を示す。
4. 境界分析: 汚染天井、効率コスト、基底感度、サンプルサイズ、処置タイプ、計算の 6 つの実証的境界をマッピングし、このツールを代替手段に対していつ使用すべきかについて実務家に指針を与える。

5. 結果の概要

• 頑健性: 本手法は、標準的なベイズおよび頻度論的手法が崩壊する 20% の汚染密度下で、妥当なカバレッジ（設定に応じて約 83〜100%）を維持する。
• 効率コスト: 攻撃的な頑健設定（例：$\delta=0.5$）を使用すると、完全にクリーンなデータにおいて構造的な効率コストが生じるが、汚染が存在しない場合に「none」（ガウス）に切り替えることで緩和される。
• 較正: ノイズ不確実性を伝播させるモジュラー・ベイズ・プーリングの使用は、単一クロスフィット事後分布がカバレッジ不足を示す高汚染シナリオにおいて、名目上のカバレッジを回復するために重要であることが示された。
• 意思決定の質: ポリシーシミュレーションにおいて、較正された事後分布は汚染下で処置決定におけるゼロの後悔をもたらすのに対し、非頑健な手法は重大なポリシエラーをもたらす。

6. 意義

本論文は、因果機械学習における重要なギャップに取り組んでいる。頑健な統計学とベイズ推論は長らく別々に存在してきたが、異質な処置効果という文脈におけるそれらの交差点は以前は探求されていなかった。

• 実用的影響: 臨床試験やマーケティングなどの分野の実務家が、データに外れ値が散りばめられていても、サブグループ効果に基づき、信頼できる不確実性の範囲で意思決定を行えるようにする。
• 理論的進展: 点推定の頑健性を超えて事後分布の頑健性へと進歩し、再降下損失を用いた一般化ベイズが較正された推論をもたらすことを証明する。
• 指針: 重尾における「ノイズ」と「シグナル」を区別するための明確で実行可能な指針を提供し、真の異質な効果を平滑化してしまうという一般的な落とし穴を防ぐ。

コードとベンチマークはオープンソースとしてリリースされており、実世界の重尾因果推論タスクにおける再現性と採用を促進する。