AI・機械学習
上級

Medusa(LLMデコーディング)(メドゥーサ)

LLMの既存Transformerモデルに複数の予測ヘッドを追加し、各ヘッドが異なる将来位置のトークンを同時に予測することで、追加学習のみで1.8〜2.8倍の推論高速化を実現する並列デコーディング手法。

0 回閲覧
0 いいね

Medusaとは

Medusa(メドゥーサ)は、Carnegie Mellon UniversityのCai et al.が2024年に発表した並列デコーディング手法である。既存のLLMに軽量な予測ヘッド(Medusa heads)を追加し、各ヘッドが1〜5トークン先の出力を同時に予測する。ドラフトモデルが不要なため、Speculative Decodingと比較してGPUメモリの追加消費が5〜10%と少なく、単一GPU環境でも導入しやすい点が特徴である。

Medusaの動作原理

Medusaは以下の3段階で動作する。

  1. 並列予測: ベースモデルの最終隠れ層の出力を入力として、K個のMedusa heads(通常K=2〜5)がそれぞれ1〜Kトークン先を予測する。各ヘッドは1層のResidual MLP(約50Mパラメータ/ヘッド)で構成
  2. ツリー候補生成: 各ヘッドの上位候補(top-k、通常k=10〜50)を組み合わせてツリー構造の候補シーケンスを生成。例えばK=3、k=10の場合、最大1,000パターンの候補が生成される
  3. ツリー検証: Tree Attentionにより全候補を1回のフォワードパスで検証し、最長の有効シーケンスを受理。平均受理長は2.5〜3.5トークン/ステップ

バージョン比較

項目Medusa-1Medusa-2
学習方式ヘッドのみ学習(ベースモデル凍結)ベースモデルとヘッドを同時学習
学習コストA100 1台で1〜2時間A100 8台で1〜2日
高速化倍率1.8〜2.2x2.3〜2.8x
品質影響ベースモデルと同一微小な品質向上(+0.3% on MT-Bench)
適用範囲任意の学習済みモデル事前学習から組み込む場合に最適
メモリ増加+5%(ヘッドのみ)+5%(ヘッドのみ)

対応モデルとベンチマーク

2026年時点で公開されているMedusaヘッド付きモデルの主要な性能データは以下のとおりである。

ベースモデルヘッド数高速化倍率MT-BenchGPU
Llama 3.1 70B32.5x8.92(変化なし)A100 80GB
Llama 3.1 8B52.1x8.15(変化なし)RTX 4090 24GB
Vicuna 33B32.3x7.45(変化なし)A100 40GB
Mistral 7B42.0x7.82(変化なし)RTX 3090 24GB
Qwen 2.5 72B32.4x9.01(変化なし)H100 80GB

学習と導入手順

Medusaヘッドの追加学習は比較的軽量で、以下の手順で実施できる。

  • データセット準備: ShareGPT(約68K会話)またはOpenOrca(約1M件)から30〜50K件をサンプリング。ドメイン特化の場合はファインチューニングデータと同一データセットを使用
  • 学習設定: バッチサイズ64、学習率1e-3、エポック数1〜3。A100 80GB 1台でLlama 3.1 70Bのヘッド学習が約2時間で完了
  • 推論設定: vLLM 0.8+では --speculative-model medusa --num-speculative-tokens 3 で有効化。temperature=0の場合はExact Matchモード、temperature>0の場合はTypical Acceptanceモードを使用
  • メモリ見積もり: 70Bモデル(FP16)で約140GB + Medusaヘッド3個で約7GB = 合計約147GB。A100 80GB x2で運用可能

Speculative Decodingとの使い分け

判断基準Medusa推奨Speculative Decoding推奨
GPUメモリ制約あり(単一GPU)余裕あり(マルチGPU)
品質要件ロスレス必須ロスレス必須
導入コスト低(ヘッド学習のみ)低(ドラフトモデル選定のみ)
最大高速化2.8x程度3.5x程度
バッチサイズ小〜中(1〜16)小(1〜8)

よくある質問(FAQ)

Q1: Medusaヘッドの学習にはどの程度のデータが必要ですか? A: 30〜50K件の会話データで十分な品質が得られる。ShareGPTデータセットが最も広く使われており、A100 1台で1〜2時間で学習完了する。ドメイン特化モデルの場合は、そのドメインのデータを使うことで受理率が5〜10%向上する。

Q2: 既存のファインチューニング済みモデルにも適用できますか? A: Medusa-1であれば任意の学習済みモデルにヘッドを追加学習するだけで適用可能。ベースモデルの重みは一切変更しないため、ファインチューニングの効果はそのまま維持される。

Q3: Medusaヘッドの数はいくつが最適ですか? A: 一般的にはK=3(3ヘッド)が最もコストパフォーマンスが良い。K=5にすると理論上の最大受理長は伸びるが、各ヘッドの予測精度が低下するため、実効高速化倍率はK=3と比べて5〜10%程度の改善にとどまる。

まとめ

  • Medusaは追加ヘッドのみで1.8〜2.8倍の推論高速化を実現するコスト効率の高い手法
  • ドラフトモデル不要でGPUメモリ増加は5〜10%に抑えられる
  • Medusa-1は任意の既存モデルに後付け可能、Medusa-2はさらに高い高速化を達成
  • vLLM 0.8+でネイティブサポートされており、導入障壁が低い
  • メモリ制約の厳しい環境ではSpeculative Decodingより有利な選択肢