Medusa and EAGLE: self-speculative decoding heads, explained

Medusa and EAGLE are self-speculative decoding methods: instead of a separate draft model, they attach a small trained module to the target model itself, which reads the target's hidden states and proposes several future tokens. The target then verifies the proposals in one forward pass, exactly as in speculative decoding. Medusa adds several independent prediction heads, one per future position, and verifies a tree of their top candidates. EAGLE adds one lightweight transformer layer that runs autoregressively over the target's features and its own previous guesses. Both remove the hardest practical problem of classic speculative decoding (finding a small model with the same tokenizer that agrees with your target) and replace it with a cheap training job. That job has to be redone for every fine-tune you deploy, which is why it belongs in a post-training workflow and not only in a serving configuration.

Where they came from#

Medusa was introduced by Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen and Tri Dao (2024). It descends from blockwise parallel decoding (Stern, Shazeer and Uszkoreit, 2018), which also predicted several positions with extra heads. Medusa's contributions were a simple residual head design, tree attention to verify many candidate continuations at once, two training regimes (Medusa-1 with a frozen backbone and Medusa-2 trained jointly), a self-distillation procedure for models whose training data is unavailable, and a "typical acceptance" rule. Medusa-1 reached over 2.2 times speedup on Vicuna-7B without changing outputs, and Medusa-2 reached 2.3 to 3.6 times.

EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency) was introduced by Yuhui Li, Fangyun Wei, Chao Zhang and Hongyang Zhang (2024). Its two observations were that autoregression is easier at the feature level (the second-to-top hidden state) than at the token level, and that feature prediction is uncertain unless the draft also sees the token that was actually sampled, so EAGLE feeds both. On MT-bench at temperature 0 it reached average acceptance lengths of 3.6 to 4.0 tokens and 2.8 to 3.0 times speedup on LLaMA2-Chat 7B to 70B. EAGLE-2 (2024) made the draft tree dynamic, expanding the branches the draft is confident in, for 3.05 to 4.26 times. EAGLE-3 (2025) dropped feature regression in favour of direct token prediction, fused low, middle and high-level target features, and trained with "training-time test" (feeding the draft its own predictions during training), reaching up to 6.5 times speedup, 4.40 times with an acceptance length of 6.13 on LLaMA-3.1-8B, and 1.38 times throughput at batch size 64 in SGLang.

The bottleneck, and what self-speculation changes#

At small batch sizes a decode step reads all of the target's weights to produce one token per sequence, so scoring several positions in one pass costs almost the same as scoring one (the roofline argument). Any drafter turns that idle compute into accepted tokens. The speedup from speculative decoding is

where $c$ is the cost of one drafting step relative to a target step. A separate 1B draft for an 8B target has $c$ around 0.1 to 0.2 once fixed per-step overheads are counted. Medusa's heads run in the same pass as the target, so their marginal cost is a few matrix multiplies ($c\approx 0$ for drafting, with the tree verification cost moved into the next target pass). EAGLE's draft is one transformer layer, a few percent of an 8B model's 32 layers. Both also reuse the target's representation of the context, which a separate small model has to rebuild with far less capacity. The price is coupling: the drafter only works on the exact model whose hidden states it was trained on.

Medusa: heads, losses and trees#

Let $h_t\in {\mathbb{R}}^d$ be the target's last hidden state at position $t$ and $V$ the vocabulary size. The $k$-th Medusa head predicts the token at position $t+k+1$ (the base LM head already predicts $t+1$):

$W_1^{(k)}$ is initialised to zero and $W_2^{(k)}$ to a copy of the LM head, so every head starts out predicting the next token and learns to look further ahead. Medusa-1 freezes the backbone and trains only the heads with a weighted cross-entropy against the ground-truth future tokens $y$:

The decaying weights reflect that far heads are less accurate and matter less, since a draft is only useful up to its first rejection. Medusa-2 trains the backbone jointly, $\mathcal{L}={\mathcal{L}}_{\text{LM}}+{\lambda }_0{\mathcal{L}}_{\text{Medusa-1}}$, with a lower learning rate for the backbone and a warm-up that first trains the heads alone. Joint training raises head accuracy but changes the model, so its outputs are no longer those of the original: the speedup is no longer free of quality risk, and Medusa-2 uses self-distillation with a KL term to the original model's predictions to contain it.

