Multi-Domain Causal Empirical Bayes Under Linear Mixing

この論文「Multi-Domain Empirical Bayes for Linearly-Mixed Causal Representations（線形混合因果表現のためのマルチドメイン経験的ベイズ）」は、高次元の観測データから低次元の因果潜在変数を学習する「因果表現学習（Causal Representation Learning: CRL）」の推定問題に焦点を当てています。特に、複数のドメイン（環境）から得られる干渉（介入）データを用いて、経験的ベイズ（Empirical Bayes: EB）の枠組みを適用し、より高精度な推定を行う手法を提案しています。

以下に、論文の技術的な詳細を要約します。

1. 問題設定と背景

• 課題: 従来の CRL の研究は、無限データ（人口レベル）における「識別可能性（identifiability）」に注力してきましたが、有限サンプルからの「推定（estimation）」については十分に探求されていませんでした。
• シナリオ: 複数のドメイン（環境）から得られるデータ。各ドメインは、基底となる因果モデルに対する「スパースな介入（sparse interventions）」によって特徴付けられます。観測データ $x$ は、ドメイン固有の因果構造を持つ潜在変数 $z$ と、ドメイン間で共有される線形混合モデル（ノイズ付き）によって生成されると仮定します。
• 目標: 有限のノイズを含む観測データから、ドメインごとの潜在因果変数 $z$ と混合行列 $A$ を同時に推定すること。これは、多数の局所潜在変数を同時に推論する「同時推論（simultaneous inference）」問題として定式化されます。

2. 提案手法：因果経験的ベイズ（Causal Empirical Bayes）

著者らは、CRL を経験的ベイズの枠組みで捉え直し、f-modeling（ marginal density を直接モデル化するアプローチ）を採用しました。

2.1 モデルの定式化

• 観測モデル: 線形ガウスモデルを仮定します。
  $$x_{ei} = A^* z_{ei} + \varepsilon_{ei}, \quad \varepsilon_{ei} \sim \mathcal{N}(0, \sigma^{*2} I)$$
  ここで、$e$ はドメイン、$i$ はサンプル、$A^*$ は混合行列です。
• 事前分布（介入構造）: 潜在変数 $z$ は、既知の因果グラフ $G^*$ と介入ターゲット $a_e$ によって定義された構造因果モデル（SCM）に従います。ドメイン $e$ における $z$ の分布 $p(z|a_e)$ は、SCM 上の介入によって変化しますが、測定プロセス（$A$ とノイズ）はドメイン間で不変であると仮定します。

2.2 推論アルゴリズム（EM 型 f-modeling）

提案手法は、Tweedie の公式とスコアマッチングを組み合わせた EM アルゴリズム（アルゴリズム 1）です。

1. 直交性の仮定と正規平均モデルへの帰着:
   混合行列 $A$ が直交列を持つ（$A^\top A = D^2$）と仮定し、観測データを投影することで、古典的な「正規平均モデル（Normal Means Model）」$y_{ei} = D z_{ei} + \zeta_{ei}$ に帰着させます。これにより、Tweedie の公式が適用可能になります。

2. E ステップ（潜在変数の推定）:

   • Tweedie の公式: 事後平均 $E[z|x]$ を、周辺分布のスコア関数 $\nabla \log f(x)$ を用いて推定します。
     $$\hat{z} \approx x + \sigma^2 \nabla \log f(x)$$
   • 因果構造付きスコアマッチング: 従来の EB ではスコア関数を非構造的に学習しますが、本手法では既知の因果グラフ $G$ を利用してスコア関数を**疎（sparse）**にパラメータ化します。
     $$[s(y, a)]_j = s_j(y_j, y_{pa(j)}, a_j)$$
     各成分 $j$ について、その親ノード $pa(j)$と介入ターゲット$a_j$ のみを入力とする関数としてスコアを学習します。これにより、高次元問題における推定の安定性と計算効率を向上させています。
   • 減衰（Damping）: 推定値の安定化のため、Tweedie 更新に減衰係数 $\eta$ を導入します。

3. M ステップ（モデルパラメータの更新）:
   E ステップで得られた潜在変数の事後モーメント（1 次および 2 次モーメント）を用いて、対数尤度を最大化するように混合行列 $A$ とノイズ分散 $\sigma^2$ を更新します。これは解析的に解くことができます。

4. 反復: 上記の E ステップと M ステップを収束するまで反復します。

3. 主要な貢献

• CRL における経験的ベイズの定式化: 多ドメイン CRL を同時推論問題として捉え、EB の f-modeling アプローチを適用する新たな枠組みを提案しました。
• 因果構造を利用したスコアマッチング: 既知の因果グラフをスコア関数の構造制約として組み込むことで、高次元データにおける推定精度と計算スケーラビリティを両立させました。
• EM 型アルゴリズムの提案: 因果構造付きスコアマッチングと Tweedie 公式、および EM 更新を統合した効率的な推定アルゴリズムを開発しました。
• g-modeling との関連性の解明: 既存の CRL 手法（VAE など）が g-modeling（事前分布を明示的にモデル化）に分類されることを示し、提案手法との対比を行いました。

4. 実験結果

合成データを用いた実験で、提案手法の性能を検証しました。

• 設定: 4 つの潜在変数を持つチェーン構造の SCM からデータを生成し、ノードへの介入をシミュレートしました。観測次元は 100、潜在次元は 4 です。
• ベースライン: 介入情報を無視した手法、PCA、Tweedie 公式を適用しない EB 手法、誤った因果グラフ（空グラフや完全グラフ）を仮定した手法などとの比較を行いました。
• 結果:
  • 真の因果グラフ（True DAG）または完全グラフ（Full DAG）を用いた提案手法（CRL f-modeling）が、他のすべての手法よりも低い相対 MSE（推定誤差）を達成しました。
  • 空グラフ（因果構造を無視）を仮定した場合、因果親の影響が強い場合に性能が低下し、不安定になることが示されました。
  • 介入情報を無視してデータをプールした場合や、PCA を適用した場合よりも、提案手法の方が優れていることが確認されました。
  • サンプル数 $N_e$ が増加するにつれて誤差が減少し、潜在次元 $d_Z$ が増加しても比較的安定した性能を示しました。

5. 意義と将来展望

• 理論的意義: 因果推論と経験的ベイズ推論を統合し、有限サンプルにおける因果表現の推定問題に対する新しい解決策を提供しました。特に、スコアマッチングに因果構造を注入するアプローチは、高次元推定におけるバイアスとバリアンスのトレードオフを改善する可能性があります。
• 実用的意義: 単細胞データ解析や、異なる実験条件（介入）下でのデータ統合など、現実の多ドメインデータ解析において、ノイズに強い因果表現の学習を可能にします。
• 今後の課題:
  • 非線形な測定モデルへの拡張（Tweedie 公式の適用が困難なため、g-modeling への移行が自然である可能性が示唆されています）。
  • 因果グラフや介入ターゲットが未知の場合への拡張（グラフ発見アルゴリズムとの組み合わせ）。
  • より複雑な介入構造（連続値の介入など）への対応。

総じて、この論文は、多ドメインデータにおける因果表現学習の推定問題を、経験的ベイズの強力な枠組みと因果構造の知識を組み合わせることで解決する、画期的なアプローチを示しています。