多トークン予測（Medusa・EAGLE）

なぜ「1ステップで複数トークン」を狙うのか

KV キャッシュと LLM 推論の最適化 で見たとおり、LLM の Decode はメモリ帯域律速です。1トークン出すだけでも全パラメータと KV を HBM から読み出す必要があり、演算ユニットは大きく余っています。ここで効くのが、投機的デコーディング と同じ発想、すなわち1回の前方計算で複数トークンを確定させて読み出しコストを割り勘することです。

標準の投機的デコーディングは、そのために別個の小さなドラフトモデルを用意しました。しかし別モデルは、語彙を揃える・メモリに2つ載せる・整合性を保つといった運用コストを伴います。これを避け、ターゲットモデル自身に先読みさせるのが本記事の主題、**多トークン予測（multi-token prediction）**です。代表格が Medusa と EAGLE で、いずれも「本体にわずかな部品を足すだけ」で複数トークンの候補を生み、それを一括検証します。

Medusa：本体に複数の予測ヘッドを足す

Transformer は通常、最終層の隠れ状態 h_t を1つの LM ヘッド（線形層＋softmax）に通して「次の1トークン」を出します。Medusa は、この h_t の上に追加の軽量ヘッドを複数（Medusa ヘッド）並べます。

• ヘッド0（元の LM ヘッド）: 次トークン t+1 を予測（従来どおり）。
• Medusa ヘッド1: 2手先 t+2 を予測。
• Medusa ヘッド2: 3手先 t+3 を予測。
• …と t+k まで拡張。

各 Medusa ヘッドは「同じ h_t から、複数手先を直接当てる」小さな全結合ブロック（残差付き）で、本体パラメータは凍結したままヘッドだけを追加学習します。これが Medusa の第一の利点です。別モデルも、本体の再学習も基本要りません。

              ┌─→ LM ヘッド(既存) ─→ t+1 の候補分布
  h_t ────────┼─→ Medusa ヘッド1 ─→ t+2 の候補分布
 (最終隠れ状態) ├─→ Medusa ヘッド2 ─→ t+3 の候補分布
              └─→ Medusa ヘッド3 ─→ t+4 の候補分布
        すべて同一の前方計算1回で並列に得られる

ここで一つ、原理的な弱点があります。各ヘッドは同じ h_t だけを見て独立に先を当てるため、t+2 を予測する際に「t+1 が実際に何になったか」を条件にできません。先の位置ほど不確実性が積み上がり、単独の当たり精度は落ちます。Medusa はこれを、各ヘッドから上位数候補を取り出して組み合わせのツリーを作り（後述）、多めの経路を一括検証することで補います。当たりの1本を引ければよい、という設計です。

Medusa には代表的に2段階の学習があります。Medusa-1 は本体を凍結しヘッドだけ学習する軽量版、Medusa-2 は本体も含めて共同微調整し受理長を伸ばす版です。手軽さの Medusa-1、性能の Medusa-2 という住み分けです。

EAGLE：token ではなく「特徴量」を先読みする

EAGLE（Extrapolation Algorithm for Greater Language-model Efficiency）は、Medusa とは別角度から多トークン予測を実現します。核心の観察はこうです。token の系列は離散でサンプリング由来のノイズを含み予測しにくいが、出力直前の「特徴量（feature、＝最終層の隠れ状態）」の系列は連続的で滑らかであり、次を外挿しやすい。そこで EAGLE は、token 列ではなく特徴量列のほうを自己回帰的に先読みします。

仕組みは次のとおりです。EAGLE は小さな自己回帰ヘッド（1層 Transformer 程度）を1つ持ち、これに「これまでの特徴量」と「直前に確定した token の埋め込み」を入力して次の特徴量を予測します。得られた特徴量をターゲット本体と同じ LM ヘッドに通せば、token の候補分布が出ます。これを数手ぶん繰り返してドラフト列（実際にはツリー）を作り、最後に本体で一括検証します。

