技術要約：Flow Annealing Posterior Sampling (FAPS)

問題提起

本論文は、科学的な逆問題、データ同化、および不確実性定量化において中心的な課題である、確率過程に対する原理的な回帰（principled regression）の問題に取り組んでいる。これらの設定では、疎でノイズを含む観測量に基づき、未知の関数または物理場に対する事後分布を推論することが目標となる。古典的なガウス過程（GP）回帰は、ガウス事前分布と単純な観測モデルに対してベイズ的な解法を提供するが、多くの科学的プロセスは非ガウス的であり、高次元であり、かつ間接的な物理的測定（例：偏微分方程式（PDE）の解）を通じて観測される。

既存のアプローチには以下の限界がある：

• Neural Processes は、定式化された条件付き予測器を学習するが、変化するクエリの離散化に対して厳密な事後一貫性を欠くことが多く、再利用可能な事前分布の学習とテスト時の推論を明示的に分離できていない。
• Neural Operators は、効率的なPDE解写像を提供するが、通常は決定論的であり、疎なデータから事後分布を生成することはない。
• 拡散ベースの事後サンプラー（例：DAPS, FunDPS）は、学習された事前分布と観測によるガイダンスを組み合わせるが、多くの場合、拡散ベースであり、計算コストが高く、特定のPDE設定に特化しており、回帰と逆問題を統一するものではない。

事前分布の密度を明示的に評価することなく、学習済みの関数空間事前分布を活用して、直接的な確率過程回帰と間接的なPDE逆問題の両方に対して尤度ガイド型の推論を実行できる、一般的な事後サンプリング・フレームワークが求められている。

手法：Flow Annealing Posterior Sampling (FAPS)

FAPSは、学習済みであるOperator Flow Matching (OFM) 事前分布に基づいた、関数空間の事後サンプリング・フレームワークである。これは、デカップルされたアニーリング手順を通じて、確率過程回帰とPDE逆問題を統一する。

1. 関数空間フロー・マッチング事前分布

本手法は、参照ガウス測度 $\gamma = \mathcal{N}(0, C_0)$ をターゲットとなるデータ分布 $\mu_1$ へ輸送する確率フローを学習した、学習済みのOFM事前分布を利用する。これは、ノイズを含む直線的な補間パスに沿った条件付き速度場 $v_t(u_t | u_0, u_1) = u_1 - u_0$ を近似するように訓練されたニューラルオペレータ $v_\theta$ を通じて達成される。結果として得られるフロー写像 $\Phi_{s \to t}$ により、参照測度からターゲットプロセスへの輸送が可能となる。

2. 統一された事後分布の定式化

FAPSは、以下の単一の観測モデルの下で、関数的回帰とPDE逆問題の両方を扱う：
$$y = A(u_1) + \epsilon$$
ここで、$u_1$ は未知の関数、$A$ は観測演算子、$\epsilon$ はノイズである。

• 関数的回帰の場合、$A$ は直接的な点評価またはマスキング演算子である。
• PDE逆問題の場合、$A$ は、疎な観測演算子と、前方PDE解写像を表す学習済みのニューラルオペレータ・サロゲート $G_\phi$ の合成である。

3. アニーリングおよびリブリージング（Re-bridging）手順

FAPSは、アニーリング・スケジュール $0 = t_0 < t_1 < \dots < t_K = 1$ に沿って、3つのステップを交互に繰り返すことで、事後分布 $\pi_1^y$ からサンプリングを行う：

1. フロー輸送（エンドポイントへ）： アニーリング・レベル $t_k$ における「ノイズを含んだ」状態 $u_{t_k}$ が与えられると、本手法は学習済みのフローをエンドポイントまで積分し、「クリーンな」アンカー $\hat{u}_k = \Phi_{t_k \to 1}(u_{t_k})$ を得る。このアンカーは、クリーンな関数に対する事前分布と整合した推定値として機能する。
2. 尤度ガイド型ランジュバン補正： アンカー $\hat{u}_k$ の周囲で、FAPSは観測尤度 $p(y|u_1)$ に基づいてサンプルを補正するために $L$ ステップのランジュバン動力学を実行する。更新ルールは以下の通りである：
   $$u_1^{(\ell+1)} = u_1^{(\ell)} + \eta \left[ -\frac{u_1^{(\ell)} - \hat{u}_k}{\lambda_k^2} + C_X \nabla_{u_1} \log p(y|u_1^{(\ell)}) \right] + \sqrt{2\eta} \zeta^{(\ell)}$$
   極めて重要な点は、勾配がターゲットプロセスからオフラインで推定された低ランク共分散行列 $C_X$ によってプリコンディショニングされていることである。このプリコンディショナは、疎な観測情報を支配的な関数空間の相関に沿って伝播させ、未観測領域における更新を大幅に改善する。
