Terminal Velocity Matching

本論文は、拡散モデルの初期時刻ではなく終端時刻での挙動を正則化する「Terminal Velocity Matching（TVM）」を提案し、Transformer 構造への最小限の改変と効率的な計算実装を通じて、ImageNet においてワンステップおよび数ステップ生成で最先端の性能を達成したことを報告しています。

要約

この論文は、**「Terminal Velocity Matching (TVM)」**という新しい画像生成の技術を提案しています。

一言で言うと、**「複雑な計算を何十回も繰り返さなくても、たった 1 回（または数回）の計算で、高画質で美しい画像を瞬時に作り出す方法」**です。

これを一般の方にもわかりやすく、いくつかの比喩を使って説明してみましょう。

---

1. 従来の方法：「泥濘（ぬかるみ）を歩く旅」

これまでの AI 画像生成（拡散モデルなど）は、**「霧の中から絵を描く」**ようなものでした。

• 仕組み: 最初は真っ白なノイズ（霧）から始めて、AI が「ここは少し赤くしよう」「ここは形を整えよう」と何度も何度も修正を繰り返します。
• 問題点: 高画質にするには、この修正を50 回〜100 回も繰り返す必要があります。まるで、目的地まで行くために、道中を何度も立ち止まって地図を確認し、方向を修正しながら歩くようなものです。
• 結果: 画像は綺麗ですが、時間がかかりすぎます。

2. 従来の「高速化」の試み：「急ぎ足で走る」

最近、もっと速く描こうとする技術（Consistency Models など）が出てきました。

• 仕組み: 目的地への「最短ルート」を直接覚えさせようとするものです。
• 問題点: しかし、これらは「出発点（ノイズ）」に注目して学習させるため、理論的な保証が弱く、高画質と高速さを両立させるのが難しかったり、学習が不安定になったりしていました。

3. 新しい方法（TVM）：「ゴール地点の速度を合わせる」

この論文の**TVM（Terminal Velocity Matching）**は、全く新しい発想でこの問題を解決しました。

比喩：「スキーのジャンプ」

想像してください。スキーのジャンプ台から飛び出し、着地する瞬間（ゴール）をイメージしてください。

• 従来の考え方: 「飛び出した瞬間（スタート）の勢い」を完璧にコントロールしようとする。
• TVM の考え方: **「着地する瞬間（ゴール）の速度」**に注目する。

「ゴール地点での速度（Terminal Velocity）」を正しく合わせれば、スタートからゴールまでの道筋（軌道）は自然に正しいものになる、というのがこの技術の核心です。

• メリット: 出発点（ノイズ）からゴール（完成画像）までを、**「一瞬でジャンプ」**させることができます。
• 結果: 従来の 50 回の計算が、たった 1 回〜4 回で済みます。まるで、霧の中を歩くのではなく、瞬時にゴール地点にテレポートしたかのような速さです。

4. なぜこれが難しいのか？（技術的な壁）

この「ゴールの速度」を正確に教えるには、AI の脳（ニューラルネットワーク）が**「滑らかさ（リプシッツ連続性）」**を持っている必要があります。
しかし、最新の AI 模型（Transformer）は、この「滑らかさ」が保たれず、学習中に暴走したり不安定になったりしていました。

TVM の解決策：
著者たちは、AI の構造を**「最小限の修正」**で直しました。

• 例えるなら、暴走しやすい車のサスペンションを少し調整し、**「どんな急カーブでも安定して走れるように」**したようなものです。これにより、学習が安定し、高画質を実現できました。

5. 驚異的なパフォーマンス

この技術を実際にテストした結果は驚異的です。

• ImageNet（有名な画像データセット）で:
  • 1 回の計算（1-NFE）: 3.29 というスコア（FID）。これは、従来の方法が 50 回計算しても出せないレベルの画質を、一瞬で達成しました。
  • 4 回の計算（4-NFE）: 2.94 というスコア。これは、従来の最高峰の AI が 500 回計算して出すレベルの画質に匹敵します。

まとめ：何がすごいのか？

1. 超高速: 画像生成が「数秒」から「瞬時」になります。
2. 高画質: 速くても画質は落ちません。むしろ、従来の方法より綺麗になることもあります。
3. シンプル: 特別な複雑な学習プロセス（カリキュラム学習など）が不要で、シンプルに実装できます。
4. 理論的な裏付け: 単なる「試行錯誤」ではなく、数学的に「この方法なら間違いない」という保証（2-ワッサーシュタイン距離の上限）を持っています。

結論として：
TVM は、AI が絵を描くとき、「何回も修正しながら描く」のではなく、「ゴール地点のイメージを完璧に把握して、一発で描き上げる」ための新しい魔法のような技術です。これにより、動画生成やリアルタイムの画像生成が、今後さらに現実的なものになることが期待されています。

技術概要

Terminal Velocity Matching (TVM) の技術的サマリー

本論文は、ICLR 2026 にて発表された「Terminal Velocity Matching (TVM)」という新しい生成モデルの学習フレームワークを提案しています。TVM は、フローマッチング（Flow Matching）を一般化し、単一のトレーニングステージから高忠実度かつ高速な（1 ステップまたは数ステップ）生成モデルを構築することを可能にします。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 背景と課題

従来の拡散モデルやフローマッチングは、高品質な画像生成において支配的なパラダイムですが、高品質な出力を得るためには通常、50 回以上のサンプリングステップ（ODE 積分）が必要であり、推論コストが高いという課題があります。
これを解決するため、Consistency Models や MeanFlow などの「数ステップ推論」を目指す研究が進められてきましたが、以下の課題が残っていました：

