技術的サマリー：Flux Matching による生成モデリング

問題定義

画像合成、タンパク質設計、ロボティクスなどの高次元ドメインにおける最先端の生成モデリングは、主にスコアベースモデルに依存しています。これらのモデルは、通常スコアマッチングを通じて、Stein スコア関数 $\nabla \log p_{\text{data}}(x)$ を学習します。学習されたスコアは、ランジェヴィン力学などの拡散プロセスを駆動し、目標分布 $p_{\text{data}}$ からサンプリングするために使用されます。

しかし、スコア関数は $p_{\text{data}}$ を生成できる唯一のベクトル場ではありません。フォッカー・プランク方程式（FPE）によれば、条件 $\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$ を満たす任意のベクトル場 $f_\theta$ は、$p_{\text{data}}$ をその定常分布として持ちます。これは、同じ目標分布を生成しつつも、異なるダイナミクス（例えば、異なる混合率、循環パターン、または指向性依存関係）を誘発する生成ベクトル場の無限の族（非保存場）が存在することを意味します。

既存のスコアマッチングの目的関数は、モデルをこの空間内の一点（$\nabla \log p_{\text{data}}$）に収束させ、その偏差が目標分布を保持していても、いかなる偏差をも罰則として課します。これにより、有用な帰納的バイアス、構造的事前分布、または（より速い混合や因果的依存関係のような）動的性質を、生成場へ直接エンコードする能力が制限されます。

手法：Flux Matching

著者らは、Flux Matchingを提案します。これは、モデルをスコアと点ごとに一致させるのではなく、目標分布を生成するベクトル場の族内の任意のベクトル場を学習することで、スコアマッチングを一般化する新しいパラダイムです。

核心的な目的関数

モデル $f_\theta$ とスコア $\nabla \log p_{\text{data}}$ の間のフィッシャーダイバージェンスを最小化する代わりに、Flux Matching は確率フラックスの発散を最小化します。場が生成場であるための条件は以下の通りです：
$$\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$$
この条件は点ごとの等式よりも緩やかであり、同じフラックス発散を持つ無数の有効なベクトル場からなる「ヌル空間」を許容します。

理論的導出

この条件を最適化するために、著者らは射影フィッシャーダイバージェンスを定義します：
$$\tilde{J}(\theta) := \mathbb{E}_{x \sim p_{\text{data}}} \left[ \| \Pi_{\text{flux}} f_\theta(x) - \nabla \log p_{\text{data}}(x) \|^2 \right]$$
ここで、$\Pi_{\text{flux}}$ は、同じフラックス発散を共有する勾配場の空間 onto ベクトル場を射影します。$\tilde{J}(\theta)$ は理論的に妥当ですが、高次元では扱いが困難です。

著者らは以下の手順により、扱い可能なFlux Matching 損失（$L_{\text{flux}}$）を導出します：

1. 残差場 $u_\theta = f_\theta - \nabla \log p_{\text{data}}$ とランジェヴィン・Stein 演算子 $r_\theta = \nabla \cdot u_\theta + u_\theta \cdot \nabla \log p_{\text{data}}$ を定義します。
2. 拡散シミュレーション $d x_t = \nabla \log p_{\text{data}}(x_t) dt + \sqrt{2} dW_t$ を用いて、データ分布全体にわたって $r_\theta$ を伝播させます。
3. 部分積分とstop-gradientトリックを適用し、高価なパスごとの微分を通じた逆伝播を回避する形式に目的関数を変換します。

最終的な損失は以下の通りです：
$$L_{\text{flux}}(\theta) := -\mathbb{E}_{t \sim q, x_0 \sim p_{\text{data}}, x_t | x_0} \left[ \frac{1}{q(t)} u_\theta(x_0)^\top \text{sg}\left( \frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t) \right) \right]$$
ここで、$\text{sg}$ は stop-gradient 演算子を表します。著者らは $\nabla_\theta \tilde{J}(\theta) = 2 \nabla_\theta L_{\text{flux}}(\theta)$ であることを証明しました。

ノイズアニーリングと実用的な実装

この手法は、ノイズレベル $\sigma$ に対する連続的な場を学習することで、拡散モデルと同様のノイズアニーリング設定に拡張されます。