3. リブリージング（Re-bridging）： 補正されたクリーンなサンプルは、フローパスによって誘導されるブリッジ・カーネルを用いて、次のアニーリング・レベル $t_{k+1}$ へとリブリージングされる。これにより、次の反復に備えて、実質的に状態にノイズが再び加えられる。

このプロセスは $t_K = 1$ になるまで繰り返され、事前分布密度を明示的に評価することなく事後サンプルが得られる。

主な貢献

1. 初の統一された関数空間フロー・マッチング事後サンプラー： FAPSは、学習済みの関数空間フロー・マッチング事前分布を用いて、確率過程回帰とPDE逆問題を統一する最初のフレームワークとして提示される。これは、直接的な観測と間接的な観測の両方に対して共通の尤度ガイド型サンプラーを提供する。
2. 低ランク共分散プリコンディショニングを用いたランジュバン補正： 本手法は、ランジュバン補正ステップのための低ランク共分散プリコンディショナを導入している。これは関数空間における支配的な相関を利用するものであり、未プリコンディショニングの手法と比較して、疎な観測下での事後更新を大幅に改善する。
3. 柔軟かつ計算効率の高い推論： FAPSは「プラグアンドプレイ」であり、学習済みの事前分布と前方サロゲートを活用する。これは、再学習なしに、変動するクエリの離散化やノイズを含む観測をサポートする。また、事前分布密度の明示的な評価を回避するため、フロー事前分布を用いた尤度ベースのサンプリングや拡散ベースのサンプラーよりも、大幅にメモリおよび時間の効率が高い。
4. 強力な実証的性能： 広範な実験により、FAPSが関数回帰において（CNP, ANP, NDP, FlowNPなどのNeural Processのバリアントを上回り）、PDE逆問題において（拡散ベースラインと比較して、同等または優れた性能を示す）高い性能を達成することが示されている。

実験結果

本論文では、以下の3つの設定でFAPSを評価している：

• ガウス過程回帰： 1D MatérnおよびGibbsカーネルのベンチマークにおいて、FAPSは、変化するクエリ解像度（128, 512, 1024）にわたって、参照事後平均および不確実性に密接に一致する。FAPSは、特に解像度が高くなりベースラインが劣化する領域において、Sliced Wasserstein Distance (SWD) および Maximum Mean Discrepancy (MMD) に関して、Neural Processesを凌駕する。
• 非ガウス的関数回帰： Navier-Stokes、ブラックホール撮像、および全球気候データセットにおいて、FAPSは支配的な空間構造を再構成し、一貫した事後サンプルを生成する。ベースラインがしばしば過度に平滑化したり不確実性を過小評価したりする中で、FAPSは最高の、あるいはそれに準ずるキャリブレーション（SSR）および再構成指標（CRPS, PSNR, SSIM）を達成する。
• PDE逆問題： Darcy flow、Poisson、Helmholtz、およびNavier-Stokesの逆問題ベンチマーク（解像度128x128、ノイズを含む128個の観測）において、FAPS（FNOおよびUNetのバリアントの両方）は、DiffusionPDE、FunDPS、DDISと比較して、CRPS、PSNR、SSIMにおいて競争力のある、あるいは最高の成績を収める。
  • 効率性： FAPSは拡散ベースのベースラインよりも大幅に高速であり（DiffusionPDEより1.73倍、DDISより2.37倍高速）、フロー事前分布の明示的な尤度評価や大規模な拡散ステップ数を回避することで、大幅に少ないGPUメモリを必要とする。
  • ゼロショット超解像： FAPSは、事前学習時に見られなかったクエリメッシュ（例：160x160）に対して逆推論を実行する能力を示しており、これはOFM事前分布の関数空間的な性質によって可能となっている。

重要性と主張

著者らは、FAPSが科学的機械学習における不確実性を考慮した推論の新しいパラダイムを確立すると主張している。学習済みの関数空間フロー・マッチング事前分布を尤度ガイド型の事後サンプラーへと変換することで、FAPSは既存の手法に対して、再利用可能で柔軟、かつ計算効率の高い代替手段を提供する。

著者らは、実用的なサンプラーが近似的な局所ガウス補正に依存していること、および完全な非漸近的収束理論が今後の課題であることを控えめに認めているが、誤差伝播の分析と経験的証拠を提供し、局所的な遷移誤差は加算的に蓄積されること、およびアニーリングメッシュとサンプリングパラメータを精緻化することで手法がターゲット事後分布に収束することを示唆している。

最終的に、本研究は、関数空間フロー・マッチング事前分布が、複雑な科学領域におけるスケーラブルなベイズ推論のための堅牢な基盤となり得ること、そして、単一かつ効率的なフレームワークの下で回帰と逆モデリングを統一できることを示唆している。