技術概要：ドリフト・フロー・マッチング（DFM）

問題定義

フロー・マッチングや拡散モデルなどの反復的生成モデルは、「テスト時スケーリング」の能力を実証しており、推論ステップ数の増加が生成品質、制御性、分布カバレッジの向上をもたらすことが示されています。しかし、これらの手法は複数の推論ステップを必要とし、計算コストが高くなる傾向があります。一方、ドリフトモデルは、ソース分布からターゲット分布への直接のプッシュフォワード写像を学習することで、効率的なワンステップ生成を実現します。それでも、その直接生成パラダイムには、テスト時に推論ステップ数を増やすことで結果を改善するための自然なメカニズムが欠けており、その柔軟性と異なる品質・効率要件への適応性が制限されています。

手法：ドリフト・フロー・マッチング（DFM）

著者らは、ワンステップ・ドリフトモデルの効率性とマルチステップ・フロー・マッチングのスケーラビリティを統合するフレームワークとして、ドリフト・フロー・マッチング（DFM）を提案します。DFM は、フロー・マッチングの確率経路に沿って、任意の現在の時刻 $t$ から任意の将来の時刻 $r > t$ へサンプルを移動させる二時刻輸送モデルを学習することでこれを実現します。

中核メカニズム

1. 二時刻周辺分布の構築:
   点ごとの条件付き速度ターゲットを回帰する標準的なフロー・マッチングとは異なり、DFM は周辺分布のペア $(p_t, p_r)$ に対して動作します。これは、ソース分布 ($X_0$) とターゲット分布 ($X_1$) の独立した結合に基づき、条件付き経路 $X^Z_t = \alpha(t)X^Z_0 + \beta(t)X^Z_1$ および $X^Z_r = \alpha(r)X^Z_0 + \beta(r)X^Z_1$ を定義します。この構築により、任意の時刻ステップにおいて周辺分布 $p_t$ と $p_r$ が誘導されます。

2. 平均速度のパラメータ化:
   DFM は、現在の分布 $p_t$ を将来の分布 $p_r$ へ輸送する平均速度場 $u_\theta(x_t, t, r)$ をパラメータ化します。輸送写像は以下のように定義されます：
   $$T^\theta_{t,r}(x_t) = x_t + (r - t)u_\theta(x_t, t, r)$$
   この定式化により、モデルは任意のステップサイズ $(r-t)$ を処理できます。$r \to t$ の極限において、平均速度は標準的なフロー・マッチングの瞬間的周辺速度を回復します。

3. ドリフトに基づく監督:
   速度ベクトルを回帰する代わりに、DFM は予測分布 $q^\theta_{t,r}$（$T^\theta_{t,r}$ によって生成される）を真のターゲット周辺分布 $p_r$ へ輸送するために、ドリフトに基づく目的関数を採用します。

   • ドリフト場: ドリフト場 $V_{q^\theta_{t,r}, p_r}$ は、ターゲットサンプルへのカーネル重み付き引力 ($V^+_{p_r}$) と、現在のモデル分布からの反発的な自己修正項 ($V^-_{q^\theta_{t,r}}$) を用いて構築されます。
   • 損失関数: 学習目的は、ストップグラディエント定式化を用いてドリフト場の大きさを最小化することです：
     $$L_{DFM}(\theta) = \frac{1}{2} \mathbb{E}_{(t,r), X_t \sim p_t} \left[ \| \hat{X}^\theta_r - \text{sg}(\hat{X}^\theta_r + V_{q^\theta_{t,r}, p_r}(\hat{X}^\theta_r)) \|^2 \right]$$
   • グループ推定: 重要なのは、ドリフト場と損失が、各サンプリングされた時刻ペア $(t, r)$ に対してグループごとに計算されなければならない点です。異なる時刻ペアからのサンプルを混合すると、各ペアごとのターゲット周辺分布 $p_r$ が異なるため、ドリフト推定が破綻します。

主要な貢献

• 統合フレームワーク: DFM は、効率的なワンステップ・ドリフトモデルと反復的フロー・マッチング手法の間のギャップを埋め、直接生成と反復的洗練の両方をサポートする単一のパラダイムを創出します。
• テスト時スケーリング: このフレームワークは、ドリフトベースの手法にテスト時スケーリング能力をもたらします。ユーザーは再学習なしに、推論ステップ数（NFE）を調整することで、サンプリング速度と生成品質のトレードオフを実現できます。
• 新しい学習原理: 本論文は、フロー・マッチングの標準である点ごとの速度回帰から、ドリフト監督による周辺分布輸送へと移行する学習原理を導入します。これは、他のワンステップ生成手法で用いられる自己蒸留法に伴う不安定性を回避します。
• 理論的保証: 著者らは、DFM がエンドポイント生成をワッサーシュタイン距離 ($W_2$) によって制御される短距離周辺輸送問題に分解し、無限小ステップ極限において標準的なフロー・マッチングを回復することを示す理論的解析を提供します。

実験結果

合成データ、MNIST、FFHQ、ImageNet-1k、およびロボティクス制御タスクにおける広範な実験は、DFM の有効性を示しています：

• 合成データ: 可視化により、DFM は単一ステップで高品質なサンプルを生成でき、NFE が増加するにつれて分布カバレッジとサンプル品質がさらに向上することが示されました。
• 条件付き生成（MNIST & FFHQ）: DFM はワンステップ推論において競争力のある性能（ドリフトモデルのベースラインと同等）を達成し、NFE を 2、5、または 10 に増加させるにつれて、指標（EMD、FID）が一貫して改善されます。
• 大規模画像合成（ImageNet-1k）: ImageNet-256x256 データセットにおいて、DFM は NFE=1 で 1.52 の競争力のある FID を達成しました。NFE を 10 に増加させることで FID はさらに 1.31 に低下し、反復的フロー・マッチング手法に匹敵する効果的なテスト時スケーリングを実証しました。
• ロボティクス制御: ロボット操作タスク（Lift、Can、PushT など）において、DFM はドリフトポリシーのワンステップ推論能力を維持しつつ、推論ステップ数の増加に伴う明確な性能向上を示し、特定のタスクにおいてワンステップ・ドリフトポリシーおよびマルチステップ・拡散ポリシーの両方のベースラインを上回りました。

意義と主張

本論文は、DFM がより柔軟な生成パラダイムを確立すると主張しており、直接生成と反復的洗練はもはや相互排他的な設計選択ではなく、単一のフレームワーク内で見なすことができるとしています。ドリフト生成モデリングとフローベースの反復的生成を接続することにより、DFM は生成プロセスが異なる品質・効率要件に応じてサンプリング計算を適応させることを可能にします。著者らは、このアプローチが、以前のドリフトモデルの主要な制限を克服しつつ、強力なワンステップ生成を達成すると同時に反復的手法のスケーラビリティの恩恵を保持する新たな方法を提供すると強調しています。