Medusa（マルチヘッドデコーディング）（メデューサマルチヘッドデコーディング）

Medusa — 複数ヘッドによる並列トークン投機生成

Medusaは2023年にCai et al.（Tsinghua University/Princeton）が発表したLLM推論高速化手法だ。従来の投機的デコーディングが別個のドラフトモデルを必要としていた問題を解決するため、メインのLLMの最終隠れ状態に複数の「Medusaヘッド」を追加し、ドラフト生成をモデル内部で完結させる。

従来の投機的デコーディングの問題

基本的な投機的デコーディングの課題:

1. ドラフトモデルの管理: 別のモデルのウェイトをVRAMに追加ロードする必要
2. モデル選択の困難: ドラフトモデルはターゲットと同じ語彙・アーキテクチャ系統が必要
3. 品質のトレードオフ: ドラフトモデルが小さすぎると受理率が低下し効果が薄れる

Medusaはこれらを「ヘッドを追加するだけ」で解決する。

Medusaヘッドのアーキテクチャ

MedusaヘッドはLLMの最終Transformer層の出力（最終隠れ状態）を入力として受け取る小型の予測ヘッドだ:

# Medusa Head: 2層の前向き層 + 出力層
class MedusaHead(nn.Module):
    def __init__(self, hidden_dim, vocab_size):
        super().__init__()
        self.fc1 = nn.Linear(hidden_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, vocab_size)
    
    def forward(self, x):
        return self.fc2(F.silu(self.fc1(x)))

複数のMedusaヘッド（通常5個）を追加し、ヘッドiはi+1番目先のトークンを予測する:

• Head 1: position t+1のトークンを予測
• Head 2: position t+2のトークンを予測
• Head k: position t+kのトークンを予測

Medusa-1 と Medusa-2

Medusa-1（ベースモデル凍結）:

• ベースLLMのウェイトは完全に凍結
• Medusaヘッドのみを少量のデータ（ShareGPT等）でファインチューニング
• 元のモデル品質を保ちながら高速化を追加できる
• LoRAと同様の「追加モジュールのみ訓練」パラダイム

Medusa-2（共同最適化）:

• ベースLLMとMedusaヘッドを同時にファインチューニング
• より高い受理率（Medusa-1比10〜20%向上）を実現するが、元のモデル性能に影響する可能性
• FP16 + LoRAで効率的に訓練可能

ツリーデコーディングとの統合

Medusaは複数ヘッドの出力を使いツリー構造の候補を形成してから一括検証:

1. 各Medusaヘッドがtop-K候補トークンを出力
2. カルテシアン積でトークンのツリーを形成（最大 K^5 個のパス）
3. ターゲットLLMがツリーAttentionで全パスを一括検証
4. 最長の受理されたパスを採用

パフォーマンス

Vicuna 13B (A100 80GB):

• Medusa-1: スループット2.2〜2.5倍向上
• Medusa-2: スループット2.5〜3.0倍向上
• 出力品質: 受理率85〜90%（設定依存）