技術要約：ベイズ逆問題における分散減少のための条件付きニューラル制御変量（Conditional Neural Control Variates）

問題設定

ベイズ推論における逆問題では、事後分布下での期待値（事後平均、分散、または予測量など）を計算する必要がある。これらの期待値は通常、モンテカルロ（MC）法を用いて推定される。しかし、関心対象となる関数 $h(x)$ が事後分布の下で大きく変動する場合、正確な推定には膨大な数のサンプルが必要となる。このコストは、各事後サンプルごとに計算負荷の高い順モデルのフォワード・ソルブ（前向き計算）を必要とする偏微分方程式（PDE）制約付き逆問題において、特に深刻である。

制御変量（Control variates）は、バイアスを導入することなくMC推定量の分散を減少させる古典的な解決策であるが、既存のスコアベースの手法（例：スタイン制御変量）は、新しい観測値 $y$ ごとに再学習が必要であるという決定的な限界を抱えている。この償却性の欠如は、各観測に対して再学習を行うコストが高すぎる場合、計算的に実行不可能にする。さらに、一般的なニューラルネットワークのダイバージェンス（発散）の計算に伴う計算コスト（通常、$O(d)$ のバックプロパゲーションまたは確率的トレース推定を必要とする）により、スタインベースの制御変量を高次元にスケールさせることは困難である。

手法

著者らは、観測値に対して償却可能（amortized）な設計を持つモジュール式の分散減少手法である**条件付きニューラル制御変量（Conditional Neural Control Variates: CNCV）**を提案する。この手法は、シミュレーションベースの推論パイプラインとシームレスに統合できる。

1. 条件付き制御変量の構成

本手法の中核は**スタインの恒等式（Stein's identity）**にあり、これは滑らかな関数 $\phi(x, y)$ とターゲット分布 $p(x|y)$ に対して以下を保証する：
$$\mathbb{E}_{p(x|y)}[\nabla_x \cdot \phi(x, y) + \phi(x, y) \cdot \nabla_x \log p(x|y)] = 0$$
この恒等式は、ゼロ平均の制御変量の族 $g(x, y) = \nabla_x \cdot \phi(x, y) + \phi(x, y) \cdot \nabla_x \log p(x|y)$ を定義する。著者らは、$\phi$ をニューラルネットワーク $\phi_\theta$ を用いてパラメータ化する。

2. 階層的カップリング層アーキテクチャ

高次元におけるダイバージェンス $\nabla_x \cdot \phi_\theta$ の計算に伴う計算上のボトルネックに対処するため、著者らは階層的アフィン・カップリング層のアンサンブル（HINTに着想を得たもの）に基づく特定のアーキテクチャを提案している。

• 構造: ネットワークは、アフィン・カップリング層を再帰的な二分木の中に構成する。各ノードにおいて、入力は分割され、下位のパーティションは上位のパーティションに基づいて変換される。
• 三角ヤコビ行列: この構造により、ヤコビ行列が下三角行列になることが保証される。これにより、ダイバージェンスを $O(d)$ 回のバックプロパゲーションを介することなく、単一のフォワードパスで正確に計算できる。
• アンサンブルと置換: 次元間の相互作用を効果的に捉えるため、入力次元にランダムな置換を適用する $L$ 個の独立したネットワークのアンサンブルを採用する。最終的な制御変量は、これらアンサンブルメンバーの平均となる。

3. 学習と推論

• 学習データ: 本手法は、パラメータとデータの結合分布 $(x, y)$ からのサンプルを必要とする。これらは、シミュレーションベースの推論パイプライン（事前分布から $x$ をサンプリングし、$y = F(x) + \epsilon$ を計算する）から自然に得られる。
• スコアのソース: 事後スコア $\nabla_x \log p(x|y)$ は、物理ベースの尤度、ニューラルオペレータ・サロゲート、または学習済み生成モデル（例：条件付き正規化流）から導出できる。
• 目的関数: ネットワークのパラメータは、制御された推定量の平均二乗誤差を最小化するように訓練される。これは、制御変量と関心対象の相関を最大化しつつ、変量自体の分散を抑制することに相当する。
• 償却（Amortization）: 一度訓練されると、CNCVモデルはあらゆる新しい観測値 $y_{obs}$ に対して、再学習なしで適用可能である。推論時には、条件付き正規化流（CNF）などを通じて事後サンプルが抽出され、スコアが評価され、学習された制御変量が生のMC推定値から差し引かれる。

