[論文レビュー] Flow Matching for Scalable Simulation-Based Inference

Maximilian Dax, Jonas Wildberger+4 more|arXiv (Cornell University)|May 26, 2023

Model Reduction and Neural Networks被引用数 9

ひとこと要約

FMPE は連続正規化フローを用いたフローマッチングをシミュレーションに基づく推論へ適用し、制約のないアーキテクチャ、扱える密度、そしてスケーラブルな性能を実現します。競争力のあるベンチマークと、訓練時間の短縮によるより優れた重力波推論を達成します。

ABSTRACT

Neural posterior estimation methods based on discrete normalizing flows have become established tools for simulation-based inference (SBI), but scaling them to high-dimensional problems can be challenging. Building on recent advances in generative modeling, we here present flow matching posterior estimation (FMPE), a technique for SBI using continuous normalizing flows. Like diffusion models, and in contrast to discrete flows, flow matching allows for unconstrained architectures, providing enhanced flexibility for complex data modalities. Flow matching, therefore, enables exact density evaluation, fast training, and seamless scalability to large architectures--making it ideal for SBI. We show that FMPE achieves competitive performance on an established SBI benchmark, and then demonstrate its improved scalability on a challenging scientific problem: for gravitational-wave inference, FMPE outperforms methods based on comparable discrete flows, reducing training time by 30% with substantially improved accuracy. Our work underscores the potential of FMPE to enhance performance in challenging inference scenarios, thereby paving the way for more advanced applications to scientific problems.

研究の動機と目的

• 複雑な科学モデルに対するスケーラブルで厳密な密度 SBI によるニューラルポスタ推定を動機づける。
• FMPE を用いた SBI の flow matching 後方推定の導入。
• FMPE が質量を網羅する後方分布と大規模なアーキテクチャへのスケーラブルな訓練を提供することを示す。
• FMPE の SBI ベンチマークと重力-wave 推論における性能を示し、実用的な訓練時間の利得を得る。

提案手法

• ベイズ推論にフローマッチングを適用し、対数尤度訓練をサンプル条件付きのフローマッチング目的 FMPE 損失に置換する。
• ODEs によって決定的なサンプルトラジェクトリを定義するベクトル場パラメータ化 v_t,x を使用。
• サンプル条件のパス p_t(θ|θ1) と対応するベクトル場 u_t(θ|θ1) を用いて訓練し、L_FMPE = E[...] ||v_t,x(θ_t) − u_t(θ_t|θ1)||^2 を最小化。
• θ へのパラメータ空間条件付けを処理する高次元 x を処理するアーキテクチャを提供し、t 条件付けはゲーティッドリニアユニットを介して。
• p_α(t) を用いた時間事前再重み付けを許し、特定の t 領域を強調する; データを用いた訓練は p(θ)p(x|θ) からのデータで訓練するというベイズの定理を使う。
• 質量をカバーする特性と、正則性仮定の下でベクトル場間の MSE と KL 発散の理論境界を示す。

実験結果

リサーチクエスチョン

• RQ1 FMPE は SBI において真の事後 p(θ|x) を扱える密度評価可能な近似として近づけられるか。
• RQ2 FMPE は p(θ|x) のサポートをベンチマークタスク全体で保持し、質量をカバーする後方分布を提供するか。
• RQ3 高次元の x や大規模ネットワークにおいて、離散正規化フローと比べて FMPE はどの程度スケールするか。
• RQ4 重力波推論のような難密な科学応用で、NPE と比較して FMPE はどんな利点を提供するか。

主な発見

• FMPE は SBI ベンチマークで NPE と競合する性能を示し、いくつかのタスクで改善を示し、質量をカバーする挙動を明確に示す。
• FMPE は扱える密度を得られ、制約のないアーキテクチャをサポートし、より大きなネットワークとより良いスケーリングを可能にする。
• 重力波推論では、FMPE は精度と訓練効率で NPE を大幅に上回り、訓練時間を約 30% 削減しながらポスタ―品質を向上させる。
• GW150914 に対する FMPE のポスターは、シンメトリー強化手法 GNPE に近い、あるいは同等の Jensen-Shannon 発散を達成し、実用的な精度向上を示す。
• 経験的分析は、FMPE が benchmark タスク全般で質量をカバーするポスターを生み出し、参照ポスターをカバーする密度推定を提供することを示す。