技術要約: NanoFLUX

問題提起

FLUXファミリー（最大17Bパラメータ）やQwen-Image（20Bパラメータ）といった大規模なテキストからの画像生成（T2I）拡散モデルは、最先端の視覚的品質を達成しています。しかし、その巨大なスケールはオンデバイス展開への大きな障壁となっており、高価なGPUインフラやクラウドベースの推論を必要とします。既存のオンデバイス・ソリューションは、同等の品質を持たない小規模なモデルに依存するか、あるいは（プルーニング、量子化、ステップ削減などの）圧縮技術を採用していますが、これらは大規模な拡散トランスフォーマー（DiT）およびそれに関連するテキストエンコーダーのアーキテクチャ上の冗長性を完全には解決できていません。課題は、生成品質を維持しつつ、モバイルデバイス向けにこれらの巨大なモデルをいかに圧縮し、推論レイテンシを最小限に抑えるかという点にあります。

手法: NanoFLUX

著者らは、17BのFLUX.1-Schnell（ティーチャーモデル）から蒸留された2.4BパラメータのT2IモデルであるNanoFLUXを提案しています。この圧縮は、最も重い2つのコンポーネント、すなわち12Bの拡散トランスフォーマー（DiT）と5BのT5-XXLテキストエンコーダーを対象とした、段階的な蒸留駆動型のパイプラインを通じて実現されています。

1. トランスフォーマーの圧縮 (12B $\to$ 2B)

著者らは、DiTのサイズを削減するために、4段階の段階的な圧縮戦略を採用しています。

• ステージ C1: アテンションヘッドのプルーニング (12B $\to$ 5B):

  • 分析: アテンションヘッドの出力に対する特異値分解（SVD）により、すべてのヘッドが独立しているわけではなく、多くのヘッドに冗長な情報が含まれていることが明らかになりました。
  • アクション: アテンションヘッドの数（$H$）を24から16へと一様に削減しました。
  • 蒸留: 残ったヘッドが削除されたヘッドの情報を集約して保持できるように、特徴レベルの損失（feature-level loss）を導入しました。学生モデル（スチューデント）は、ティーチャーモデルの最初の16個のヘッドで初期化されました。

• ステージ C2: ヘッド次元の削減 (5B $\to$ 3B):

  • 分析: 5Bモデルの機能に対して再びSVDを適用しました。
  • アクション: ヘッド次元（$d_H$）を128から96に削減し、上位の特異成分を保持しました。これにより、3Bパラメータのモデルが得られました。

• ステージ C3: ブロック・マージングによる深度プルーニング (3B $\to$ 2.5B):

  • 分析: トランスフォーマーブロックの入力・出力間のコサイン類似度を計算したところ、シングルストリーム（SS）ブロックは高い類似度（>0.85）を示し、機能的な冗長性があることが判明しました。
  • アクション: これらのブロックを単に削除するのではなく、15個の冗長なSSブロックの連鎖を、パラメータを平均化することで単一のブロックへとマージしました。これにより、トランスフォーマーの深さが38ブロックから24ブロックへと減少しました。

• ステージ C4: 静的レイヤー正規化 (2.5B $\to$ 1.8B):

  • 分析: モデルによって予測された適応型レイヤー正規化（AdaLN）係数は、サンプル間での分散が低く、CLIP投影からの動的なコンディショニングが冗長であることを示唆していました。
  • アクション: 動的なAdaLN係数を、事前計算された静的な平均値（Static-LN）に置き換えました。これにより、約0.7Bのパラメータが削減されました。モデルは、品質を回復するために、固定された係数を用いて短期間のファインチューニングが行われました。

2. 段階的なトークン・ダウンサンプリング

詳細を損なうことなく、自己アテンションの二次関数的な複雑さに対処するために：

