アーリーイグジット（LLM）（アーリーイグジット）

アーリーイグジット（Early Exit）in LLM

アーリーイグジット（Early Exit）は、Transformerの中間層に補助出力ヘッドを設け、モデルが十分な確信度を持った時点で残りの層の計算をスキップして早期に出力する手法。

基本構造

標準Transformerは L 層を必ず全て通過する:

Input → Layer1 → Layer2 → ... → LayerL → Output

Early Exitモデルは各層またはN層ごとに出力ヘッドを追加:

Input → Layer1 → [Exit1?] → Layer2 → [Exit2?] → ... → LayerL → Output
           ↓ 確信度高                    ↓ 確信度高
         早期出力                       早期出力

確信度の測定

エントロピーベース: softmax出力のエントロピーが低い（特定トークンに集中）→ 確信度高

import torch

def should_exit(logits: torch.Tensor, threshold: float = 0.5) -> bool:
    probs = torch.softmax(logits, dim=-1)
    entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=-1)
    return entropy.item() < threshold

最大確率: top-1トークンの確率が閾値以上

return probs.max().item() > 0.9

代表的研究

CALM (Confident Adaptive Language Modeling)

• Meta AI 2022年発表
• 各位置で独立にEarly Exit判断
• 平均計算量を 37-50% 削減

DejaVu

• 入力の 70% でスパースアテンション・FFN利用
• LLaMA-7B で 2倍高速化

SkipDecode

• バッチ内の異なるトークンが異なる層数で処理可能
• Continuous Batchingと組み合わせ

トレーニング方法

補助ヘッドを含む総合損失:

L = L_final + λ * Σ L_exit_i

各Exit層の損失も最終層同様にCross Entropyで訓練。λでEarly Exit重みを制御。

精度とのトレードオフ

課題

• バッチ処理との不整合: バッチ内で異なるリクエストが異なる層でExitするとGPU効率が落ちる
• 学習コスト: 補助ヘッドの追加で事前学習/ファインチューニングが必要
• 生成タスクへの適用: 各トークン生成に適用するため、複雑な実装が必要

生成タスクでは各トークンごとにEarly Exit判断が必要で、Continuous Batchingとの統合が研究課題となっている。