Medusa Heads

Medusaヘッドは、大規模言語モデルの最終隠れ層に追加する軽量なデコーディングヘッドです。それぞれのヘッドは、モデル自身の出力層よりも先のトークンを予測するように学習します。1回のフォワードパスで、モデルは通常の次トークン予測に加え、その先のいくつかの位置の予測も生成します。次のパスでツリー注意機構がそれらの予測を検証し、モデルが合意する最長の連続部分を採用します。Medusaは投機的デコーディングの一種で、別のドラフトモデルが不要なため、追加のネットワークを学習、ホスト、同期する必要はありません。Tianle Cai、Tri Daoらは、2024年1月にMedusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Headsでこの手法を発表しました。

ヘッドによる先読みの仕組み#

標準的なテキスト生成では、位置tの最終隠れ状態を言語モデルのヘッドに渡して、位置t + 1のトークンをスコアリングします。Medusaのヘッドkは同じ隠れ状態を読み取り、位置t + k + 1のトークンをスコアリングします。各ヘッドは、SiLU活性化と残差接続を備えた単一のフィードフォワード層に続き、語彙への射影を行います。射影層は言語モデルのヘッドのコピーとして初期化され、フィードフォワード層はゼロで初期化されます。そのため、学習前はすべてのヘッドがベースモデルの次トークン予測を再現します。論文では5つのヘッドで十分であると報告されています。このPyTorchの例では、3つのヘッドを構築し、1つの隠れ状態に対する出力をベースモデルのヘッドと比較します。

import torch
import torch.nn.functional as F

hidden_size, vocab_size, num_heads = 256, 1000, 3
lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)  # 凍結されたモデルの出力層

# 各Medusaヘッド：ゼロ初期化された単一の残差SiLU層と、LMヘッドのコピー
w1 = [torch.nn.Linear(hidden_size, hidden_size) for _ in range(num_heads)]
w2 = [torch.nn.Linear(hidden_size, vocab_size, bias=False) for _ in range(num_heads)]
for layer, out in zip(w1, w2):
    torch.nn.init.zeros_(layer.weight)
    torch.nn.init.zeros_(layer.bias)
    out.weight.data.copy_(lm_head.weight.data)

# すべてのヘッドが同じ最終隠れ状態を読み取ります。ヘッドkは位置t + k + 1を予測します
h = torch.randn(1, hidden_size)
head_logits = [out(h + F.silu(layer(h))) for layer, out in zip(w1, w2)]
print([torch.equal(logits, lm_head(h)) for logits in head_logits])

出力[True, True, True]から、ゼロ初期化された残差層により、新しい各ヘッドがベースモデルの予測を正確にコピーした状態で開始することがわかります。学習によってヘッドkは自身が担当する将来の位置に向けて調整されますが、lm_headとバックボーンは変更されません。

ツリー注意機構と受け入れ#

各ヘッドの最上位の予測を組み合わせて、候補となる続きの系列を作ります。ヘッド1から上位2トークン、ヘッド2から上位3トークンを選ぶと、直積から共通の接頭辞を持つ6つの候補が得られ、ツリーを形成します。

次のフォワードパスでは、候補トークンをすべて一度に入力に追加し、アテンションマスクによって各トークンがツリー内の自身の祖先だけを参照できるようにします。これにより、モデルはすべての分岐を並列にスコアリングします。モデルが受け入れる最長の分岐を採用し、その最後の隠れ状態を次のラウンドのヘッドに入力します。すべてのヘッドが同じ隠れ状態を読み取るため、ヘッド2はヘッド1のどの分岐でも同じトークンを提案します。著者らは、最も精度の高い組み合わせから構築した64ノードの疎なツリーのほうが、密な256ノードのツリーより高い高速化率を達成すると報告しました。

貪欲デコーディングでは、候補トークンがモデルの最上位の選択と一致した場合に受け入れます。サンプリングでは、Medusaは典型的な受け入れ方式を提供します。これは、ベースモデルでの確率が分布のエントロピーに応じたしきい値を超える候補を採用する方法です。分布の厳密な一致を犠牲にして、より多くのトークンを受け入れます。

Medusa-1とMedusa-2の学習#

Medusa-1では、新しいヘッドだけを凍結したバックボーン上で学習するため、モデル出力の品質が保たれます。論文では、生成速度が2.2倍以上になると報告されています。Vicuna-7Bのヘッドの学習には、60,000件のShareGPT会話を使い、単一のA100 GPUで約5時間かかりました。これは、freeze=23がUltralytics YOLO26のファインチューニング時に設定するのと同じ分割です。レイヤー0～22を固定し、レイヤー23の検出ヘッドだけを新しいデータで学習させます。

Medusa-2では、ヘッドとバックボーンを一緒にファインチューニングします。バックボーン本来の能力を保護する学習レシピを使い、ヘッドの精度を上げて2.3～3.6倍の高速化を実現します。元の学習データがない場合は、自己蒸留によって、モデル自身が生成したテキストでヘッドを学習します。

Medusaは、ターゲットモデル自体からトークンをドラフトする複数の手法の1つです。EAGLEは小さな自己回帰層を使い、モデルの内部特徴からドラフトを生成します。DeepSeek-V3で使われているマルチトークン予測は、事前学習中に追加の予測モジュールを学習し、サービングエンジンが後にドラフターとして再利用します。これら3つの手法はすべて、インタラクティブなチャットボットのようにバッチサイズが小さく、フォワードパス中にGPUがほとんどメモリー待ちになる場面で、推論レイテンシを最も大きく削減します。