技術的サマリー：ワッサーシュタイン勾配流によるワンステップ生成モデリング

問題定義

現代の生成モデル、特に拡散モデルやフローベース手法は、最先端のサンプル品質を達成していますが、多段階の反復サンプリングに依存しているため、高い計算コストと遅延という課題を抱えています。ワンステップまたは少数ステップの蒸留手法は存在しますが、それらは往々にして反復生成用に設計された経路からの学習信号を継承するか、収束保証が欠如しモード崩壊を起こしやすいヒューリスティックな力学（例：Drifting Models）に依存しています。中心的な課題は、安定した収束とグローバルなモードカバレッジを確保しつつ、参照分布からターゲットデータ分布への直接的な単一ステップマッピングを生成器が学習できるような学習力学を設計することです。

手法：W-Flow

著者は、学習力学の設計と生成器のパラメータ化を分離するフレームワークW-Flowを提案します。中核となるアイデアは、モデル分布のターゲット分布への進化を**ワッサーシュタイン勾配流（WGF）**を用いて規定し、その後、静的なニューラルネットワークを訓練してこの多段階進化を単一の推論ステップに圧縮することです。

1. ワッサーシュタイン勾配流による学習力学

ノイズ除去ステップや速度場を直接学習するのではなく、W-Flow は学習軌跡を確率測度の空間における連続時間進化として定義します。

• エネルギー汎関数: 進化はエネルギー汎関数 $F(q) = D(q \| p)$ の最急降下によって駆動されます。ここで、$q$ は現在のモデル分布、$p$ はターゲットデータ分布です。
• 速度場: 連続時間速度場 $V_t$ は、エネルギー汎関数の第一変分から導出されます：$V_t(x) = -\nabla \frac{\delta F}{\delta q}(q_t)(x)$。
• 離散化: 連続流は陽的オイラー法を用いて離散化されます。各学習ステップ $k$ において、モデル分布 $q^{(k)}$ は局所輸送マップ $g^{(k)}(x) = x + \eta V^{(k)}(x)$ によって更新されます。ここで $\eta$ はステップサイズです。

2. Sinkhorn 発散による具体化

速度場が高次元において計算的に扱いやすく、幾何学的に良好な振る舞いをするようにするため、著者はエネルギー汎関数をSinkhorn 発散（$S_\epsilon$）、すなわちエントロピー正則化最適輸送（OT）メトリックを用いて具体化します。

• 速度場の定式化: 誘起される速度場は、2 つの重心射影の差によって与えられます：
  $$V^\epsilon_{q,p}(x) = T^\epsilon_{q,p}(x) - T^\epsilon_{q,q}(x)$$
  ここで、$T^\epsilon_{q,p}$ は生成分布 $q$ からターゲット $p$ への質量輸送を表し、$T^\epsilon_{q,q}$ は現在の分布との整合性を強制する自己輸送項を表します。
• グローバルな調整: 局所カーネル相互作用（MMD や Drifting Models など）とは異なり、Sinkhorn 速度場はグローバルな OT プランから導出されるため、粒子の移動がバッチ全体で調整され、モード崩壊が防止されます。
• 2 バッチ推定量: 粒子がゼロコストで自分自身にマッチする自己マッチングアーティファクトを回避するため、著者は2 バッチ戦略を採用します。自己輸送項 $T^\epsilon_{q,q}$ は、単一バッチのコスト行列の対角成分をマスクするのではなく、現在の分布からの独立した第 2 バッチのサンプルを用いて計算されます。

3. 生成器の訓練

生成器 $f_\theta$ は、WGF によって定義された離散粒子更新を模倣するように訓練されます。

• 目的関数: 損失関数は、生成器の出力と Sinkhorn 更新から導出された「ターゲット」粒子位置との二乗距離を最小化します：
  $$\mathcal{L}(\theta) = \mathbb{E}_{z \sim p_{ref}} \| f_\theta(z) - \text{sg}(f_\theta(z) + \eta V^\epsilon(\hat{q}_\theta, \hat{p})) \|^2$$
  ここで、$\text{sg}(\cdot)$ は停止勾配演算子を示し、更新された粒子位置を固定された回帰ターゲットとして扱います。
