Rethinking Cross-Layer Information Routing in Diffusion Transformers

技術サマリー：拡散トランスフォーマーにおけるクロスレイヤー情報ルーティングの再考

問題定義
拡散トランスフォーマー（DiT）は、トークナイズ化、アテンションメカニズム、条件付け、目的関数に関する広範な研究を経て、現代の視覚生成の基盤として確立されています。しかし、層を超えて情報が蓄積される仕組みを支配する「残差ストリーム」は、その元のトランスフォーマー設計からほとんど変更されていません。著者らは、NLP の慣習から受け継がれたこの標準的な事前正規化された残差加算が、拡散プロセスの時間変動的なダイナミクスには不適切であると特定しました。

著者らは、深さとノイズ除去タイムステップに沿ったクロスレイヤー情報流の体系的な実証分析を通じて、DiT における従来の残差ルーティングの 3 つの具体的な症状を診断しました。

1. 単調な前方マグニチュードの膨張: 隠れ状態のマグニチュードは、最初のブロックから最後のブロックまで約 100 倍に成長します。これにより、より深いブロックは影響力を維持するために、ますます大きな生出力を生成することを強要され、大規模言語モデル（LLM）で見られる「PreNorm の希釈」現象を反映しています。
2. 急激な後方勾配の減衰: 勾配は最初の数ブロックの後に著しく減衰し、より深い層には実質的に弱い最適化シグナルしか残されません。
3. 顕著なブロックごとの冗長性: 隣接するトランスフォーマーブロックは非常に類似した表現（コサイン類似度 > 0.9）を生成し、標準的なルーティングが深さ全体で特徴を多様化させることに失敗していることを示しています。

さらに、分析により、特定の層にとって最適な情報源はノイズ除去タイムステップに応じて体系的に変化することが明らかになりました（高ノイズでは粗い構造的特徴から、低ノイズでは細かい詳細特徴へ）。均一で時間非依存的な重みを適用する標準的な残差接続は、これらの変化する要件に適応できません。

手法：拡散適応ルーティング（DAR）
これらの問題に対処するため、著者らは標準的な残差接続のドロップイン置換となる拡散適応ルーティング（DAR）を提案しました。固定された加法的再帰とは異なり、DAR はサブレイヤー出力の履歴に対して学習可能で、タイムステップ適応的、かつ非増分的な集約を実行します。

主要な設計要素には以下が含まれます。

• 深さに対する Softmax アテンション: 単位重みで以前の出力を合計する代わりに、DAR は先行するサブレイヤー出力（$v_i$）の softmax 重み付き集約を計算します。集約重みは、現在の隠れ状態から導出されたクエリ（$q_l$）と、正規化された以前の出力から導出されたキー（$k_i$）によって決定されます。
• タイムステップ認識: ノイズ除去タイムステップ $t$ が重要な制御シグナルであることを認識し、クエリパラメータ化は $t$ を認識するように設計されています。著者らは 3 つの変種を検討しました。
  • 静的: 層ごとの学習可能ベクトル（時間非依存）。
  • 動的: クエリは前のサブレイヤー出力の投影であり、隠れ状態に符号化されたタイムステップ情報を暗黙的に継承します。
  • 明示的注入: 静的クエリにタイムステップ埋め込み $e(t)$ を追加します。
    実験により、タイムステップ認識が重要であることが示され、動的および明示的変種は時間盲の静的ベースラインを上回りました。
• チャンク化された集約: 全ての $L$ 個の以前の層にアテンションする計算コストを管理するため、DAR はチャンク化された変種をサポートします。サブレイヤーはチャンクに分割され、ルーターは以前のチャンクの要約と、現在のチャンク内の生出力にアテンションします。理論的分析は、最適なチャンクサイズが存在することを示唆しており（SiT-XL/2 については経験的に $S=4$ と判明）、ルーティングの精度と圧縮コストのバランスを取ります。
• 互換性: DAR は純粋に深さ次元に沿って動作し、等方的で均一なトランスフォーマースタックを維持します。手動での層ペアリングや学習目的の変更を必要とせず、REPA（表現アライメント）などの他の現代の拡張機能と互換性があるように設計されています。

主要な結果
SiT-XL/2 アーキテクチャを用いた ImageNet 256 × 256 での評価結果は以下の通りです。

• 品質の向上: DAR の静的変種は（クラスターフリーガイダンスなしで）7.56 の FIDを達成し、ベースラインの SiT-XL/2（9.67 FID）に対して2.11 ポイントの改善を示しました。
• 収束速度: DAR は、ベースラインの収束品質を約8.75 倍少ないトレーニング反復で達成します。
• REPA との相乗効果: REPA の上にスタックすると、DAR は初期段階で2 倍のトレーニング加速をもたらします。これは、クロスレイヤールーティングの改善が表現アライメント目的と直交していることを示しています。
• U-Net スキップとの比較: DAR は、手動設計された U-Net 型のスキップ接続を使用するモデル（U-DiT-L など）を上回ると同時に、均一なトランスフォーマースタックのスケーラビリティの恩恵を維持します。
• 蒸留: 大規模な T2I 事後トレーニング（Qwen-Image における分布マッチング蒸留）において、DAR は積極的な少数ステップ蒸留で失われがちな高周波の詳細（シャープなエッジ、微細なテクスチャ）を保持するのに役立ちます。

意義と主張
本論文は、クロスレイヤー情報ルーティングが、既存の表現アライメント目的やアーキテクチャのスケーリングと直交して動作する、拡散モデリングにおいて未探索の設計軸であると主張しています。DiT における標準的な残差ストリームの特定の時空間病理を診断することで、著者らは、固定された加算を適応的かつタイムステップ条件付きの集約に置き換えることが、収束速度と最終的な生成品質の両方を大幅に向上させることを実証しました。

この研究は、クロスレイヤールーティングを継承された慣習から明示的な設計次元へと昇華させます。標準的な残差の硬直性が拡散モデルの制限要因であることを示唆し、特により深いスタックやより複雑な事後トレーニング体制へとスケールする際に、DAR のような適応的ルーティングメカニズムが深層 DiT アーキテクチャの潜在能力を解き放つために不可欠であることを示しています。