技術概要：Flowing with Confidence (FMwC)

問題定義

生成モデル、特にフローマッチング（FM）モデルは、言語、視覚、科学の分野において最先端のパフォーマンスを達成しています。しかし、これらは（ハルシネーションした引用、解剖学的に誤った図、または熱力学的に不安定な結晶などの）出力を生成しますが、信頼性の指標は本来的に存在しません。Deep Ensembles や MC-Dropout などの既存の不確実性定量化ソリューションは、サンプルごとに $k$ 個の別々のモデル、または $k$ 個の確率的軌跡を実行することに依存しています。このアプローチは $k\times$ の計算コストを伴うため、大規模展開におけるサンプルごとの信頼性の評価は過度に高価となります。さらに、尤度ベースの手法は、深層生成モデルにおいて分布内入力と分布外入力を区別する際に信頼性が低いことが示されています。ここで扱われる中心的な問いは、フローマッチングモデルが推論コストを増加させることなく、また尤度へのアクセスを必要とせずに、生成された出力に対する自身の不確実性を定量化することを学習できるかどうかです。

手法：Flow Matching with Confidence (FMwC)

FMwC は、標準的な決定論的フローマッチングと同じ計算コストで、フローマッチングモデルにサンプルごとの軌道レベルの信頼性スコアを付与する単一パスアプローチです。この手法は、変分適応的ドロップアウト（VAD）を、常微分方程式（ODE）軌道を通じた閉形式のモーメント伝播と統合します。

1. 変分定式化

単一の決定論的速度場 $v_{t,\theta}(x_t)$ を学習する代わりに、FMwC はネットワーク重み $\omega$ を変分事後分布 $q_\psi(\omega)$ から引き出された確率変数として扱います。モデルはエビデンス下限（ELBO）を最大化します：
$$\text{ELBO}(\psi) = \mathbb{E}_{\omega \sim q_\psi} [\log \mathcal{L}(D | \omega)] - \text{KL}[q_\psi(\omega) \parallel \pi(\omega)]$$
ここで、尤度項は標準的な条件付きフローマッチング回帰損失に対応します。重要なのは、事後分布下での期待対数尤度の勾配が、重みサンプル上で評価された標準的な回帰損失の勾配と一致することであり、これによりシミュレーション不要なトレーニングが可能となります。

2. VAD による状態依存事後分布

不確実性がフローの局所幾何学に敏感であることを保証するために、FMwC は変分適応的ドロップアウト（VAD）を採用します。この枠組みにおいて、層 $l$ の重みに関する事後分布はガウス分布 $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$ であり、ここでドロップアウトスケール $\alpha_l(x, t)$ は小さな推論ネットワーク $E_\gamma$ によって予測されます。これにより、事後分布の広がりは幾何学的な曖昧さがある領域（速度場の適合が困難な領域）で広がり、それ以外の領域では狭まるよう、ELBO を通じてのみ学習されます。

3. 解析的分散伝播

FMwC の中核的な革新は、速度場の分散を、単一のフォワードパスにおいて、ニューラルネットワーク層と ODE 積分軌道の両方を通じて解析的に伝播させることです：

• ネットワーク内での伝播： 局所再パラメータ化トリックを用いて、活性化の平均（$\mu_l$）と分散（$\sigma_l^2$）を層ごとに伝播させます。線形層では、分散の再帰は閉形式です。非線形関数については、モーメントをモンテカルロサンプリングを必要とせず、ガウス・エルミート求積法またはテイラー線形化を用いて解析的に計算します。
• ODE 内での伝播： オイラーステップの各時点で、ネットワークは事後平均速度 $\mu_t$ とその分散 $\sigma_t^2$ を出力します。積分器は $\mu_t$ を用いて位置を進め、$\sigma_t^2$ を記録します。これにより、単一のスカラーではなく、分散の軌道 $\{\sigma_t^2\}_{t=0}^T$ が得られます。

4. 信頼性読み出し

分散軌道は、下流タスクのためのスカラー信頼性スコアに集約できます：