• 分類器フリーガイダンス（CFG）: このフレームワークは、速度場を直接修正することで CFG をサポートします。分布を線形混合するのではなく、著者は速度ガイダンスを提案します。これは、条件付き速度場と無条件速度場の差に比例する項を加えるものです。このアプローチは、KL 発散を使用する際に標準的な指数傾斜ターゲット分布を回復し、Sinkhorn 発散に対して幾何学的に整合性のある拡張を提供することが示されています。

主要な貢献

1. 原理的なワンステップフレームワーク: W-Flow は、学習力学をワッサーシュタイン勾配流に基づけることで、Drifting Models などのヒューリスティックな手法には見られない明確な変分解釈を提供し、ワンステップ生成モデリングのための新たなパラダイムを確立します。
2. Sinkhorn に基づく力学: Sinkhorn 発散による WGF の具体化は、効率的で最適輸送に基づく更新規則を提供し、カーネルベースまたは局所相互作用手法と比較して、グローバルな分布の不一致を捉え、モードカバレッジを向上させます。
3. 理論的収束: 著者は、適切な正則性仮定の下で、W-Flow の有限サンプル粒子力学が連続時間分布力学に収束することを証明し、学習過程に対する理論的根拠を提供します。
4. 2 バッチ推定量: 自己輸送項を推定するための新しい 2 バッチ戦略が導入され、単一バッチ推定量や対角マスクヒューリスティックに内在する自己マッチングバイアスを排除します。

実験結果

著者は、ImageNet 256×256およびFFHQにおいて W-Flow を評価しました。

• 最先端のパフォーマンス: W-Flow は、ImageNet におけるクラス条件付きワンステップ生成において、XL スケールで FID 1.29、L スケールで FID 1.35という新しい最先端の性能を達成しました。これは、Drifting Models や MeanFlow などの以前のワンステップ手法を大幅に凌駕し、LightningDiT や SiT などの多段階拡散モデルをも上回っています（ただし、関数評価回数（NFE）は 1 回のみ）。
• サンプリング効率: 同様の FID スコアを持つ多段階拡散モデルと比較して、W-Flow は単一 H100 GPU 上で約100 倍高速なサンプリングを実現します（例：SiT-XL/2 対して 77.88 画像/秒 vs 0.93 画像/秒）。
• モードカバレッジ: 不均衡な分布（例：95% 高齢者の顔、5% 子供の顔）を用いた実験において、W-Flow は少数派のモードを正常に捉えますが、Drifting Models はモード崩壊を起こし、少数派クラスを完全に無視します。
• 収束速度: W-Flow は訓練中に著しく速く収束し、Drifting Model を 1280 エポック訓練した場合の性能をわずか 384 エポックで達成します。
• ドメイン転送: この手法は、ドメイン転送タスク（例：FFHQ における年齢変換）において強力な能力を示し、ヒューリスティックなベースラインと比較して、より短く、より一貫性のある輸送軌跡を生成します。

意義と主張

本論文は、ワッサーシュタイン勾配流が、高速かつ高忠実度な生成モデリングのための原理的かつ効果的な基盤を提供すると主張しています。学習力学の設計と生成器のパラメータ化を分離することで、W-Flow は、単一ステップ生成器が複雑な多段階反復プロセスに匹敵する性能を達成するように訓練可能であることを示しています。著者は、最適輸送力学によって提供されるグローバルな調整が、訓練を安定させ、包括的なモードカバレッジを確保する鍵であり、既存のワンステップ手法の重要な限界を克服していると強調しています。結果は、リアルタイムのクリエイティブツールやシミュレーションなど、低遅延かつ高品質な生成を必要とするアプリケーションにおいて、W-Flow がスケーラブルな代替手段となり得ることを示唆しています。