Trees. A single top-1 guess per head is fragile. Medusa takes the top $s_k$ candidates of head $k$ and verifies continuations from their Cartesian product (or a pruned subset) in one pass, flattening the tree into a sequence and using an attention mask in which each node sees only its ancestors (Figure 2). If head $k$ contains the target's token among its top $s_k$ with probability $h_k(s_k)$, and hits are roughly independent across heads, the expected accepted length per step is

Worked example. Suppose three heads have top-1 hit rates of 0.6, 0.4 and 0.3. A chain of top-1 guesses yields $1+0.6+0.6\times 0.4+0.6\times 0.4\times 0.3=1.91$ tokens per step with 3 candidate tokens to verify. Now take the top 3 from head 1 (say hit rate 0.8) and the top 2 from head 2 (0.55), keeping top-1 for head 3: $1+0.8+0.8\times 0.55+0.8\times 0.55\times 0.3=2.37$ tokens per step, for a tree of $3+3\times 2+3\times 2\times 1=15$ candidate tokens. Verification of 15 tokens instead of 3 is still nearly free at batch 1 but not at batch 64, which is why Medusa searches for sparse trees that maximise expected acceptance for a fixed node budget (its paper compares a 64-node sparse tree against 256-node dense trees). The hit rates here are illustrative; measure your own per head.

Typical acceptance. Exact speculative sampling over a tree requires the multi-draft rejection schemes of SpecInfer. Medusa instead proposes a cheaper, lossy rule: accept candidate $x$ if $p_{\text{orig}}(x)>\min (ϵ,\delta \exp (-H(p_{\text{orig}})))$, where $H$ is the entropy of the target's distribution and $ϵ,\delta$ are thresholds. At high entropy more candidates qualify, at low entropy only the likely ones. It raises acceptance but does not preserve the target's sampling distribution, so treat it as a quality knob to evaluate, not a free speedup. With greedy decoding, acceptance is simply "the candidate equals the target's argmax", which is exact.

EAGLE: autoregression over features#

Let $f_i$ be the target's second-to-top-layer feature at position $i$ and $t_i$ the token there. EAGLE's draft model concatenates each feature with the embedding of the next token, reduces the $2d$-dimensional vector to $d$ with a fully connected layer, and passes it through one decoder layer to predict the next feature, reusing the target's frozen embedding and LM head to turn features into tokens:

Feeding the token $t_{i+1}$ is the fix for feature uncertainty: the same feature can lead to different sampled tokens, and without knowing which one was sampled the next feature is ambiguous. Training combines a regression loss on features and a classification loss on tokens,

with uniform noise $U(-0.1,0.1)$ added to input features so the draft tolerates its own errors at drafting time. At inference the draft runs $m$ steps on its own predicted features to build a tree of depth $m$, which the target verifies with standard speculative sampling, so EAGLE is lossless. EAGLE-2 grows the tree dynamically using the draft's confidence, which the authors found approximates acceptance probability well. EAGLE-3 removes ${\mathcal{L}}_{\text{reg}}$ entirely, feeds a fusion of low, middle and high-layer features, and during training rolls the draft forward on its own outputs (training-time test) so that training matches multi-step drafting; this is what let it keep improving with more training data.

Side by side

Toy run: independent heads, sequential heads, and the data they see#

To isolate two effects, sequential conditioning and training-data alignment, we reused the 64-token bigram "fine-tuned target" from the speculative decoding article: a base model plus a strong shift on half of the contexts. We trained two kinds of drafter with the same rank-8 capacity. The Medusa-style drafter has four independent heads, head $k$ mapping the current token to the token $k$ steps ahead. The EAGLE-style drafter is one next-token head applied autoregressively to its own drafts. Each was trained either on 200,000 tokens of "fixed" text sampled from the base model (the analogue of training on a public chat dataset the fine-tune did not write) or on 200,000 tokens sampled from the fine-tuned target itself (self-distillation). We then ran exact speculative sampling with four draft tokens per step against the fine-tuned target for 30,000 steps. The script is seeded; the numbers are its output.