• 終点信頼性： 最終ステップ（$t=1$）における分散。
• 積分信頼性： 軌道に沿った速度標準偏差の合計。
• 時間的比率： 後期ウィンドウの分散と初期ウィンドウの分散の比率。これは、収束するサンプル（後期分散が低い）と曖昧なままのサンプルを区別するのに役立ちます。
• 学習されたヘッド： 小さな教師ありヘッドが、完全な分散系列をタスク固有のスコアにマッピングできます。

主要な貢献

1. 単一パス信頼性推定機： FMwC は、標準的な FM の推論予算に匹敵する単一の決定論的フォワードパスのコストで、フローマッチングに対する較正されたサンプルごとの信頼性スコアを提供する最初の手法です。これは、学習された入力依存ノイズを ODE を通じて解析的に伝播させることで達成されます。
2. 再トレーニング不要な推論の改善： 信頼性スコアは、再トレーニングなしで実用的な応用を可能にします：
   • フィルタリング： 低信頼性の出力を破棄することで、サンプル品質を単調に向上させます。
   • 編集： モデルが構造にコミットする正確なタイムステップ（分散のピークである $t^*$）を特定し、$t^*$ でノイズを注入することで、化学的スワップ（結晶内）など、制約を尊重した標的編集を可能にします。
   • 適応的ステップ： 高分散（曖昧さ）の領域に ODE 計算ステップを集中させ、均等なステップでは失敗する低ステップ予算下でも高品質なサンプルを回復します。
3. 幾何学的解釈可能性： 本論文は、FMwC の信頼性スコアが学習された速度場の局所発散と相関することを確立しています。これにより、生成プロセスへの窓が開かれ、認識的不確実性がフローの幾何学的不安定性（分岐点）と結びついています。

実験結果

著者は、FMwC を 3 つのモダリティ（2 次元密度推定（Checkerboard）、画像生成（MNIST）、および新規無機結晶生成（Riemannian Flow Matching））で評価しました。

• サンプル品質： FMwC は、標準的な FM と比較してサンプル品質を維持または向上させます。Checkerboard と MNIST において、それぞれ誤配置率と FID スコアが低く、結晶においては同等の品質を維持しました。
• 推論コスト： FMwC は、より大きなバックボーン（例：MNIST で約 1.003 倍、結晶で約 1.08 倍）において、決定論的ベースラインに対して無視できるオーバーヘッドしか発生させません。対照的に、アンサンブルおよび MC-Dropout ベースラインは $5\times$ の計算を必要とします。
• フィルタリング性能： FMwC スコアに基づくフィルタリングは、$k=5$ のアンサンブルおよび MC-Dropout と同等の精度 - 再現曲線下面積（AUPRC）スコアを達成しましたが、単一軌道のコストで実現しました。例えば、Checkerboard データセットにおいて、時間的比率スコアは AUPRC 0.55 を達成し、アンサンブルベースライン（0.55）と一致しましたが、計算コストは 1/5 でした。
• 編集の成功： 分散のピーク（$t^*$）でノイズを注入して結晶を編集する際、FMwC はランダムなタイムステップ選択と比較して（化学的スワップで 39% 対 22% など）、有意に高い成功率を達成しました。これは、高分散領域が重要な意思決定点に対応することを示しています。
• 発散との相関： 分散軌道 $\sigma_t^2$ は、速度場の発散の大きさ（$|\nabla \cdot v_{t,\theta}|$）と強い相関を示しました。Crystal データセットでは、積分相関はスピアマン $\rho$ 0.96 に達し、信頼性スコアの幾何学的解釈を検証しました。

意義と主張

本論文は、FMwC がサンプルの生成器を、信頼性の尺度を伴うサンプルの生成器へと変換すると主張しています。重みに対する入力依存事後分布を学習し、分散を解析的に伝播させることで、FMwC は科学および安全性が重要な応用において、信頼性認識型の生成を実用的なデフォルトにします。著者は、信頼性スコアがアンサンブルによって測定されるモデル間の不一致とは異なり、単一のトレーニング済みフローの内在的な幾何学的性質であると強調しています。これにより、既存の不確実性定量化手法の過度なコストなしに、不安定な瞬間を標的とした「外科的」な生成プロセスのガイダンスが可能になります。この研究は、学習された速度場の幾何学を理解し活用することが、より信頼性が高く、解釈可能で、制御可能な生成モデルにつながることを示唆しています。