• 分布一致の保証の欠如: 多くの手法は軌道の微分を学習するだけで、生成分布とデータ分布の一致に対する理論的な保証（分布レベルの保証）が不足している。
• スケーラビリティの問題: 分布一致を保証する手法（例：IMM）は、学習ステップごとに複数の粒子を必要とし、大規模モデルへの拡張が困難。
• 安定性の問題: 既存の Diffusion Transformer (DiT) 構造はリプシッツ連続性を持たず、高速な 1 ステップ学習において学習が不安定になりやすい。

2. 提案手法：Terminal Velocity Matching (TVM)

TVM は、フロー軌道の初期速度ではなく、**終端速度（Terminal Velocity）**に注目して学習を行う新しい枠組みです。

2.1 核心的なアイデア

• 終端速度の一致: 任意の 2 つの拡散タイムステップ $t$ と $s$ ($s < t$) 間の遷移をモデル化します。従来のフローマッチングが $t=s$ における瞬間的な速度を一致させるのに対し、TVM は軌道の終端（$s$）における速度が、真の速度場と一致することを目的とします。
• 理論的保証: TVM の学習目的関数は、データ分布とモデル分布間の 2-Wasserstein 距離の上限を理論的に保証します（モデルがリプシッツ連続である場合）。これにより、分布レベルの一致が保証されます。
• 単一ネットワークによる学習: 瞬間的な速度場 $u_\theta$ と、任意の時間間隔での変位マップ（1 ステップ生成）$f_\theta$ の両方を、単一の 2 時刻条件付きニューラルネットワーク $F_\theta(x_t, t, s)$ でパラメータ化して学習します。

2.2 学習目的関数

TVM の損失関数は、以下の 2 つの項を同時に最小化します：

1. 終端速度誤差: モデルが予測する終端速度 $\frac{d}{ds}f_\theta(x_t, t, s)$ が、モデル自身で推定した速度場 $u_\theta$ と一致すること。
2. フローマッチング誤差: 境界ケース（$t=s$）において、従来のフローマッチング損失が満たされること。

これにより、1 ステップ生成と多ステップ生成の両方を単一のモデルで効率的に学習できます。

2.3 実装上の工夫と課題解決

TVM を実用的かつ大規模に適用するために、以下の技術的革新がなされています。

• 半リプシッツ制御（Semi-Lipschitz Control）:
  • 問題：標準的な Transformer（Scaled Dot-Product Attention + LayerNorm）はリプシッツ連続性を持たず、TVM の学習を不安定化させます。
  • 解決：最小限のアーキテクチャ変更として、QK-Norm に RMSNorm を採用し、AdaLN（Adaptive LayerNorm）のモジュレーションパラメータにも RMSNorm を適用することで、リプシッツ性を制御し、学習の安定化を図りました。
• JVP 対応 Flash Attention カーネル:
  • 問題：TVM の損失計算にはヤコビアン・ベクトル積（JVP）の計算が必要であり、特にそのバックワードパス（勾配伝播）が Transformer において計算コストとメモリ消費のボトルネックとなります。
  • 解決：JVP のフォワードパスとバックワードパスを融合した新しい Flash Attention カーネルを開発。これにより、PyTorch の標準実装と比較して最大 65% の高速化と大幅なメモリ削減を実現しました。
• スケーリングパラメータ化と CFG 学習:
  • Classifier-Free Guidance (CFG) 重み $w$ に対して、ネットワーク出力を $w$ に比例してスケーリングするパラメータ化を導入。
  • 学習中に CFG 重みをランダムにサンプリングし、勾配爆発を防ぐために損失を $1/w^2$ で重み付けすることで、安定した学習を可能にしました。

3. 実験結果

ImageNet データセットにおける評価結果は、既存の手法を大きく上回る性能を示しています。

• ImageNet-256×256:
  • 1-NFE (1 ステップ): FID 3.29（MeanFlow の 3.43 を上回る）。
  • 4-NFE (4 ステップ): FID 1.99（従来の拡散モデル DiT の 2.27 を上回る）。
• ImageNet-512×512:
  • 1-NFE: FID 4.32。
  • 4-NFE: FID 2.94（DiT の 3.04 を上回る）。
• 特徴: 再学習なしで 1 ステップから多ステップまで自然に補間可能であり、トレーニングカーリキュラムや損失関数の複雑な修正を必要としません。

4. 主要な貢献

1. 理論的枠組みの確立: フロー軌道の「終端速度」に焦点を当てた新しい学習目標を提案し、それが 2-Wasserstein 距離の上限となることを証明。これにより、分布一致の保証を持つ単一ステップ生成モデルの理論的基盤を提供しました。
2. 大規模な実用化: Diffusion Transformer のリプシッツ連続性の欠如という構造的な課題を解決し、JVP 対応の Flash Attention を実装することで、大規模モデルでの安定した学習と効率的な推論を実現しました。
3. SOTA 性能の達成: ImageNet において、ゼロから学習した 1 ステップ/数ステップモデルとして最高水準の性能を達成し、従来の多ステップ拡散モデルと同等以上の品質を少ない計算コストで実現しました。

5. 意義と将来展望

TVM は、生成モデルの分野において「高品質・高速・スケーラビリティ」を同時に達成するための重要な進展です。

• 理論と実践の融合: 分布一致の理論的保証を持ちながら、実用的な実装（単一ステージ学習、単一ネットワーク）を可能にしました。
• 効率化: 推論ステップ数を劇的に削減できるため、動画生成やリアルタイム応用など、計算リソースが限られる分野での応用が期待されます。
• 今後の展望: 本手法は、より大規模なモデルや高解像度データ、さらには動画生成への拡張において、新しい標準となり得る可能性を秘めています。

総じて、TVM は生成モデルの設計原則に新たな洞察を与え、理論的に裏付けられた設計が実用的な性能向上につながることを示した画期的な研究です。