Two lessons transfer. First, independent heads hit a ceiling that more data cannot fix: even exact marginal distributions give only 0.47 acceptance at depth 2 under sampling, because head 2 must hedge over every possible token at depth 1 while the target has already committed to one. That is the motivation for Hydra (Ankner et al., 2024), whose sequentially dependent heads improved throughput by up to 1.31 times over Medusa, and for EAGLE's design. Medusa recovers much of the gap in practice with trees, greedy decoding and typical acceptance, which this toy does not model. Second, drafter training data must come from the model being served. Training the sequential drafter on the fine-tuned target's own outputs instead of base-model text raised tokens per step from 3.04 to 4.56. In the EAGLE paper, training on target-generated data instead of fixed ShareGPT data improved LLaMA2-Chat 7B only slightly (2.78 to 2.88 times), because a chat model's outputs are close to ShareGPT; a domain fine-tune is far from it, and the toy shows what happens then.

Training heads for a fine-tuned model#

Every Medusa or EAGLE drafter is bound to the hidden states of one checkpoint. After SFT, DPO or RL, the features and the output distribution both move, and heads trained for the base model lose acceptance. A practical recipe:

1. Freeze the exact serving checkpoint. Merge LoRA adapters if you serve merged weights; if you serve base plus adapter, capture hidden states with the adapter active. If you serve an FP8 or 4-bit model, capture features from that model or at least validate on it.
2. Build self-distillation data. Take 50k to 100k prompts from your domain (real traffic, the SFT prompts, or synthetic ones) and generate responses with the fine-tuned model at the serving temperature. Medusa's self-distillation used about 100k samples; EAGLE-3 regenerates responses with the target. Plain ground-truth SFT targets are a worse choice: they are not what the model will say.
3. Train the drafter with the backbone frozen (Medusa-1, EAGLE, EAGLE-3). This keeps the target's outputs exactly as they were, so no quality re-evaluation is needed beyond a sanity check. For reference, EAGLE was trained with AdamW at a learning rate of 3e-5, betas (0.9, 0.95) and gradient clipping at 0.5; Medusa-1 heads for Vicuna-7B trained in about 5 hours on one A100.
4. Measure acceptance length on held-out prompts at the serving temperature and batch size, and compare tokens per second with and without speculation under realistic concurrency.
5. Retrain on every new checkpoint. The cost is a small fraction of the fine-tune itself; build it into the release pipeline.

Joint training is a different decision. Medusa-2, or training multi-token prediction into the model as Gloeckle et al. (2024) and DeepSeek-V3 did, changes the model. Gloeckle et al. found 4-token prediction made models up to 3 times faster at inference and improved code benchmarks; DeepSeek-V3's MTP modules can be kept for speculative decoding or discarded. If you add such heads during your own fine-tune, the drafter comes for free with every checkpoint, but the main model's quality must be evaluated like any other training change. Medusa's guidance for joint training with self-distillation is to use LoRA without quantisation on the backbone, so the teacher is not a quantised model.

Many adapters. With multi-LoRA serving, one drafter trained on base features sees shifted features for every adapter. Options: one drafter per high-traffic adapter, a drafter trained on a mixture of adapters' features, or a model-free drafter (n-gram prompt lookup) for the long tail. See multi-LoRA serving.

Templates and tool calls. If the fine-tune introduced new special tokens or a structured output format, include examples in the drafter's data; formats are highly predictable and often give the best acceptance of all, which combines well with constrained decoding.

Tuning knobs#

Failure modes and how to spot them#

Stale heads. Acceptance length drops after a new fine-tune ships while outputs look fine. Log mean acceptance length per model version; alarm on a drop.

Feature mismatch. Heads trained on bf16 features served on a quantised model, or on base features served with an adapter. Acceptance is lower than in offline evaluation; re-capture features from the served configuration.

High-load slowdowns. Trees of dozens of candidates per sequence turn verification compute-bound at large batch. Measure throughput at your real concurrency, not batch 1.

