A Few-Step Generative Model on Cumulative Flow Maps

技術的概要：累積フローマップを用いた数ステップ生成モデル

問題定義
拡散モデルやフローマッチングを含む生成モデルは、通常、現在の状態 $x(t)$ に条件付けられた局所的な状態更新を予測する瞬間的フローマップ（または瞬間的ダイナミクス）を学習する形式で定式化されます。データを生成するためには、これらのモデルは単純なソース分布から複雑なデータ分布へサンプルを移動させる有限時間輸送である累積フローマップを近似するために、多数の微小時間ステップにわたる反復数値積分を必要とします。この多ステップサンプリングプロセスは、多大な計算コストを伴います。

近年の取り組みでは数ステップおよび単一ステップ生成が探求されてきましたが、既存のアプローチには限界があります：

1. 蒸留ベースの手法は、複雑な教師 - 生徒トレーニングパイプラインを必要とすることが多いです。
2. コンシステンシーモデルや平均フロー（平均速度を学習するもの）は、特定の定式化（例：$u$-予測フローマッチング）向けに主に設計されており、DDIM、EDM、または $x_1$-予測フローマッチングなど、コンピュータグラフィックスで広く使用されている他のフレームワークへ自然に一般化しません。
3. 長距離の累積マップを直接学習することは困難です。なぜなら、標準的な学習目的に必要な条件付き対応物が累積場には存在せず、根本的な学習の障壁が生じるからです。

手法：累積フローマップ（CFM）
著者は、局所的な瞬間的更新と有限時間輸送を、明示的な累積場パラメータ化を介して接続する統合フレームワークである**累積フローマップ（CFM）**を提案します。

• 統合抽象化：この手法は、瞬間的場（例：速度）$m_t(x)$ を用いた抽象関数 $F[m_t(x), x, t, t+h]$ を介して、瞬間的フローマップ $\psi_{t \to t+h}$ を形式化します。この抽象化は、多様なパラメータ化（$u$-FM、$x_1$-FM、DDIM、EDM）を統合します。
• 累積パラメータ化：CFM は、瞬間的マップの合成の極限として累積フローマップ $\psi_{t \to r}$ を定義することで、この概念を長距離輸送へ拡張します。モデルは、$\psi_{t \to r}(x) = F[m_{t \to r}(x), x, t, r]$ となるような累積パラメータ化場 $m_{t \to r}(x)$ を学習します。これにより、モデルは時間 $t$ から任意の未来時間 $r$ へのサンプリングを、単一または数ステップで進行させることができます。
• 場方程式による学習：重要な課題として、自己整合的な方法で条件付き累積場 $m_{t \to r}(x | X_1)$ が存在しないことが挙げられます。これを克服するため、著者は場方程式に基づく再定式化（定理 3）を導出します。彼らは、累積場を瞬間的場とその微分を用いて以下のように表現します：
  $$m_{t \to r}(x) = G(t, r)m_t(x) + H(t, r)E[\dots]$$
  この定式化により、学習可能な条件付き瞬間的場 $m_t(x | X_1)$ を用いて累積場を監督する代替損失関数（式 7）の構築が可能になります。この損失には、ヤコビアン - ベクトル積（JVP）または離散近似を介して計算される微分項（$\partial_t m$ および $\partial_x m$）が含まれます。
• 実装：このアプローチは、最小限のアーキテクチャ変更を必要とします。時間埋め込みを、目標時間 $r$ のための追加エンコーダで拡張し、平均化された埋め込み $(emb_t + emb_r)/2$ を使用します。蒸留や容量の増加は不要です。

主要な貢献

1. 瞬間的 - 累積フローマップ抽象化：本論文は、累積フローマップを、瞬間的マップを合成することで得られる有限時間輸送として形式化し、多ステップおよび数ステップ生成を単一の数学的フレームワークの下で統合します。
2. $u$-予測を超えた一般化：CFM は、平均フローの概念を、$u$-および $x_1$-フローマッチング、EDM、DDIM を含む広範な生成定式化へ一般化します。これにより、幾何分布モデリングやピクセル空間画像生成など、以前の手法（平均フローなど）が適用不可能であった設定において、数ステップ生成が可能になります。
3. モデル非依存の学習：場方程式に基づく目的関数を導出することで、この手法はアーキテクチャ変更や蒸留なしで累積フローマップの学習を可能にし、既存モデルにおけるサンプリングステップを大幅に削減します。

実験結果
著者は CFM を 5 つの多様なグラフィックタスクで評価し、生成品質を維持または向上させながら、大幅な高速化（10 倍～200 倍）を実証しました：

• 画像生成（CelebA-HQ）：CFM-DDIM を使用すると、モデルは 128 ステップのベースラインと同等の FID スコアで 1 ステップおよび 4 ステップ生成を達成し、コンシステンシー蒸留およびトレーニング手法を上回りました。
• 幾何分布モデリング：GeoDist タスクにおいて、CFM-EDM は 60 ステップのベースラインと比較して、チャーマー距離の劣化なしに 6 倍～10 倍の高速化を達成しました。注目すべきは、$u$-予測手法はこのドメインで数ステップ生成をサポートできないのに対し、CFM-EDM は成功したことです。
• 関節位置予測（PDT）：RigNet データセットに適用された CFM-DDIM は、推論を 1000 ステップから 5 ステップに削減（200 倍の高速化）しつつ、元の 1000 ステップ PDT と同等の関節予測精度（CD-J2J、IoU、Precision、Recall）を維持しました。
• 画像条件付きスケッチ生成：ControlSketch データセットにおいて、CFM は 1 または 4 ステップ（50 倍の高速化）でベクトル化されたスケッチを生成し、忠実度（MS-SSIM、DreamSim）は 50 ステップの SwiftSketch ベースラインと一致しました。
• 3D SDF 生成：疎な条件付き SDF 再構成（64 点）において、CFM-$x_1$-FM は、Functional Diffusion に対して同等の再構成品質で 6～16 倍の高速化を達成しました。

意義と主張
本論文は、CFM が画像、点群、陰関数場など多様なデータ表現および生成定式化を含むコンピュータグラフィックス応用に特に適した一般的な数学的フレームワークを提供すると主張しています。その主な意義は以下の点にあります：

• 効率とアーキテクチャの分離：蒸留やアーキテクチャの大規模な見直しを必要とせず、学習目的関数と時間埋め込みのみを変更することで、最大 200 倍の推論加速を実現します。
• 統合フレームワーク：瞬間的ダイナミクスと長距離輸送の間のギャップを埋め、以前の「平均速度」アプローチが失敗したモデル全体にわたる数ステップ生成のための原理的な解決策を提供します。
• 実用性：この手法はモデル非依存であり、事前学習された多ステップモデルを加速するために適用することも、ゼロから新しい数ステップモデルを学習するために使用することもできます。

著者は、CFM が広範なモデルに適用可能である一方で、現在の評価は 5 つの応用と 4 つの代表的な定式化に限定されており、ImageNet などの大規模データセットへの拡張は今後の課題であると述べています。