技術的概要：画像生成における球面フローマッチングのための潜在幾何学の整合化

問題提起

画像生成における潜在フローマッチング（および拡散モデル）は、通常、ガウスノイズをユークリッド空間内の直線的な経路に沿って変分オートエンコーダ（VAE）の潜在空間へ輸送することによって機能します。しかし、このアプローチは、高次元 VAE 潜在空間の 2 つの重要な構造的性質を見落としています：

1. 測度の集中：ガウスノイズの事前分布とエンコードされたデータ潜在変数の両方は、空間の体積を埋めるのではなく、薄い球殻に集中します。これらの殻上の 2 点を結ぶ直線（弦）は、どちらの分布も占めていない内部半径を通り抜けるため、モデルはほとんど占有されていない潜在空間の領域を通る遷移を学習することを強制されます。
2. 方向性の支配：デコーダの出力は、潜在トークンの長さ（半径成分）よりも、その方向（角度成分）に対して圧倒的に敏感です。同じクラスの近傍トークンから半径と方向を交換する実験は、方向を維持しつつ半径を変更しても知覚的な変化は最小限である一方、方向を変更すると意味的なシフトが顕著に生じることを示しています。

標準的な線形フローマッチングはこれらの事実を無視し、トレーニング速度目標の相当部分（トークナイザーによっては 50〜90%）を半径方向の運動に割り当てています。デコーダは半径に対して感度が低いため、この監督は非効率的であり、基礎となる幾何学を尊重する手法と比較して、収束が遅く、画像品質（FID）が最適でない結果をもたらします。

手法

著者らは、拡散アーキテクチャを変更したり、補助エンコーダを必要としたりすることなく、潜在幾何学をデータの球面構造に整合させる球面フローマッチングを提案します。このアプローチは以下の 4 つの主要な構成要素からなります：

1. トークン単位の球面射影：
   この手法は、事前学習済み VAE の各潜在トークン $z$ を、固定半径の超球面 $S^{d-1}(\sqrt{d})$ 上に射影します。これは、凍結されたエンコーダとデコーダの間に挿入された $L_2$ 正規化ステップによって達成されます：
   $$z_{i,j} \leftarrow \sqrt{d} \cdot \frac{z_{i,j}}{\|z_{i,j}\|}$$
   これにより、すべてのトークンが、標準ガウス事前分布の集中半径に一致する半径を持つ球面上に存在することが保証されます。

2. 球面事前分布：
   生ガウスノイズを使用する代わりに、ノイズの終点 $z_0$ は、等方性ガウスノイズ $\epsilon \sim \mathcal{N}(0, I)$ をサンプリングし、それを球面上に射影することで構成されます：
   $$z_0 = \sqrt{d} \cdot \frac{\epsilon}{\|\epsilon\|_2}$$
   これにより、射影されたデータの幾何学と一致する、球面上の一様分布が得られます。

3. デコーダの微調整：
   幾何学的制約に対応するため、エンコーダと射影層は凍結したまま、デコーダを数エポック（実験では 5 エポック）微調整します。再構成目的関数は、標準的なピクセル損失（L1、LPIPS、敵対的損失）を維持しますが、射影された潜在変数が決定論的であるため KL 項は除外されます。

4. Slerp 輸送と接線速度：
   ノイズとデータ端点間の輸送経路は、球面上の測地線弧に従う**球面線形補間（slerp）**に置き換えられます。

   • 測地線経路：経路 $z_t$ はすべての $t \in [0, 1]$ において球面上に留まります。
   • 速度目標：条件付き速度 $u_t$ は、測地線の時間微分であり、接空間 $T_{z_t}S^{d-1}$ 上に射影されます。これにより、速度は純粋に角度方向であり、構成上、半径方向成分がゼロであることが保証されます。
   • トレーニングとサンプリング：モデルはこの接線速度を予測するようにトレーニングされます。推論中は、サンプリャが速度を積分するために指数写像を使用し、各ステップで半径方向の射影補正を必要とせずに軌道が球面上に留まることを保証します。

主要な貢献

本論文は、標準的な潜在フローマッチングにおける「半径 - 殻の不一致」を特定・定量化し、それを解決するための幾何学的介入を導入します：

1. 不一致の特定：著者らは、直線的な経路がデコーダが感度を持たない半径方向の運動をモデルに予測させることで、容量を浪費し収束を遅らせていることを実証します。
2. トークン単位の射影：エンコーダの再トレーニングや補助表現エンコーダの使用を必要とせず、デコーダの微調整のみを使用して事前学習済み VAE 潜在変数を固定半径の球面に制約する手法を導入します。
3. Slerp ベースのフローマッチング：接線速度目標を用いた slerp 測地線に沿ってフローマッチングをトレーニングすることを提案し、軌道全体が球面多様体を尊重することを保証します。
4. 実証的検証：この手法は、アーキテクチャの変更なしに、複数のトークナイザーファミリー（FLUX.2、VA-VAE、REPA-E FLUX.1）およびモデルスケール（SiT-B から SiT-XL）にわたって性能向上を示します。

結果

クラス条件付き ImageNet-256 生成に関する実験は、一貫した改善を実証しています：

• FID の改善：球面-slerp 手法は、すべてのテスト構成において、バニラ線形フローマッチングよりも低い FID スコアを達成します。例えば、FLUX.2 を使用した SiT-B/2 において、FID は 26.35（線形）から 20.55（球面-slerp）に低下します。
• トレーニング効率：この手法は、バニラ線形ベースラインと比較して、目標 FID 30 に到達するために必要なトレーニングステップ数を約2.2$\times$削減します。
• 汎化性：異なるガイダンススケール（CFG）およびモデルサイズ（B および XL）にわたって利益が観察され、この利点が特定のハイパーパラメータ調整ではなく、潜在幾何学そのものに固有であることを示しています。
• アブレーション研究：
  • バニラと球面パイプライン間でデコーダを交換すると、両方向で性能が低下し、学習されたフローが特定の潜在幾何学と密接に結合していることを確認します。
  • バニラ潜在変数に対して「殻分解」された経路（半径と方向を個別に補間する）を使用しても同様の利益は得られず、半径の自由度を完全に除去することが重要であることを確認します。
  • 微調整された球面デコーダの再構成品質（rFID）は競争力のあるままであり、幾何学的制約が VAE の画像再構成能力を著しく損なわないことを示しています。

意義

本論文は、生成モデルにおける潜在サポート幾何学を重要な設計軸として確立します。潜在分布の形状が生成器が通過しなければならない空間を定義し、この幾何学を無視すること（ユークリッド線形性を仮定すること）が非効率性につながることを主張します。

この研究の意義は、画像生成の品質とトレーニング速度を向上させる能力にあります、かつ以下の変更を行わずに達成されます：

• 拡散アーキテクチャの変更（例：SiT バックボーン）。
• 補助エンコーダの追加（凍結された DINOv2 特徴量上のリーマン幾何フローマッチングとは異なり）。
• VAE エンコーダのゼロからの再トレーニング。

既存の事前学習済み潜在変数を単に球面上に射影し、フロー経路を結果として生じる測地線に整合させることで、この手法は最小限の計算オーバーヘッドで最先端の結果を達成します。著者らは、将来の研究において、より強い圧縮領域におけるこれらの幾何学的制約の探求と、この球面構造にネイティブに一致するトークナイザーの設計を提案しています。