Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers

技術要約：クロススペース蒸留（Cross-Space Distillation）

問題提起
現代の一ステップ拡散モデルは、分布ベースのタイムステップ蒸留を通じて高品質な生成を実現しています。しかし、これらの手法は、教師モデル（Teacher）と生徒モデル（Student）が同一の潜在空間に存在し、潜在解像度やVAEのパラメータ化を共有するという、制限的な**共有空間制約（Shared-Space constraint）**に依存しています。この仮定により、高解像度（例：1024×1024）で高度なVAEを利用する現代の高容量な教師モデル（例：Stable Diffusion 3.5、Flux、Kolors）から、低解像度（例：512×512）でコンパクトかつ展開に適した生徒モデル（例：Stable Diffusion 1.5）への知識転送が妨げられています。

標準的な蒸留目的関数は、以下の理由により、このような設定では機能しません。潜在テンソル（$z_S$ および $z_T$）が本質的に不適合であるためです：

1. クロス解像度（Cross-Resolution）： 空間グリッド形状の不一致（例：$h \times w$ 対 $H \times W$）。
2. クロスVAE（Cross-VAE）： オートエンコーダーのパラメータ化の違いによる潜在空間の不整合。
3. クロスアーキテクチャ／メカニズム（Cross-Architecture/Mechanism）： UNetとDiTバックボーンの違い、あるいは標準的な拡散モデルとフローマッチングの定式化の違い。

ネットワークプルーニングや、新しいコンパクトなバックボーンをゼロから構築するといった代替手法は、硬いスパース性閾値、広範な再学習の必要性、およびエコシステムとの互換性の喪失といった課題に直面します。

手法：クロススペース蒸留と「Bridge」
著者らは、この設定を**クロススペース蒸留（Cross-Space Distillation）として定式化し、生徒のバックボーンを変更することなく、生徒の潜在変数を教師の潜在空間へとマッピングするように設計された軽量モジュールであるBridge ($B_\phi$)**を導入しました。

• アーキテクチャ： 効率的に不一致を処理するために、Bridgeは2つのステージに分解されています：

  1. 空間的事前分布（Spatial Prior / 固定）： 生徒のVAEデコーダーの最初の $n$ 個のデコーディングブロック（$D_S^{(n)}$）を固定します。これらのブロックは、事前学習済みの空間アップサンプラーとして機能し、生徒の潜在変数 $z_S$ を低解像度から教師の高解像度（$H \times W$）へと拡張します。これにより、アップサンプリングのための学習パラメータを追加することなく、生徒が持つ既存の空間的知識を活用できます。
  2. 学習可能なプロジェクター（Learnable Projector $g_\phi$）： コンパクトで学習可能なネットワーク（SwinIRとして実装）が、アップサンプリングされた生徒の特徴空間から教師の潜在表現空間へとマッピングし、教師と互換性のある潜在変数 $\hat{z}_T$ を出力します。

• 学習目的関数： Bridgeは、同じ画像を両方のVAEでエンコードすることによって得られる、ペアとなった潜在変数 $(z_S, z_T)$ の合成データセットを用いて学習されます。最適化は、以下の2つの損失の加重和を最小化します：

  1. 潜在再構成損失（Latent Reconstruction $L_{rec}$）： 予測された潜在変数 $\hat{z}_T$ と正解の教師潜在変数 $z_T$ の間の $\ell_1$ 損失です。著者らは、この文脈において $\ell_1$ が $\ell_2$ よりも優れた安定性と構造保存性を提供することを指摘しています。
  2. アテンション忠実度（Attention Fidelity $L_{attn}$）： $\hat{z}_T$ を入力した際の、凍結された教師デノイザーの自己アテンション分布と、$z_T$ を入力した際の分布との間の逆KLダイバージェンス損失です。これにより、マッピングされた潜在変数が、真の教師潜在変数と同じ内部構造的応答（長距離依存性）を誘発するようにし、ピクセル単位の損失では捉えきれない微細な詳細を保持します。