主な貢献

1. 償却可能な制御変量: 著者らは、ニューラル制御変量を条件付きの設定へと拡張し、観測値ごとに再学習する必要がない、あらゆる事後サンプラーおよびスコアソースと統合可能な手法を構築した。
2. 計算可能な高次元ダイバージェンス: ランダムな入力置換を伴う階層的カップリング層のアンサンブルを利用することで、単一のフォワードパスで正確なダイバージェンス計算を可能にし、高次元問題へのスケーラビリティを実現した。
3. 実証的検証: 本手法は、複雑さが段階的に増す4つの逆問題（ガウス型、ローゼンブロック、非線形、およびPDE制約付きダルシー流）を用いて検証されており、解析的なスコアがサロゲートに置き換えられた場合でも大幅な分散減少を実現している。

実験結果

著者らは、性能を示す指標として分散減少係数（Variance Reduction Factor: VRF）を報告している（値が低いほど性能が良いことを示す）。

• 定型問題:
  • ガウス事後分布: CNCVは、次元 $d \in \{2, 4, 8, 16\}$ における平均推定において、0.009から0.083のVRFを達成した。スコアが条件付き正規化流（CNF）によって近似されている場合でも、手法は堅牢性を維持した。
  • ローゼンブロック（非ガウス型）: 本手法は、多様な観測値（左裾、右裾、リッジ）に対して、0.08から0.23の範囲のVRFで成功裏に汎化性能を示した。これは、非ガウス型のターゲットに対して著しく性能が低下した古典的な多項式スタイン制御変量を上回る結果であった。
  • 非線形順モデル: サンプリングにデノイジング拡散モデルを用い、スコアにCNFを用いた場合、CNCVは0.57の平均VRFを達成した。
• PDE制約付きダルシー流:
  • 対数浸透率推定を含む100次元の問題において、CNCVは $\sim 0.18$ のVRFを達成した。
  • これは、生のMCと比較して、平均二乗誤差（MSE）を $\sim 5.5$ 倍減少させたことに相当する。
  • 極めて重要な点として、サンプルが事前学習済みのCNFから抽出されているため、推論コストは無視できる程度（$\sim 0.5$秒）であり、同等の精度を得るために必要な生のMCのコスト（PDEの解法に $\sim 100$秒）と比較して極めて効率的である。
• 堅牢性: 手法はスコア近似誤差に対して頑健であることが証明された。学習されたスコアが17.5%の相対誤差を持っていても、VRFは解析的スコアを用いた場合と同等の水準を維持した。

意義と主張

本論文は、CNCVが、高次元かつPDE制約付きの環境におけるベイズ逆問題の根本的なボトルネックに対処することを主張している。

• 償却性: 従来のニューラル制御変量手法は、すべての観測に対して再学習を必要とするが、CNCVは結合データを用いて一度のみ訓練を行い、未知の観測に対して汎化する。著者らは、およそ20個以上の異なる観測を必要とする問題においては、CNCVの償却コストが非償却型のベースラインよりも有利になると指摘している。
• スケーラビリティ: カップリング層の三角ヤコビ行列構造を活用することで、高次元におけるダイバージェンス推定に伴う $O(d)$ の計算障壁を克服している。
• 柔軟性: 手法はスコアのソースに依存せず、物理ベースの解析的スコアと学習された生成モデルの両方に対して効果的に機能する。
• 不偏性: 著者らは、本手法は分散を減少させるものの、基礎となる事後近似を変更するものではないと強調している。もしスコアとサンプルが同一の学習ソース（例：CNF）から得られているのであれば、推定値は学習された事後分布下での期待値に対して厳密に不偏である。学習されたスコアと真の事後分布との間の不一致は、制御変量のメカニズムではなく、純粋に学習されたスコアの質に起因するものである。

著者らは、CNCVは非ガウスの設定において、最適な制御変量が非線形であり、かつ償却が必要となる場合に特に価値があり、複雑な逆問題における効率的な不確実性定量化への実用的な道筋を提供するものであると結論付けている。