• スコア近似: $\nabla \log p_{\text{data}}$ は未知であるため、著者らはミニバッチ上で非パラメトリックなカーネル密度推定（KDE）近似を用いてスコアを推定します。
• 分散低減: パスごとの項 $\frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t)$ を推定するために、著者らは単一の連鎖に依存するのではなく、ガウス遷移確率で重み付けされたミニバッチ内の全サンプルを用いたクロスチェーン推定子を採用します。

特定の場の選択

Flux Matching は場の族を学習します。望ましい特性を持つ特定の場を選択するために、著者らは 2 つの戦略を提案します：

1. アプリケーション固有の損失: 混合速度や構造制約などのための補助損失を Flux Matching の目的関数に追加します。
2. モデルのパラメータ化: $f_\theta$ に直接特定の構造を強制するアーキテクチャ的制約（例えば、因果的アテンションマスク）を適用します。

主要な貢献

1. Flux Matching パラダイム: 確率フラックスの発散を一致させることで、スコアを超えたベクトル場を学習する生成モデリングフレームワークの導入。
2. 効率的な損失導出: 非保存ベクトル場を許容しつつ、フィッシャーダイバージェンスの幾何学を保持するスケーラブルな Flux Matching 損失の導出。
3. ノイズアニーリング拡張: ノイズアニーリング分布への手法の拡張により、複雑な高次元画像データセットへの応用を可能にすること。
4. 実証された柔軟性: Flux Matching が標準的なスコアマッチング下では不可能な応用を可能にすることを示す実証的証拠：
   • 高速混合: サンプリング収束を加速するベクトル場を学習すること。
   • 解釈可能性: ブラックボックスなニューラル近似なしに、構造化された生物学的に意味のあるモデル（例：RNA バイロシティ）を適合させること。
   • 構造化ダイナミクス: 指向性依存関係（例：時系列自己回帰）を生成場へ直接埋め込むこと。

結果

本論文は、Flux Matching を 5 つの設定で評価しています：

1. 制御可能な場（玩具データ）: 2 次元ガウス混合分布において、Flux Matching は目標分布を保持しつつ異なるダイナミクス（回転、循環など）を持つ場に対してゼロ損失を割り当てますが、スコアマッチングはこれらの偏差を罰則します。
2. 解釈可能な場（RNA バイロシティ）: 単一細胞ゲノミクスに適用し、Flux Matching は勾配降下法を通じて scVelo 動力学モデルのパラメータを適合させます。5 つのデータセット全体で、一貫性と境界横断の正しさにおいて標準的な EM 法に基づく scVelo 適合を上回り、性能向上がモデルの表現力増加ではなく最適化手順に由来することを示しています。
3. 制限のない生成（画像）: CIFAR-10 および CelebA において、Flux Matching は単独の目的関数として機能します。競争力のある FID スコアを達成しますが、現在、デノイジング・スコアマッチング（DSM）と比較してトレーニング中に 3〜4 倍遅く、2〜3 倍多くのメモリを使用します。ただし、高次元分布にスケーリング可能です。
4. 高速混合: 混合代理変数を最適化することで、Flux Matching は、バニラ Flux Matching や DSM よりも著しく少ないサンプリングステップで妥当な FID スコアに達する場を学習します。
5. 構造の埋め込み: 軌道生成タスク（非線形バネ）において、Flux Matching は時系列自己回帰を強制するために因果的アテンションマスクを正常に組み込みます。これに対し、DSM に同じマスクを適用すると、スコア関数の対称なヤコビアンが指向性依存関係と衝突するため、性能が低下します。

意義と主張

著者らは、Flux Matching がベクトル場自体を固定された目標ではなく設計選択として扱うことで、生成モデリングに「新たな次元」を開くことを主張しています。

• サンプリングを超えて: 主な価値は、標準的な画像生成（ここで DSM は高度に最適化されている）のための DSM の代替ではなく、有用な特性を持つダイナミクスを可能にすることにあります。
• 柔軟性: 帰納的バイアス、構造的事前分布、および機械的制約を生成プロセスに直接課すことを実務者に可能にします。
• 応用: この柔軟性によって解き放たれる具体的なユースケースを強調しています：非可逆ダイナミクスによる高速サンプリング、RNA バイロシティのような解釈可能な科学モデル、および指定された因果的または時系列構造を持つ生成モデル。

論文は結論として、Flux Matching は現在計算コストが高いものの、非保存で構造化された生成場を学習する能力は、出力の分布と同様に生成のダイナミクスが重要なアプリケーションにとって強力なツールであると述べています。