技術概要：LLM 蒸留のための多粒度軌道整合（MTA）

問題定義

既存の大規模言語モデル（LLM）向けの知識蒸留（KD）手法は、主に固定された静的な層における出力確率分布（ログit）または中間隠れ状態の整合に焦点を当てています。局所的な信号の転送には効果的ですが、これらのアプローチはネットワークの深さ全体にわたり均一なトークンレベルの整合戦略を採用することが多いです。

本論文は、この均一な扱いが、人間の言語の階層性およびトランスフォーマーモデルの機能的組織と整合していないと主張しています。言語学研究およびトランスフォーマーの解釈可能性研究は、以下のことを示しています：

1. 下位層は、表面レベル、語彙的、および事実検索の特徴を強調します。
2. 上位層は、抽象的な意味情報と構成的推論を段階的に符号化します。

すべてのトークンを孤立して扱い、同じ整合目的関数を均一に適用する現在の手法は、語彙単位から抽象的な意味構造へと変化する表現の進化的軌道を捉えることに失敗しています。この限界は、学生モデルが下流タスクに一般化する能力を制限します。

手法：多粒度軌道整合（MTA）

これらの限界に対処するため、著者は教師モデルと学生モデルの表現軌道を層適応戦略を用いて整合させるモジュール型フレームワークであるMTAを提案します。MTA は、蒸留がトランスフォーマー層の深さに依存する意味的役割を尊重すべきであるという仮説に基づいて動作します。

中核コンポーネント

1. 層適応型多粒度スパン：

   • 下位層：語彙的基盤と微細な依存関係を維持するために、単語レベル（Word Spans）で整合を実行します。
   • 上位層：構成的意味と高次構造関係を捉えるために、整合を句レベル（名詞句および動詞句）へ移行させます。
   • スパン構築：意味単位は構文パーサー（例：spaCy）を使用して抽出されます。トークンの重要度重みは、自己ループを除外した双方向自己注意機構を介して計算され、トークン表現をスパン表現に集約します。

2. 動的構造整合損失（$L_{DSA}$）：

   • この目的関数は、教師と学生の間の幾何学的整合性を強制します。
   • 個々のトークンベクトルを一致させるのではなく、$L_{DSA}$ は特定の層内の意味スパン（単語または句）間のペアワイズ関係距離の不一致を最小化します。
   • 構成トークンの集約された重要度に基づいて整合を重み付けることで、意味的に顕著なスパンを優先します。
   • 数学的には、表現空間のトポロジカル構造を整合させます：
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     ここで、$d(\cdot)$ はコサイン距離であり、$U$ は重み付けされたスパン表現を表します。

3. 隠れ表現整合損失（$L_{Hid}$）：

   • 構造制約を補完するため、特徴レベルの損失が特定の隠れ状態を整合させます。
   • 学生モデルの次元は通常、教師モデルの次元よりも小さいため、学習可能な線形プロジェクターが学生状態を教師の空間にマッピングします。
   • この損失は、最も情報量の多いトークンに焦点を当て、投影された学生状態と教師状態の間の重み付けコサイン距離を最小化します。

4. 最適化：

   • MTA はプラグインモジュールとして設計されています。総損失は、既存フレームワーク（DistiLLM、FDD など）のベース蒸留損失（$L_{Base}$）に、提案された構造損失と隠れ損失を組み合わせます：
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

主要な貢献

1. 限界の特定：著者は、均一な中間整合が表現の階層的進化を無視しており、結果として最適ではない知識転送につながると特定しました。
2. 層適応型フレームワーク：著者は MTA を導入しました。これは、下位層での単語レベル整合から上位層での句レベル構造整合へ動的に移行し、言語のボトムアップ構成的性質を反映します。
3. 汎用モジュール：MTA は、トークン化を共有する最先端の蒸留手法（DistiLLM、DistiLLM-2、FDD）と互換性のある汎用最適化レイヤーとして実証されました。
4. 実証的検証：広範な実験により、MTA が多様なモデルアーキテクチャ（GPT-2、Qwen1.5、OPT）およびベンチマークにおいて一貫して性能を向上させることが示されました。

実験結果

著者は、3 つの学生 - 教師ペアを使用して指示追従タスクにおける MTA を評価しました：

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

性能指標：

• ROUGE-L：MTA は、すべてのデータセット（Dolly、SelfInst、Vicuna、S-NI）において、ベースライン（FDD、DistiLLM、DistiLLM-2）を一貫して上回りました。例えば、GPT-2 1.5B $\to$ 120M の設定において、DistiLLM に MTA を追加することで、平均 ROUGE-L は 20.21 から 21.45 に向上しました。
• LLM-as-a-Judge：GPT-4o-mini を用いて事実の正確性、一貫性、有用性を評価したところ、MTA 強化モデルはベースラインに対して一貫してスコア向上を示しました。

アブレーション研究：

• 損失コンポーネント：$L_{DSA}$（構造）と $L_{Hid}$（特徴）の両方が正の寄与をし、それらの組み合わせが最良の結果をもたらしました。
• 粒度：「フルレベル」（層適応）戦略は、静的な戦略（単語のみまたは句のみ）を上回りました。これは、下位層での語彙的基盤と上位層での構成的意味の整合が重要であることを確認しました。
• 層選択：この手法は異なる層スケジュールに対して頑健ですが、テストされた構成ではハイブリッド割り当て（1 単語レベル層：2 句レベル層）が最適であることが判明しました。

効率性：

• MTA は、主に構文スパン抽出（最適化により約 0.27 秒/ステップに削減）および幾何学的マッチングに起因するトレーニングオーバーヘッドを導入します。
• スパン抽出はテスト時に使用されないため、推論コストは変更されません。
• 時間一致比較により、性能向上はトレーニング予算の増加ではなく、整合目的関数によるものであることが確認されました。

意義と主張

本論文は、LLM の階層的表現軌道を明示的にモデル化することで、知識蒸留に対する意味的に忠実なアプローチを MTA が提供すると主張しています。静的なトークンレベルの一致を超えて、動的な多粒度構造整合へ移行することで、MTA は学生モデルが教師の内部情報処理ダイナミクスをよりよく再現することを可能にします。

著者は、MTA を既存の蒸留フレームワークの代替ではなく、圧縮された LLM の一般化性と堅牢性を強化する補完的なプラグアンドプレイモジュールとして位置付けています。この研究は、将来のモデル圧縮技術の設計において、言語的階層性と深さに依存する特徴の進化を尊重することの重要性を強調しています。