Learned harmonic mean estimation of the marginal likelihood for multimodal posteriors with flow matching

技術要約：マルチモーダルな事後分布に対するフローマッチングを用いた学習型調和平均推定

問題提起
周辺尤度（またはベイズ証拠）は $z = p(d|M)$ と表記され、ベイズ因子を介した定量的評価を可能にする、ベイズモデル比較における基本的な量である。しかし、この量を計算するには、全パラメータ空間にわたる多次元積分を評価する必要があり、これは複雑で高次元なモデルに対して計算量的に極めて困難である。マルコフ連鎖モンテカルロ（MCMC）やハミルトニアンモンテカルロ（HMC）法は、高次元における事後分布サンプルの生成には優れているが、本質的に周辺尤度を提供するものではない。逆に、ネステッドサンプリングは証拠を提供するが、特定の入れ子状の階層構造にサンプリング戦略を制限してしまう。

「学習型調和平均推定器」は、使用されるサンプリング手法に関わらず、事後分布のサンプルのみを用いて周辺尤度を推定するために以前導入された。この手法は、事後分布内に含まれ、かつ事後分布よりも裾（テイル）が薄い内部ターゲット密度 $\phi(\theta)$ を学習することに依拠している。しかし、離散的な正規化フローや特注の密度推定器を用いた従来の構成では、高度にマルチモーダルな事後分布の扱いに苦慮することがあった。これらのアーキテクチャは、ヒューリスティックな修正（例：基底分布としてガウス混合モデルを手動で構築するなど）なしには複雑なトポロジーを捉えることができなかったり、モード間の密度ブリッジに非無視できる確率質量を割り当ててしまい、推定量の分散を無限大にしてしまったりするという課題があった。

手法
本研究は、**フローマッチング（Flow Matching）**を学習型調和平均の枠組みに統合することで、従来の密度推定器の限界に対処する。手法は以下の通りである：

1. 学習型調和平均推定器： 推定器は、事後分布のサンプル $\theta_i \sim p(\theta|d)$ と学習されたターゲット密度 $\phi(\theta)$ を用いて、逆数となる周辺尤度 $\rho = z^{-1}$ を近似する：
   $$\hat{\rho} = \frac{1}{N} \sum_{i=1}^{N} \frac{\phi(\theta_i)}{L(\theta_i)\pi(\theta_i)}$$
   推定量の分散が有限であるためには、$\phi(\theta)$ が事後分布の高確率領域内に集中し、かつ事後分布よりも速く減衰する必要がある。

2. フローマッチング・アーキテクチャ： 著者らは、離散的な正規化フローの代わりに、フローマッチングによって訓練された**連続正規化フロー（CNF）**を採用している。

   • 定義： CNFは、単純な基底分布（標準ガウス分布）とターゲット事後分布との間の変換を、時刻依存の速度場 $v_t(\theta)$ を通じてパラメータ化する。
   • 訓練目的関数： ネットワークは、シミュレーションフリーの回帰目的関数を用いて訓練される。条件付き確率パスは、ノイズサンプル $\theta_0$ と事後分布サンプル $\theta_1$ の間の直線として定義される：$\theta_t = (1-t)\theta_0 + t\theta_1$。ネットワークは、予測された速度場と真の条件付きベクトル場 $(\theta_1 - \theta_0)$ の間の平均二乗誤差を最小化する。
   • 利点： このアプローチにより、アーキテクチャに制約のない高容量のニューラルネットワーク（Swish活性化関数を用いた多層パーセプトロン）の使用が可能となり、従来のCNF法で必要とされた高価なODEシミュレーションを回避できる。

3. 温度スケーリング： ターゲット密度 $\phi(\theta)$ が事後分布よりも薄い裾を持つことを確実にするため、温度パラメータ $T \in (0, 1)$ を適用してフローの基底分布の分散をスケーリングする。これにより、フローの分布が集中され、調和平均推定器の包含条件を満たすようになる。

主な貢献

• 新規アーキテクチャ： 内部密度推定器として、フローマッチングで訓練された連続正規化フローを導入したこと。
• マルチモーダリティへの対応： このアーキテクチャが、基底分布へのヒューリスティックな修正や微調整を必要とせずに、複雑で高度にマルチモーダルな事後分布構造を正確に捉えられることを示したこと。
• スケーラビリティ： サンプリング戦略に依存しない（MCMC、HMC、NUTS、SBIなどと互換性がある）手法を提供し、中程度の高次元空間でも動作可能であること。

実験結果
著者らは、2つの困難なベンチマークを用いて手法を検証した：

1. Rastrigin関数（2次元）：

   • 多くの局所モードを持つ標準的なベンチマーク。
   • フローマッチングモデル（6層、256隠れユニット）は、温度 $T=0.98$ においてすべてのモードを正常に捉えた。
   • 100回の実験ランの結果、推定された周辺尤度の値は正確であり、数値積分によって得られた真値と密接に一致していた。誤差推定も信頼できるものであった。

2. 20次元ガウス混合モデル：

   • 非重複のモードと相関が導入された、5つのガウス成分からなる20次元の混合モデル。
   • モデル（10層、256隠れユニット）は、40,000個のサンプルを用いて訓練された。
   • 学習されたターゲット分布は、事後分布内に適切に集中しており、困難なトポロジーを捉えていることが示された。
   • 100回のランにおける推定周辺尤度は、解析的な真値と一致しており、この複雑な中次元シナリオにおいても、本手法が正確かつ堅牢な誤差推定を行えることが確認された。

意義と主張
本論文は、フローマッチングで訓練されたCNFの表現力を活用することで、学習型調和平均推定器が、高度に複雑でマルチモーダルな事後分布に対しても頑健に証拠を推定できると主張している。主な意義は、マルチモダリティを扱うために以前必要であったヒューリスティックな調整（手選定のガウス混合モデルなど）への依存を排除したことにある。著者らは、このアプローチにより、加速された、あるいはスケーラブルなあらゆるサンプリング手法からのサンプルを用いて正確な証拠計算が可能になり、天体物理学のようなデータ豊富な分野におけるベイズモデル比較の計算上のボトルネックを解消できると断言している。

著者が述べているように、今後の課題には、重力波天体物理学などの実世界の課題への適用や、画像処理やフィールドレベルの宇宙論的推論のような超高次元タスクに向けたさらなるアーキテクチャの洗練が含まれる。