[論文レビュー] OneFlowSBI: One Model, Many Queries for Simulation-Based Inference

Mayank Nautiyal, Li Ju+6 more|arXiv (Cornell University)|Jan 30, 2026

Generative Adversarial Networks and Image Synthesis被引用数 0

ひとこと要約

OneFlowSBI は、クエリ対応マスクを用いた結合分布 p(θ, y) の単一フローメッチングモデルを学習し、再訓練なしで事後分布、尤度、任意の条件付けを可能にし、ベンチマークと実世界の問題で競争力のある精度を示す。

ABSTRACT

We introduce extit{OneFlowSBI}, a unified framework for simulation-based inference that learns a single flow-matching generative model over the joint distribution of parameters and observations. Leveraging a query-aware masking distribution during training, the same model supports multiple inference tasks, including posterior sampling, likelihood estimation, and arbitrary conditional distributions, without task-specific retraining. We evaluate extit{OneFlowSBI} on ten benchmark inference problems and two high-dimensional real-world inverse problems across multiple simulation budgets. extit{OneFlowSBI} is shown to deliver competitive performance against state-of-the-art generalized inference solvers and specialized posterior estimators, while enabling efficient sampling with few ODE integration steps and remaining robust under noisy and partially observed data.

研究の動機と目的

• 欠測やノイズのある観測を扱い、単なる事後分布を超える多様なクエリをサポートする一般化されたシミュレーションベース推論の動機づけ。
• パラメータと観測の単一の結合モデルを開発し、マスキングを介して柔軟な条件付きクエリを可能にする。
• 高効率で高忠実度の条件付きサンプリングを実現するマスク対応のフローメッチング目的を導入し、微小な ODE 積分ステップで済む。
• ベンチマークと実世界の高次元問題で専門的な SBI 手法に対して競争力のある性能を示す。
• ノイズと部分観測に対して頑健性を示しつつ、多クエリ利用ケースのサンプリング効率を維持する。

提案手法

• z = [θ; y] 上で動作する単一のフローメッチング生成モデルを用いて結合分布 p(θ, y) をモデリングする。
• 観測済みと未観測の座標を示す二値マスク m を用い、マスク付き線形補間子 z_t = m ◦ z_1 + (1−m) ◦ (t z_1 + (1−t) z_0) を定義する。
• 未観測座標のみを対象とするマスク付き残差目的で時変ベクトル場 v_t^φ(z_t, m) を訓練し、座標ごとの重み付けを持つ L_OneFlowSBI を得て、パラメータと観測の次元をバランスする。
• 直線輸送幾何を強制し、サンプリングのための解ける ODE を生むマスク付き条件付きフローメッチング損失を最適化する。
• 再訓練なしに事後、尤度、任意の条件付きを実現するよう、後方、尤度、および部分条件付けモードをブレンドするマスキング分布 p(m) を採用する。
• 観測座標をクランプし未観測座標を輸送するよう結合状態上のマスク制約付き ODE を積分してサンプルを推論する。

Figure 1 : OneFlowSBI architecture. The flow network operates on the joint state $\mathbf{z}_{t}=(\bm{\theta}_{t},\mathbf{y}_{t})$ , where observed and latent components are controlled by a binary mask $\mathbf{m}$ . Conditioned on time $t$ , the network learns a masked velocity field, yielding a si

実験結果

リサーチクエスチョン

• RQ1 単一のフローメッチングモデルが結合分布 p(θ, y) を表現し、タスク固有の再訓練なしにマスキングによる多様な推論クエリをサポートできるか。
• RQ2 標準 SBI ベンチマークで、異なる予算とモダリティにおいて、専門の後方推定器や汎用ソルバと比較して OneFlowSBI はどの程度の性能を示すか。
• RQ3 ノイズがあり部分観測のデータに対して頑健で、ハイディメンショナル問題でのサンプリング効率を維持できるか。
• RQ4 単一の訓練済みモデルから、後方、尤度、事前、周辺、混合条件付きなどのマルチクエリ推論をどの程度実行できるか。
• RQ5 マスキング戦略と ODE離散化ステップが推論精度と効率に与える影響はどの程度か。

主な発見

• OneFlowSBI は十の SBIBM ベンチマークと二つの高次元実世界問題で競争力のある後方忠実度を達成する。
• 観測マスクを変えるだけで再訓練なしに後方サンプリング、尤度評価、および任意の条件付きをサポートする。
• マスク付き直線補間によって誘導される直線輸送幾何により、少数の ODE ステップでもサンプリングの精度を維持する。
• OneFlowSBI は加法的観測ノイズおよび欠測データに対して頑健性を示し、ノイズが高い場合や欠測が多い場合にのみ劣化する。
• ベンチマーク全体で、OneFlowSBI は異なる予算とモダリティの下で専門の後方推定器や汎用ソルバと同等または上回ることが多い。
• 実世界のタスクでは、画像のブレを抑えることや浅水問題の逆問題において、マルチクエリ推論と解釈可能な後方予測チェックが効果的に機能する。

Figure 2 : SBIBM benchmark results. C2ST across ten benchmark tasks and three simulation budgets ( $10{,}000$ ; $20{,}000$ ; $30{,}000$ ), comparing OneFlowSBI with Simformer (Dense), Simformer (Undirected), and NPE. Lower C2ST indicates higher posterior fidelity.