Turning the TIDE: Cross-Architecture Distillation for Diffusion Large Language Models

以下は、論文「Turning the TIDE: Cross-Architecture Distillation for Diffusion Large Language Models」の詳細な技術的サマリーです。

1. 問題定義

拡散大規模言語モデル（dLLMs）は、並列デコードや双方向コンテキストなど、自己回帰（AR）モデルに比べて顕著な利点を提供します。しかし、最先端の dLLM は競争力のある性能を達成するために 80 億〜1000 億パラメータ以上の膨大なパラメータ数を必要とし、これが展開の障壁となっています。

知識蒸留はモデルを圧縮するための標準的な解決策ですが、dLLM に対する既存の方法（CDLM、DDD など）は、単一のアーキテクチャ内（同じ教師と生徒のアーキテクチャ/トークナイザ）でのステップ圧縮にのみ焦点を当てています。これらは、教師と生徒が以下の点で異なるアーキテクチャ間蒸留には対応していません。

1. アーキテクチャとアテンション: （例：MoE 対 Dense、双方向対因果）。
2. トークナイザ: 異なる語彙は、標準的なトークンレベルの尤度マッチングを数学的に適用不可能にします。
3. 時間的ダイナミクス: 拡散モデルでは、教師の信頼性はノイズレベル（タイムステップ）に基づいて劇的に変動しますが、AR モデルでは教師は常に完全な左コンテキストを参照します。

2. 手法：TIDE フレームワーク

著者らは、アーキテクチャ間 dLLM 蒸留のための最初の統合フレームワークであるTIDE（Turning the TIDE）を提案します。これは時間的、空間的、語彙的な障壁を克服するために設計された 3 つのモジュールコンポーネントで構成されています。

A. TIDAL（Time-Iteration Dual-Axis Lambda Modulation）

• 課題: 教師の信頼性は拡散タイムステップ（$t$）によって変動します。高ノイズ（$t \approx 1$）では教師は推測を行い、低ノイズ（$t \approx 0$）では教師は信頼できます。さらに、生徒の容量は訓練中に進化します。
• 解決策: TIDAL は蒸留強度（$\lambda_t$）を 2 つの軸に沿って調節します。
  1. 拡散タイムステップ: $\lambda_t = \lambda_{train} \times (1 - t)$。ノイズが増加するにつれて、ノイズの多い信号の学習を避けるために教師への依存度が低下します。
  2. 訓練の進行: $\lambda_{train}$ はコサインスケジュールに従い、初期は低く（生徒優位）、訓練が安定するにつれて高く（教師優位）なります。
• メカニズム: 目標分布は生徒と教師のロジットの補間であり、混合された目標を通じた勾配の流れを防ぐために計算グラフから切り離されています。

B. COMPDEMO（Complementary Demonstration-Conditioned Denoising）

• 課題: 重度のマスク（高ノイズ）下では、教師が十分なコンテキストを受け取れず、予測が劣化します。
• 解決策: マスクされた位置（$M$）を 2 つの相補的な部分集合（$M_A, M_B$）に分割します。
  • パス 1: 教師は $M_A$ が復元（クリーン）され、$M_B$ がマスクされた状態を見ます。
  • パス 2: 教師は $M_B$ が復元（クリーン）され、$M_A$ がマスクされた状態を見ます。
• メカニズム: 両方のパスからの教師のロジットをマージします。これにより教師のコンテキストが豊かになり、実質的なタイムステップが低下し、高ノイズレベルであっても生徒に対して高品質なデモンストレーション信号が提供されます。

C. Reverse CALM（Cross-Token Alignment via Reverse Likelihood）