EAGLE の先読み（1ステップ）:
  入力 = [ 過去の特徴量 f_(≤t) , 直前に確定した token の埋め込み e_t ]
  ↓ 小さな自己回帰ヘッド(1層 Transformer 相当)
  次の特徴量 f_(t+1) を予測
  ↓ ターゲットと同じ LM ヘッド(重み共有)
  token t+1 の候補分布
  → サンプルした token を1つ遅らせて次ステップの入力へ混ぜる

EAGLE が Medusa と決定的に違う点が2つあります。第一に、Medusa の各ヘッドが h_t だけを見て独立に先を当てるのに対し、EAGLE は先読みヘッドを自己回帰で回すので「1つ先で何を出したか」を次の予測に反映できる。積み上がる不確実性を系列内で扱えるわけです。第二に、EAGLE は次に確定した token を1ステップ遅らせて（shift して）特徴量入力に足し込む設計をとります。特徴量だけを外挿すると「実際にどの token が選ばれたか（サンプリングの結果）」が抜け落ちますが、遅延 token を混ぜることでこのサンプリング由来の不確実性を吸収し、先読み精度を上げています。

Medusa と EAGLE を並べて捉える

同じ「別モデルを持たない多トークン予測」でも、候補の作り方が根本的に異なります。

大づかみには、Medusa は「安く広く候補をばらまいてツリーで拾う」、**EAGLE は「滑らかな特徴量を丁寧に外挿して当てにいく」**という対比で覚えると要点を外しません。一般に同規模の追加コストなら EAGLE 系のほうが受理長（＝実効速度）を稼ぎやすいと報告される一方、Medusa は構造が単純で実装・学習が軽い、という実務的トレードオフがあります。

ツリー検証：候補を「木」に束ねて1回で確かめる

多トークン予測に共通する武器がツリー検証（tree attention）です。素朴な投機は「候補を1本の列」として検証しますが、先頭で1つ外すとそれ以降が丸ごと無駄になります。そこで各位置で複数候補を持たせ、それらの組み合わせを木（複数経路）として1回の前方計算にまとめて検証します。

鍵は専用の注意マスクです。木の各ノード（候補 token）が「自分の祖先ノードだけ」を参照できるようにマスクを組むと、独立した複数の候補列を1つのバッチ化された系列として同時に流せます。1回の前方計算で全経路の p が並列に手に入り、その中から受理判定を最も長く通る1本を採用します。

1本の列(直列):    A → B → C → D     ← 先頭Aを外すと全滅

ツリー(複線化):        ┌ B1 ┐
                   A ─┤    ├ …    複数の (B,C,…) 経路を
                       └ B2 ┘      1回の前方計算でまとめて検証
   専用の注意マスクで「各ノードは祖先のみ参照」を強制する

木を広く深くするほど1ラウンドで探索できる経路が増え、受理長の期待値が伸びます。ただし木を大きくすれば検証1回で処理する token 数（＝計算量）も増えるため、演算がまだ余っている範囲で木を広げるのが得です。ここが「無限に広げれば速い」とはならない理由で、EAGLE-2 の動的ツリーは、この木の形を文脈の信頼度に応じて最適化して受理長と計算量のバランスを取る発想にあたります。

まとめ：本体を活かして Decode を割り勘にする

多トークン予測は、Decode がメモリ帯域律速で演算が余る構造を、別モデルなしで突く手法です。Medusa は本体にヘッドを足して複数手先を並列に候補化し、EAGLE は滑らかな特徴量を自己回帰で外挿して本体の LM ヘッドで token 化します。いずれも候補をツリーに束ねて1回の前方計算で一括検証することで、実効速度を稼ぎます。

実務では、受理長を実測してヘッド／先読みヘッドの学習と木の形を詰めるのが出発点です。検証の数理は 投機的デコーディング に、Decode がなぜ帯域律速かは KV キャッシュと LLM 推論の最適化 に、受理長を左右するサンプリング設定は デコーディング戦略 に接続しており、あわせて読むと推論高速化の打ち手が線でつながります。