技術的サマリー：潜在フローマッチングにおける安定性の探索と活用

1. 問題定義

生成モデル、特に拡散モデルとそのフローマッチング（FM）の対応モデルの学習は、依然として計算コストが高価である。これらのモデルは通常、膨大なデータセット、広範な計算リソース、長い学習期間、および条件付き設定においては広範な注釈を必要とする。フローマッチングは、確率微分方程式（SDE）と比較してサンプリングステップを削減する傾向がある常微分方程式（ODE）を介したサンプル生成のための決定論的定式化を提供するが、フルスケールのデータセットでこれらのモデルを学習させるコストは依然としてボトルネックとなっている。

この研究が扱う中心的な問いは、生成品質を犠牲にすることなく、データセットのサイズとモデル容量を大幅に削減できるかどうかである。具体的には、著者らは、データセットのサブサンプリング（プルーニング）、アーキテクチャの変更、および変更された学習構成を含む主要な摂動下での潜在フローマッチング（LFM）モデルの安定性を調査する。彼らは、LFM の輸送軌道が本質的に安定している場合、この性質を利用してより効率的な学習および推論アルゴリズムを導き出せる可能性を提唱する。

2. 手法

2.1. 理論的基盤：フローマッチングにおける安定性

この研究は、整流化フロー（Rectified Flow）が速度場に対して閉形式の式を許容するという観察に基づいている。有限のデータセットにおいて、最適な速度場は学習サンプルの重み付き和であり、その重みは現在の状態とスケーリングされた学習サンプルとの間の二乗距離に対するソフトマックス関数によって決定される。著者らは、このソフトマックスが輸送段階の早期に「鋭く」なることを指摘しており、つまり、特定の軌道に対して単一の学習サンプルがしばしば速度場を支配することを意味する。

この理論的洞察は、輸送軌道が特定の「支配的」サンプルの存在によって主に決定されることを示唆している。したがって、非支配的サンプルを除去（プルーニング）することは、大部分の軌道を変化させないはずである。著者らはこれを形式化し、割り当て指標を定義する：ノイズ $x_0$ から始まる軌道について、それがどの学習サンプル $x_i$ で終了するかを決定する。彼らは、その軌道の支配的サンプルが保持されている限り、データセットの大部分が除去された場合でも、この割り当ては不変であると仮説を立てる。

2.2. データプルーニング戦略

この安定性を利用するために、著者らは判別モデルから適応されたが、FM 目的（共有ノイズ経路とタイムステップを利用する）向けに再定式化された 3 つのデータプルーニング基準を提案する：

1. 勾配ベースのスコアリング（G）：サンプルを損失の勾配ノルムの大きさでランク付けする。
2. 損失ベースのスコアリング（L）：サンプルを共有ノイズ経路に沿って計算された損失の大きさでランク付けする。
3. クラスタリングベースのスコアリング（C）：サンプルを事前学習された埋め込み空間（例：CLIP）でクラスタリングする。著者らは 2 つの具体的な変種を提案する：
   • 比例（$C_p$）：クラスタサイズに比例してサンプルを選択する。
   • バランス型（$C_b$）：分布の偏りを軽減するために、各クラスタから等しい数のサンプルを選択する。
   • カーネルベース（$C_b^\kappa$）：潜在空間におけるクラスタのガウスカーネルの平均に一致するように、貪欲にサンプルを選択する。

著者らはまた、安定性の主張の堅牢性をテストするために、これらの基準の「逆」（最低スコアを選択する）も評価する。

2.3. 粗から細（C2F）推論

モデル容量全体にわたる軌道の安定性を活用し、著者らは 2 段階の推論戦略を提案する：

1. 粗モデル：軽量アーキテクチャ（例：DiT-S/2）をプルーニングされたデータサブセットで学習させる。これは軌道の最初の段階（$t \in [0, t_0)$）を処理する。
2. 細モデル：フルデータセットで事前学習された高容量アーキテクチャ（例：DiT-XL/2）が、2 番目の段階（$t \in [t_0, 1]$）を処理する。

「継ぎ目」の点 $t_0$ で滑らかな遷移を確保するために、粗モデルは連続性損失を用いて微調整される。この損失は、$t_0$ における粗モデルと細モデルの速度場間の差を最小化する。細モデルの軌道は、ODE 反転（後方積分）を介して目標状態 $x_{t_0}$ を生成するために使用され、これにより粗モデルが細モデルの経路に整合するように学習することが保証される。

3. 主要な貢献

