技術要約：Drifting Transformerの蒸留におけるRepresentation Autoencoder（Drift-RAE）

問題提起

Representation Autoencoder（RAE）は、学習済みDINO特徴量から得られる、ラベルごとにクラスター化された意味的に豊かな潜在空間を利用することで、フローベースの生成モデルの性能を大幅に向上させてきましたが、これらのモデルを効率的な1ステップまたは数ステップの生成器へと蒸留することは依然として困難です。

著者らは、根本的な幾何学的ミスマッチを特定しています。すなわち、RAEの潜在空間は、従来のVAE空間と比較して、その基礎となる常微分方程式（ODE）の軌跡において深刻な異方性（anisotropy）と大きな曲率（curvature）を示します。既存の軌跡ベースの蒸留手法は、滑らかでほぼ直線的なティーチャー（教師）の軌跡を前提としています。RAE空間における高度に湾曲した軌跡は、これらの手法を不安定にするか、あるいは非効率にします。さらに、最近提案されたDrifting Modelは、曲がった軌跡のマッチングを回避する分布レベルのアプローチを提供しますが、その元の定式化では、集中した特徴量を提供するために、補助として別途訓練されたMasked Autoencoder（MAE）を必要とし、多大な計算オーバーヘッドを導入します。

手法

本論文は、補助的な特徴抽出器を必要とせずに、RAE潜在空間に特化してDrifting Modelのパラダイムを適応させた蒸留フレームワークであるDrift-RAEを提案します。

1. 幾何学とダイナミクスの理論的分析

著者らはまず、潜在空間を定量的に分析し、RAEの特徴量がSD-VAEの特徴量よりも大幅に異方的（低いParticipation RatioおよびSpectral Entropy）であることを示しました。彼らは、高次元空間において正例が過度に分散している場合、引き込み成分（attractive component）のドリフト場が消失し、退化した場を招くことを理論的に証明しました（定理1）。

• 重要な知見: 元のDrifting Modelは、分散した特徴を集中させるためにMAEを必要としますが、RAEの潜在空間は自然に意味的に集中しています。これにより、Drifting Modelは補助的なMAEモジュールなしで、直接RAE潜在空間上で動作することが可能になります。

2. Drift-RAEフレームワーク

本手法は、学習済みのフローモデル $v_\theta$ を、1ステップの生成器 $G_\theta(z) = z - v_\theta(z, 1)$ へと蒸留します。ODEの軌跡を一致させる代わりに、生成器は、生成されたサンプルを実データ分布へと押し込み、他の生成されたサンプルから遠ざけるドリフト目的関数を最小化するように訓練されます。

• 目的関数: ロスは、正例（実データ）と負例（生成データ）から導出された正規化されたドリフト場 $\tilde{V}$ を用いて、RAE潜在表現上でトークン単位で計算されます。

3. 安定性のための実践的な修正

Drifting ModelをDiffusion-GANの枠組みに適合させ、訓練の安定性を向上させるために、著者らは3つの具体的な修正を導入しています。

• Softmaxの次元: 著者らは、元のDrifting Modelで使用されていた双方向のsoftmaxトリックを削除しました。代わりに、サンプルインデックスに対する単方向のsoftmaxを保持しており、これにより、基礎となるlog-sum-expポテンシャルの勾配方向とより良く一致します。
• 入力摂動: 実データと生成データの間の非交差的な多様体によって引き起こされる勾配消失を防ぐため、生成されたサンプルには、ドリフト場を計算する前にラプラス分布からサンプリングされたノイズによる摂動が加えられます。
• 負例の部分的なデタッチ（Detachment）: 生成分布の近似を行う負例の数を、GPUメモリ予算を超えることなく増やすために、著者らは、より大きな負例プール（$N_{total\_neg}$）を使用し、そのうちのサブセット（$N_{neg}$）のみが勾配を保持し、残りはデタッチされる手法を提案しています。これにより、生成分布の経験的な近似が改善されます。

主な貢献

1. 幾何学的分析: RAE潜在空間が高曲率かつ異方的であり、標準的な軌跡ベースの蒸留を不安定にすることを定量的に示したこと。
2. 理論的正当化: ドリフトモデルが過度に分散した空間では失敗すること（元の定式化においてMAEが必要となる理由）を示すとともに、RAEの自然に集中した意味空間においては有効であることを示したこと。
3. Drift-RAE: 補助的なMAE特徴抽出器を必要とせずに、RAE空間において最先端の1ステップ生成を実現する新しい蒸留手法。
4. 安定性の強化: Diffusion-GANの原理に適合させ、収束を改善するための、Drifting目的関数（softmaxの次元、入力摂動、およびサンプルのデタッチ）の修正に関する理論的および経験的な洞察。

実験結果

本手法は、DiTDH-XLおよびDiTDH-Lバックボーンを用いて、ImageNet 256 $\times$ 256データセットで評価されました。

• 性能: Drift-RAEは、単一の生成ステップを用いて、DiTDH-XLでFID 1.77、DiTDH-Lで2.12を達成しました。
• 効率性: 本モデルは、わずか10,000蒸留ステップ（約16エポック）でこれらの結果に到達します。
• 比較:
  • RAE空間における従来の蒸留手法（例：MF-RAE）を上回っています。
  • 計算コストの高い補助MAEを使用していないにもかかわらず、元のDrifting Modelと同等のFIDおよび改善されたFDDINOv2スコアを達成しています。
  • 他の潜在空間（例：ピクセル空間、SD-VAE）で訓練された1ステップおよび数ステップの生成器とも競争力のある性能を維持しています。

意義と主張

本論文は、Driftingが生成モデルを潜在空間で蒸留するための効果的かつ有望なフレームフレームワークであることを主張しています。主な意義は、RAEの固有の意味的集中を利用することで、RAEベースのモデルの蒸留における不安定性を解決し、補助的な特徴抽出器への依存を排除した点にあります。

著者らは、理論的分析が簡略化された高次元モデルに基づいていることや、RAE空間において（蒸留なしで）Drifting Modelをゼロから訓練することは依然として未解決の問題であることを認め、限界についても謙虚に述べています。また、多くの同クラスの正例を必要とすることがテキスト・トゥ・イメージ生成におけるスケーラビリティの課題となる可能性を指摘しており、これを今後の研究方向として示唆しています。