• 課題: 教師と生徒が異なるトークナイザ（$V_T \neq V_S$）を使用する場合、トークンレベルの KL 発散は定義されません。標準的なチャンクレベルのアライメント（CALM など）は、通常「Forward」バイナリクロスエントロピー（BCE）損失を使用します。
• Forward CALM の問題: 勾配には $p_t^c / p_s^c$ という項が含まれます。生徒が教師が自信を持っているチャンクに低い確率を割り当てた場合（$p_s \to 0, p_t \to 1$）、勾配が発散します。
• 解決策: Reverse CALM は BCE 損失の方向を逆転させます：$L = -[p_s \log p_t + (1-p_s) \log(1-p_t)]$。
• メカニズム:
  • 有界勾配: 勾配係数は $\log(p_t / (1-p_t))$ となり、これは固定された教師のみに依存し、有界になります。
  • 双端ノイズフィルタリング: 整合性の取れていないチャンク（$p_t \approx 0.5$）は係数をゼロにし、低い生徒の確率はノイズを抑制します。これはスカラー空間における「モード探索」目的として機能します。

3. 実験設定

• 生徒モデル: 0.6B パラメータのブロック拡散モデル（BD3LM）。
• 教師モデル:
  1. LLaDA2.0-mini (16B MoE): 異種アーキテクチャ、異なるトークナイザ（クロストークナイザパイプライン）。
  2. WeDLM-8B-Instruct (8B Dense): 共有トークナイザファミリー（共有トークナイザパイプライン）。
• ベンチマーク: 推論（GSM8K, MATH, BBH）、知識（MMLU, MMLU-Pro）、常識（HellaSwag）、コード（HumanEval, MBPP）を含む 8 つのタスク。

4. 主要な結果

• 全体的な性能: 蒸留された 0.6B モデルは、すべてのベンチマークで蒸留されていない 0.6B ベースラインを平均 +1.53 ポイント上回りました。
• コード生成（HumanEval）: 蒸留されたモデルは性能において劇的な飛躍を達成しました。
  • 48.78（蒸留）対 32.30（同サイズの AR ベースライン）。
  • これは +16.48 ポイントの獲得を意味し、蒸留された dLLM がコードのような構造化された生成タスクにおいて卓越していることを示しています。
• パイプラインの特异性:
  • クロストークナイザ: Reverse CALM 戦略が優れており、不一致な語彙に対する有界勾配の必要性を確認しました。
  • 共有トークナイザ: TIDAL + COMPDEMO 戦略が優れており、正確なトークンアライメントを活用した漸進的カリキュラム学習を実現しました。
• 効率性:
  • メモリ: 16B 教師と比較して 22 倍削減（1.4 GB 対 31.3 GB）。
  • 速度: 16B 教師よりも 5 倍高速な推論。
  • オーバーヘッド: 蒸留は蒸留されていない生徒と比較して無視できるオーバーヘッド（スループット低下約 2.6%）しか追加しません。

5. 意義と貢献

1. 初のアーキテクチャ間フレームワーク: TIDE は、異なるアーキテクチャとトークナイザを持つ dLLM 間の知識転移を可能にする先駆的な作業であり、以前は未開拓であったギャップを埋めます。
2. 新規コンポーネント:
   • TIDAL は拡散蒸留の時間的不安定性を解決します。
   • COMPDEMO はマスク拡散におけるコンテキスト不足の問題を解決します。
   • Reverse CALM は勾配発散を回避する、数学的に安定したクロストークナイザ蒸留の解決策を提供します。
3. 実用的な展開: このフレームワークは、高性能な dLLM が、拡散の双方向性と並列性の利点を維持しつつ、汎用ハードウェアで実行可能な小規模な展開モデル（0.6B）に圧縮可能であることを証明しています。
4. コード生成のブレイクスルー: 結果は、強力な教師によって導かれた拡散デコードの並列性が、プログラミングのようなグローバルな整合性を必要とするタスクに固有に適していることを示唆しています。

結論

TIDE は、大規模で高価な拡散モデルと、小規模で展開可能なモデルの間のギャップを成功裡に橋渡しします。時間的信頼性、コンテキスト不足、語彙の不一致という固有の課題に対処することで、効率的な dLLM 展開の新たな基準を確立し、特にコード生成における拡散モデルの可能性を浮き彫りにしています。