1. LFM 安定性の実証的証明：著者らは、独立して学習された LFM モデルが、データセットのサブサンプリング、アーキテクチャの変更、および異なるランダム初期化の下で、極めて一貫性のある輸送軌道に収束することを示す。この安定性は、FM 目的の理論的性質と関連付けられている。
2. 実用的なプルーニングアルゴリズム：彼らは FM に適した 3 つのプルーニング基準（勾配、損失、クラスタリング）を導入する。彼らは、積極的なプルーニング（例：データの 50〜95% を除去）が性能を劣化させず、むしろ限られた計算予算下で収束速度と生成品質を向上させる可能性を実証する。
3. 粗から細への推論加速：彼らは軌道を分割する 2 モデル生成戦略を提案する。このアプローチは、安定性の仮定がモデルアーキテクチャ全体で成り立つ場合、生成品質を維持しながら、大幅な推論速度向上（約 2.15 倍）を達成する。
4. バランス型プルーニングによるバイアス軽減：著者らは、バランス型クラスタリングプルーニング（$C_b$）が、明示的なラベル依存のバランス調整を必要とすることなく、学習セットにおけるセマンティッククラスタの表現を均等化することで、生成データにおけるバイアス（例：性別、肌色）を軽減できることを示す。

4. 実験結果

手法はCelebA-HQ、FFHQ、およびImageNet-1Kで評価された。

• プルーニング下での安定性：CelebA-HQ において、50% でプルーニングされた非重複サブセットで学習されたモデルは、同じノイズシード $x_0$ から開始した場合、視覚的に類似した画像を生成した。プルーニングモデルと非プルーニングモデル間の ArcFace コサイン類似度は 0.79 から 0.83 の範囲にあり、ランダムなペア（0.37）よりも著しく高かった。
• 性能指標：
  • CelebA-HQ：バランス型クラスタリング（$C_b$）を使用してデータの 50% をプルーニングした結果、フリーchet 開始距離（FID）はプルーニングなしの 24.24 から 22.80 にわずかに改善された。
  • ImageNet：データの 75% をプルーニング（$pr=0.75$）した結果、収束が速くなり、フルデータセットと比較して FID が維持、またはわずかに改善された。95% のプルーニングでも、モデルは初期段階でより速く収束したが、一定数の反復後は性能が劣化した。
• 摂動に対する堅牢性：
  • アーキテクチャ：モデル容量を増加させる（DiT-S/2 から DiT-XL/2 へ）ことは、軌道の変化を最小限に抑えた（類似度 $s=0.81$）。U-Net アーキテクチャへの切り替えは類似度を低下させた（$s=0.55$）が、粗い属性は保持された。
  • 潜在空間：安定性は、潜在空間のパラメータ化の変更（例：VAE シードの変更または可逆変換の適用）に対して敏感であり、これにより安定性が破綻した。
  • 目的関数：フローマッチングからスコアベース拡散への切り替えは軌道を完全に変化させ、安定性が FM 目的に特有であることを示した。
• 粗から細への効率性：CelebA-HQ において、$t_0=0.7$ の C2F アプローチは、FID を維持しつつ、推論時間を 93.95 ms/画像（細モデルのみ）から 43.53 ms/画像（約 2.15 倍の速度向上）に削減した。
• 公平性：バランス型プルーニング（$C_b$）は、ラベル非依存であるにもかかわらず、属性（性別、年齢、肌色）全体にわたる一様分布への KL 発散を、プルーニングなしのモデルよりも効果的に削減した。

5. 意義と主張

著者らは、観察された安定性が、特定の学習アルゴリズムのアーティファクトではなく、固定された潜在表現下での整流化フローマッチングパラダイムの特徴的な性質であると主張する。

• 効率性：主な意義は、品質を犠牲にすることなく学習時間とデータ要件を削減できる点にある。著者らは、ImageNet のような大規模データセットではプルーニングが収束を加速し、CelebA-HQ のような小規模データセットでは生成品質を向上させることを実証する。
• 推論速度向上：C2F アプローチは、生成の初期の粗い段階を軽量モデルにオフロードし、軌道の安定性に依存して高容量モデルが最終的な詳細のみを微調整すればよいようにすることで、推論コストを削減するための原理的な方法を提供する。
• バイアス制御：この研究は、データセットプルーニング、特にバランス型クラスタリングによるものが、複雑なラベル依存の再重み付けを必要とすることなく、生成データにおける分布の偏りを制御および軽減するメカニズムとなり得ることを示唆しており、より公平な生成モデルへの道筋を提供する。

本論文は、安定性がデータとモデル容量の変化に対して堅牢である一方で、潜在空間の構造と生成目的自体の変化には敏感であることを結論付けている。この洞察により、実践者は高忠実度の生成を維持しながら、よりリソース効率の良いパイプラインを設計することが可能となる。