Tuning Language Models by Mixture-of-Depths Ensemble

技術要約：Mixture-of-Depths Ensemble (MoDE) による言語モデルのチューニング

問題提起
従来のTransformerベースの大規模言語モデル（LLM）は、ファインチューニングにおける予測および損失計算において、最終層の表現のみに依存しています。「ロジット・レンズ（logit lens）」のような解釈可能性ツールは、後半の層の隠れ状態がすでに高度に形成されたタスク関連の予測を含んでいることを示してきましたが、これらの中間表現は通常、明示的な学習信号としてではなく、暗黙的な副産物として扱われています。その結果、後半の層が持つ潜在的な予測能力が十分に活用されておらず、標準的なファインチオニングは最終層のみを最適化するため、中間表現は間接的にしか最適化されません。さらに、中間層を活用する既存の手法（効率化のための早期脱出や表現ステアリングなど）は、層固有のロジットを最終的な予測の直接的かつ学習可能なコンポーネントとして扱うことはありません。

手法：Mixture-of-Depths Ensemble (MoDE)
著者らは、LLMの最後の $k$ 層の出力を、単一モデル内の「エキスパート」のアンサンブルとして扱うチューニングフレームワークであるMixture-of-Depths Ensemble (MoDE) を提案しています。異なるトークンを異なるモデルへとルーティングするMixture-of-Experts (MoE) とは異なり、MoDEは、同一モデル内の異なる深さにあるロジットの重み付き結合へと最終予測をルーティングします。

このフレームワークは、以下の3つのコアコンポーネントで構成されています。

1. 後半層ロジット抽出 (Late-Layer Logit Extraction): $n$ 層のモデルに対し、最後の $k$ 層 ($K = \{L_{n-k+1}, \dots, L_n\}$) を選択します。各層の隠れ状態に対して語彙ヘッド $\phi(\cdot)$ を適用し、ロジット $\ell_i(x)$ を生成します。
2. 学習可能なルーティング・ネットワーク (Learned Routing Network): 単一の線形層とsoftmaxで実装された軽量なルーター $G(x)$ が、最後の $k$ 層の直前の隠れ状態 ($x = H_{n-k}$) に基づいてルーティングの重みを計算します。最終的なロジットは、以下の重み付き和として計算されます：
   $$\ell(x_t | x_{<t}) = \sum_{i=0}^{k-1} G(x)_i \cdot \ell_i(x)$$
   推論時の早期脱出（early exit）を可能にするため、ルーターはスパース（Top-K）にすることも可能です。
3. 正規化と蒸留による適応 (Adaptation via Normalization and Distillation): 事前学習済みのヘッドを中間層に直接適用すると、初期の損失が高くなることがよくあります。重いパラメータを追加することなく（各層ごとのアフィン行列とは異なり）これを解決するために、著者らは以下を導入しています。
   • 層ごとの正規化 ($N_k$): 事前学習済みアーキテクチャ（例：LayerNorm）に適合する学習可能な正規化モジュールを各後半層に追加し、表現をヘッドの入力空間に適応させます。
   • 蒸留損失 ($L_{distill}$): 最終層が「教師」として機能します。KLダイバージェンス損失を適用することで、より前の後半層の出力分布を最終層の分布に引き寄せ、学習を安定させ、アンサンブルをタスク目的関数に整合させます。

総損失は $L_{MoDE} = L_{task} + \lambda L_{distill}$ となります。MoDEは、既存のパラメータ効率の高いファインチューニング（PEFT）手法（例：LoRA, DoRA）の上にレイヤーとして重ねることができ、あるいは独立して使用できるプラグイン・メカニズムとして設計されています。

主な貢献

• 学習可能な後半層シグナル: 蒸留によって個別に最適化された後半層のロジットが、最終層と同等の損失に収束し、補完的な問題セットを解決できることを示しました。
• MoDEフレームワーク: 後半層のロジットを学習可能なルーターを通じてアンサンブルする、低オーバーヘッドのチューニングメカニメントです。これはモデルサイズの約+0.04%という極めて微量な学習可能パラメータしか追加せず、大幅に大きな学習可能モジュール（例：全層へのLoRA）に匹敵する性能を実現します。
• 効率性と柔軟性: フレームワークはスパースなTop-Kルーティングをサポートしており、最小限の精度低下で最大1.6倍の生成速度向上（早期脱出）を可能にします。また、既存のPEFT手法と組み合わせることで、その性能を向上させることができます。

実験結果
実験は、算術推論、常識推論、および指示従順（instruction following）のタスクを用いて、LLaMA2-7BおよびLLaMA3-8Bに対して行われました。

• 算術推論: MoDEはベースとなる手法（LoRA, DoRA）の性能を一貫して向上させました。軽量版（最後の $k$ 層を除く全層にLoRAを適用し、さらにMoDEを使用）は、全層LoRAよりも大幅に少ないパラメータを使用しながら、全層LoFAの性能の0.3〜0.5ポイント以内に到達しました。
• 常識推論: MoDEは、重みベースの手法（LoRA, DoRA）に対して加算的な利得を提供しました。表現ベースの手法（ReFT）は単独でも強力な性能を示しましたが、MoDEは重みベースのアプローチと組み合わせた際に信頼性の高い改善をもたらしました。
• 指示従順: MoDEはMT-Benchにおいて小さくも一貫した利得をもたらしましたが、推論タスクと比較すると改善は限定的でした。これは、指示従順が、アンサンブルによって部分的にしか捉えられない後半層のフォーマットに強く依存していることを示唆しています。
• アブレーション解析と分析:
  • ルーターは自然に最後の1〜2層に重みを集中させますが、蒸馏と正規化のコンポーネントにより、より前の後半層も効果的に寄与できるようになります。
  • アンサンブルサイズ $k \approx 3\text{--}4$ が最適です。$k$ が大きくなると、容易なデータセットには役立ちますが、複雑な推論タスクでは悪影響を及ぼします。
  • スパースなTop-Kルーティングは、精度の大部分を維持しながら、大幅な高速化を可能にします。

意義と主張
著者らは、MoDEを、後半層のロジットを直接チューニングに利用可能にできることを示すメカニズム研究として位置づけています。論文の主張は以下の通りです：

1. 後半層の表現は、単純な適応（正規化と蒸馏）を通じて解き放つことができる、重要な予測ポテンシャルを有している。
2. これらの中間シグナルは、同等の性能を持ちながら、大幅に大きな学習可能モジュール（全層LoRAなど）を代替することができる、パラメータ効率の高い代替案となる。
3. MoDEは、推論時のプロービングや重み編集のアプローチとは異なる、中間層の表現を実用的な学習シグナルへと変えるための補完的な方向性を提供する。

本論文は、固定された $k$ への依存、指示従順における改善の小ささ、およびデコーダーのみのLLMに限定された評価といった限界についても触れており、その範囲について控えめな姿勢を保っています。今後の課題として、動的な $k$ の選択、より広範なモデルファミリーへの適用、および大規模モデルへの拡張が示唆されています。