• 蒸留と推論： 学習後、Bridgeは生徒にアタッチされます。蒸留中、生徒の出力はBridgeを通過して $\hat{z}_T$ を生成し、これにより標準的な分布ベースの目的関数（例：Variational Score Distillation、Adversarial Distillation）を教師に対して直接適用できるようになります。決定的なのは、生徒のバックボーンは変更されないという点です。

主な貢献

1. クロススペース蒸留の定式化： 本論文は、異なる潜在解像度およびVAE空間をまたいで蒸留が行われる新しいレジームを定義し、現代の高解像度モデルからコンパクトなレガシーモデルへの知識転送における重要なギャップを解決しました。
2. Bridgeモジュール： 生徒のアーキテクチャを変更することなく、標準的な蒸留目的関数を可能にする、軽量でモジュール化されたインターフェースです。これは、固定された空間的事前分布と学習可能なプロジェクターを組み合わせ、安定したアライメントのためにアテンション忠実度目的関数を利用しています。
3. 実用的な有効性： この手法により、コンパクトな生徒（例：SD 1.5）が、一ステップの推論、低レイテンシ、および広範なエコシステムの互換性を維持しながら、多様な現代的教師（SDXL、Kolors、PixArt-Σ、Flux、SD 3.5）の視覚的事前分布を継承することを可能にします。
4. モデルマージ能力： この手法は、多様な教師を同じバックボーンと潜在空間を持つ生徒へと蒸留するため、得られたチェックポイントはパラメータ平均によるマージが容易であり、複数の教師ソースからの知識を集約した単一のモデルを作成できます。

結果
実験により、コンパクトな一ステップ生徒モデルの生成品質が大幅に向上することが示されました：

• 定量的な利得： Bridgeを介してSD 3.5 Medium（教師）をSD 1.5（生徒）に蒸留することで、HPSv3スコアが5.37から9.42へ、ImageRewardが-0.29から0.62へと向上しました。5つの異なる教師からのチェックポイントをマージすることで、SD 1.5のHPSv3はさらに10.53まで上昇しました。
• 定性的な改善： 視覚的な比較において、Bridgeを用いた生徒は、初期化されたベースラインやプルーニングされたモデルと比較して、より鮮明なテクスチャ、強い構造的一貫性、および少ないアーティファクトを生成することが示されました。
• 効率性： Bridgeはわずかなオーバーヘッド（約5Mパラメータ）しか追加せず、一ステップの推論を可能にします。また、推論時の解像度アップグレードも可能であり、512×512の生徒のサンプルを、生徒を再学習させることなく、Bridgeを通じて1024×1024にマッピングし、教師のデコーダーでデコードできます。
• プルーニングとの比較： Bridgeを用いた蒸留は、ネットワークプルーニング（例：SD 3.5を30%のスパース性にプルーニング）を上回る性能を示しましたが、プルーニングはぼけやアーティファクトを生じやすい傾向があります。

意義
本論文は、クロススペース蒸留が、大規模で現代的な教師の生成能力を、新しいアーキテクチャをゼロから訓練するコストなしに、効率的で展開可能なバックボーンへと転送する可能性を解き放つものであると主張しています。異種混合の潜在表現のアライメントを再利用可能なインターフェース問題として扱うことで、本研究は以下を可能にします：

• エコシステムの互換性： SD 1.5のようなコンパクトなモデルが、既存のツールやハードウェアとの互換性を維持したまま、最先端の品質へとアップグレードできます。
• 知識の統一： 構造的に異なる複数の教師（例：FluxとSD 3.5）の知識を、単一の軽量な生徒モデルへと統合する能力。
• モジュール性： Bridgeは、蒸留と推論時の解像度向上（Resolution Enhancement）の両方に再利用可能なプラグアンドプレイのコンポーネントとして機能し、生徒の生成能力をその潜在表現の制約から切り離します。

著者らは、この研究を、テキスト・トゥ・イメージを超えてビデオやマルチモーダル生成にも適用可能な、ヘテロジニアス（異種混合）な潜在空間をまたぐ蒸留のより広範な探求への一歩として位置付けています。