Quality drift from lossy acceptance or joint training. Typical acceptance and Medusa-2 can change outputs. Run your task evaluation with the exact serving configuration.

Temperature sensitivity. Acceptance falls as sampling temperature rises because the target's distribution flattens. Benchmark at the temperature you actually use.

When to use which#

Implementation sketch#

Training a Medusa-1 head set on frozen features takes a few lines on top of any causal LM:

import torch, torch.nn as nn, torch.nn.functional as F

class MedusaHead(nn.Module):
    def __init__(self, d, lm_head):
        super().__init__()
        self.w1 = nn.Linear(d, d); nn.init.zeros_(self.w1.weight); nn.init.zeros_(self.w1.bias)
        self.out = nn.Linear(d, lm_head.out_features, bias=False)
        self.out.weight.data.copy_(lm_head.weight.data)          # start as a next-token head
    def forward(self, h):
        return self.out(F.silu(self.w1(h)) + h)

def medusa1_loss(heads, hidden, input_ids, mask, decay=0.8):
    """hidden: (B, T, d) last hidden states of the frozen fine-tuned model; mask: (B, T) completion tokens."""
    loss = 0.0
    for k, head in enumerate(heads, start=1):                  # head k predicts token t + k + 1
        logits = head(hidden[:, : -(k + 1)])
        target = input_ids[:, k + 1 :]
        m = mask[:, k + 1 :].float()
        ce = F.cross_entropy(logits.transpose(1, 2), target, reduction="none")
        loss = loss + (decay ** k) * (ce * m).sum() / m.sum()
    return loss
# hidden comes from model(input_ids, output_hidden_states=True).hidden_states[-1] under torch.no_grad(),
# on self-distilled responses generated by the same fine-tuned checkpoint.

For EAGLE and EAGLE-3, the official SafeAILab/EAGLE repository and SGLang's SpecForge provide training scripts that capture target features and train the draft layer. Serving is a configuration switch in vLLM:

from vllm import LLM
llm = LLM(
    model="path/to/finetuned-8b",
    speculative_config={"method": "eagle3", "model": "path/to/eagle3-head-for-this-finetune",
                        "num_speculative_tokens": 5},
)

In POFE, drafter training is a step that runs after a fine-tune finishes: it self-distills on the run's prompt set, trains the head against the exact checkpoint you will deploy, and reports acceptance length next to the model's evaluation results.

Further reading#

Cai et al. (2024) for Medusa, including self-distillation and typical acceptance. Li et al. (2024) for EAGLE, then the EAGLE-2 and EAGLE-3 papers for dynamic trees and training-time test. Ankner et al. (2024) for why sequential dependence matters. Gloeckle et al. (2024) for training multi-token prediction into the model itself. Miao et al. (2023) for exact tree verification.

Frequently asked questions#

References#

1. Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, Tri Dao (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arxiv.org/abs/2401.10774
2. Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang (2024). EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. arxiv.org/abs/2401.15077
3. Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang (2024). EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees. arxiv.org/abs/2406.16858
4. Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang (2025). EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test. arxiv.org/abs/2503.01840
5. Zachary Ankner, Rishab Parthasarathy, Aniruddha Nrusimha, Christopher Rinard, Jonathan Ragan-Kelley, William Brandon (2024). Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding. arxiv.org/abs/2402.05109
6. Mitchell Stern, Noam Shazeer, Jakob Uszkoreit (2018). Blockwise Parallel Decoding for Deep Autoregressive Models. arxiv.org/abs/1811.03115
7. Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve (2024). Better & Faster Large Language Models via Multi-token Prediction. arxiv.org/abs/2404.19737
8. DeepSeek-AI (2024). DeepSeek-V3 Technical Report. arxiv.org/abs/2412.19437
9. Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, et al. (2023). SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification. arxiv.org/abs/2305.09781
10. SafeAILab (2025). EAGLE official implementation (SafeAILab/EAGLE). github.com/SafeAILab/EAGLE
11. SGLang project (2025). SpecForge: training speculative decoding models for SGLang. github.com/sgl-project/SpecForge

Related reading