• メカニズム: ResNetベースのダウンサンプリングモジュール（$D$）とアップサンプリングモジュール（$U$）をトランスフォーマー内に挿入しました。
• ハイブリッド戦略: モデルはハイブリッド解像度スキームで動作します。初期の拡散タイムステップ（グローバルな構造を捉える段階）では、ダウンサンプリングされたトークンを使用し（トークン数を4分の1に削減）、後半のタイムステップ（細部を精緻化する段階）では、フル解像度のトークンを処理します。
• 学習: 学習を安定させるために、「段階的トークン・ダウンサンプリング（PTD）」戦略が使用されました。モジュールは段階的に挿入・訓練され、イテレーションが進むにつれてダウンサンプリングのポイントがネットワークの前方へと移動していきます。
• 結果: 43個のトランスフォーマーブロックのうち23個がダウンサンプリングされたトークンで動作しており、これにより高解像度の処理が必要な箇所での品質を維持しながら、レイテンシを大幅に削減しています。

3. テキストエンコーダーの蒸留 (5B $\to$ 330M)

• ターゲット: 5BのT5-XXLエンコーダーを、330MのT5-Largeへと蒸留すること。
• 手法: 2段階の蒸馏プロセスを用います：
  1. 埋め込みのマッチング: 学生モデルにMLPを追加してティーチャーの出力次元に一致させ、プロンプト埋め込み間のMSEを最小化します。
  2. ブロック単位の蒸留: 学生モデルは、特定のトランスフォーマー層におけるティーチャーのプロンプト隠れ状態に一致するように訓練されます。決定的なのは、エラーの蓄積を軽減するために、ジョイントアテンションを介してデノイザーの初期層からの視覚的信号を活用してテキストエンコーダーをガイドする手法です。
  3. 安定性: 不安定な教師あり学習を避けるため、高ノイズのタイムステップ（$t < \hat{t}$）では勾配を切り離し（detach）、後半のタイムステップと最初の3つのトランスフォーマー層に蒸留を集中させます。

主な結果

• モデルサイズ: 最終的なNanoFLUXモデルは2.4Bパラメータ（2BのDiT + 330Mのテキストエンコーダー）であり、ティーチャーである17BのFLUX.1-Schnellから7倍の削減を実現しています。
• パフォーマンス:
  • モバイルデバイス上で、NanoFLUXは512 $\times$ 512の画像を約2.5秒で生成します。
  • 定量的ベンチマーク（One-IG, DPG, GenEval, HPDv3）において、NanoFLUXはティーチャーモデルに匹敵する性能を維持しており、特定の指標（例：HPSv3が11.45から10.41へ低下）において軽微な劣化が見られるのみです。
  • 定性的結果は、本モデルが元の17Bモデルと同等の画像品質とプロンプトへの忠実度を保持していることを示しています。
• 効率性: 段階的なトークン・ダウンサンプリング戦略により、推論レイテンシと生成品質のバランスに成功しており、リソース制約のあるハードウェア上での効率的な動作を可能にしています。

意義と主張

本論文は、NanoFLUXが、FLUXのようなモデルを具体的にターゲットとした、大規模なT2I生成のための包括的な大規模モデル圧縮の研究を行った最初のフレームワークであると主張しています。その意義は以下の点にあります：

1. オンデバイス生成の実現可能性: クラウドインフラに依存することなく、高品質なT2I生成がモバイルデバイス上で実現可能であることを実証しました。
2. アーキテクチャ的洞察: 単なる汎用的な量子化やステップ削減に頼るのではなく、特定のアーキテクチャ上の冗長性（アテンションヘッド、ブロックの深さ、正規化のダイナミクス）を特定し、活用することで、大規模な生成モデルをスケールダウンするための実用的な経路を提供しました。
3. 汎用性: 将来の大規模生成モデルにも適用可能な、段階的な蒸留パイプラインと知見を提供しています。

著者らは、本研究を、計算リソースの限られたユーザーに対して高品質な生成AIへのアクセスを広げるためのステップとして位置づけています。