Medusa,
explained.

Rung 4's speed-up is real and its operational cost is a second model — one to obtain, keep aligned with the target, hold in memory, and version alongside it. For many deployments that is the reason it is not used.

The change is to remove it. Attach extra prediction heads to the existing model, each trained to predict a token further ahead than the last, so a single forward pass emits not just the next token but guesses at several following ones. The draft comes from the model itself, at the cost of a few small heads, and there is nothing separate to align because the heads sit on the representation the model already computed.

Head accuracy necessarily decays with distance — predicting four tokens ahead from the current state is harder than one — so instead of committing to one guessed continuation, the top candidates from each head are combined into a tree of possible continuations, and a single pass with an appropriately masked attention pattern verifies many branches at once. Stated in rung 4's units, the accepted length per pass comes from covering several plausible continuations rather than from one accurate draft — breadth substituting for the draft model's depth.

Recall the bottleneck: language models generate one token at a time, and each token forces the whole model to be read from memory, leaving the hardware's compute mostly idle. Speculative decoding fills that idle compute by having a small, fast draft model guess several tokens ahead, which the big model then verifies in parallel. It works — but the draft model is a real operational cost. You have to find or train one that is both cheap and well-aligned with the target, serve it alongside, and maintain the pairing as models change.

Medusa removes the second model. Its insight: the big model already computes a rich hidden state at each position — enough, with a little extra machinery, to predict not just the next token but the one after, and the one after that. So instead of a separate drafter, give the model a few lightweight heads that read that hidden state and guess further ahead. The model drafts for itself.

A standard model has one output head predicting the next token. Medusa adds K more, each a small feed-forward layer on the same final hidden state, trained to predict the token at position +2, +3, …, +K+1. They are cheap to add (in Medusa-1 the backbone is frozen and only the heads train; Medusa-2 fine-tunes both) and cheap to run (a few extra small layers). In one forward pass, the model now emits its usual next token plus K guesses for the tokens beyond it.

The catch is accuracy. Predicting the very next token is relatively easy; predicting five tokens ahead, with all the uncertainty of the words in between, is much harder. So head accuracy fades with distance — the first head is fairly reliable, later heads much less so. That decay is exactly what governs how many tokens Medusa can actually keep per pass, and it's the number the math below tracks.

Because the far heads are unsure, betting everything on their single top guess is wasteful — often it's wrong, and the whole speculation stops there. Medusa hedges: each head proposes its top few candidates, and the combinations form a tree of possible continuations (this head's option A followed by that head's options X or Y, and so on). A single guess per head accepts if it's right with probability a; offering s candidates accepts if any of them is right — probability 1 − (1 − a)^^s, which climbs fast.

The clever part is verifying this whole tree in one forward pass. Medusa builds a special attention mask — tree attention — so that each candidate token only attends to its own ancestors in the tree, letting many candidate paths be scored simultaneously without interfering. One pass, many hypotheses checked, and the longest fully-correct path is accepted. It is how Medusa turns uncertain heads into reliable speedups.

Verification keeps the output honest. The tree of candidates is scored by the true model, and Medusa accepts the longest path whose tokens the model agrees with — falling back to the guaranteed correct next token when the heads miss. With greedy acceptance, a position is only kept if every head before it was also right, so each head's contribution is the cumulative product of accuracies. Summing those, plus the always-correct base token, gives the expected tokens per pass:

Because each term multiplies one more (sub-1) accuracy, later heads add less and less — a fading geometric-like series. With four heads at accuracies 0.8/0.65/0.5/0.35 the model produces ~2.7 tokens per pass instead of 1; perfect heads would give K+1. Tree candidates raise each head's effective accuracy toward 1, pushing the sum higher. Crucially, the accepted tokens are exactly the model's own, so under greedy decoding the output matches the base model's own (with the typical-acceptance rule used for sampling, it is only close to it, not guaranteed identical). The runnable version below computes both formulas.

Medusa (Cai, Li, Geng, Peng, Lee, Chen and Dao) keeps the base model and its usual next-token head, and adds K decoding heads on top of the last hidden state h_t. Head k guesses the token k + 1 positions ahead, so head 1 predicts position t + 2 while the original head predicts t + 1. The procedure:

1. Build the heads. Each head is a single feed-forward layer with a residual connection followed by its own vocabulary projection: p^^(k) = softmax(W₂^^(k)(SiLU(W₁^^(k)h_t) + h_t)). W₂ is initialized from the base model's output head and W₁ at zero, so at the start every head reproduces the base model's next-token prediction and training only has to shift it.
2. Medusa-1: train the heads on a frozen model. The backbone does not change, so the base model's own outputs are untouched. Each head gets a cross-entropy loss against the true token k + 1 places ahead, weighted by λ_k = 0.8^^k so the more predictable near heads dominate. The authors report that heads for a 7B model train in a few hours on one GPU, and the frozen backbone can even be quantized meanwhile.
3. Medusa-2: train heads and backbone together. For higher head accuracy the backbone is fine-tuned jointly, with the ordinary next-token loss kept in the objective so the model's own predictions do not degrade. Two precautions make this work: the heads are warmed up alone first, and they get a higher learning rate than the backbone.
4. Self-distill when the training data is missing. For a chat model whose fine-tuning data is not public, the model generates its own training set from a collection of prompts. For Medusa-2 the next-token loss is then replaced by a KL divergence to the original model's predictions; with LoRA adapters, that teacher is simply the same weights with the adapters switched off.
5. Propose a tree. At inference each head offers its top few tokens. Every combination of one choice per head is a candidate continuation; instead of the full product, the paper keeps a sparse tree built from head accuracies measured on a calibration set, adding the branches most likely to be accepted.
6. Verify the tree in one pass. The tree is flattened into one sequence with an attention mask that lets each candidate see only the prefix and its own ancestors, and position indices are set by depth. One forward pass of the base model scores every node, and the heads read the same pass's hidden states to propose the next tree.
7. Accept the longest valid path. Under greedy decoding a token is valid when it equals the base model's argmax at that node, and the output matches ordinary greedy decoding. For sampling, Medusa can use speculative sampling's rejection rule, which preserves the distribution, or its own typical acceptance, which is faster but approximate.

Typical acceptance

Typical acceptance keeps a candidate whenever the base model gives it probability above min(ε, δ · exp(−H)), where H is the entropy of the base model's distribution at that point. The first token of each step is always the base model's own greedy choice, so every step yields at least one token. Because the threshold drops when the model is uncertain, lower-probability candidates can pass in open-ended contexts, while in a confident, low-entropy context only near-certain tokens clear it.

The limitation later work attacked

Every Medusa head reads the same state h_t and never sees the tokens chosen by the heads before it, so head 2's guess for position t + 3 does not know which token head 1 proposed for t + 2. Hydra makes the heads depend on earlier guesses, and EAGLE drafts autoregressively in the model's feature space, conditioning each step on the sampled token; both report higher acceptance for that reason.

• ht — the backbone's last hidden state at position t; W1 (d × d, zero-initialized) and W2 (d × vocabulary, copied from the output head) — head k's only parameters.
• yt+k+1 — the true token k + 1 positions ahead, the target for head k.
• λk — decaying weights: far heads are less accurate and matter less, because a far guess is used only if every nearer one was accepted.
• LLM, λ0 — the backbone's own next-token loss (a KL to the original model under self-distillation) and its balance against the head losses.
• ε, δ, H — a hard probability floor, the scale of an entropy-dependent floor, and the entropy of the base model's prediction.

At the extremes: with ε and δ near zero almost any candidate passes, so steps get longer and outputs drift from the base model; with high thresholds only near-argmax tokens pass, and decoding approaches greedy decoding.

The paper evaluated Vicuna-7B, 13B and 33B and Zephyr-7B, chat models that differ in size and in how they were trained, on MT-Bench at batch size 1 on A100 GPUs. Speedup is wall-clock generation time against standard decoding of the same model, and quality is checked with MT-Bench's GPT-4-judged scores.

The ablations explain where the speed comes from and where it stops. Sparse trees built from calibration accuracies accepted more tokens per step than dense Cartesian trees with the same number of nodes. Adding nodes kept raising the tokens accepted per step, but past a point real speed fell, because a bigger tree makes every verification pass more expensive: the gain is bounded by the spare arithmetic in the pass. Raising the typical-acceptance threshold improved quality and reduced speed, a direct dial between the two. For Medusa-2, the paper reports that fine-tuning heads and backbone together directly, without the two-stage warm-up, degraded generation quality, which is why the recipe includes it, and self-distillation was what made the method usable on models whose training data was unavailable.

All of these numbers are for one request at a time. Because the acceleration spends the same spare arithmetic as speculative decoding, the gains should be expected to shrink as batch size grows.

The honest comparison against rung 4 has costs on both sides. The heads are weaker predictors than a real small model, because they are shallow functions of a state computed for a different purpose, so acceptance per position is lower. Tree verification recovers that by checking more candidates, which spends the arithmetic headroom the whole approach depends on — and at high batch sizes, where that headroom is gone, the tree's cost stops being free before a single draft's does. And unlike a drop-in draft model, the heads must be trained for the specific model, so this is not something you switch on for an arbitrary checkpoint.

The acceptance criterion is also worth reading carefully. The variant that makes this fastest relaxes rung 4's exact distributional guarantee in favour of accepting tokens that are merely sufficiently likely — which is usually unobjectionable and is a real difference in kind from a method that changes nothing. Anyone who chose rung 4 because of the guarantee should not assume it carries over.

How it held up: self-drafting became a standard family, and the tree-verification idea outlived the specific design — verifying a branching set of continuations in one masked pass is now common in speculative decoding generally, including with external drafts. Read rungs 4 and 5 together as the ladder's closing argument: once memory traffic is the constraint, the way to go faster is to get more tokens out of each weight read, and where the guesses come from is an engineering choice.

Medusa and draft-model speculative decoding are two routes to the same goal; each has a personality:

Where a well-matched draft model can sometimes propose longer accurate runs, Medusa wins on operational simplicity — one model, a handful of heads, no second system. For many deployments that trade is exactly right.

Medusa is part of a fast-moving family of self-speculative decoding methods — including EAGLE, Hydra, and lookahead decoding — all aiming to get speculative decoding's speed without the burden of a separate draft model. They differ in how they draft (extra heads, feature-space prediction, n-gram pools) but share the skeleton established by speculative decoding: propose cheaply, verify in parallel, keep what's correct.

Its contribution to that lineage is a clean, practical recipe: reuse the model's own hidden state through lightweight heads, and hedge their uncertainty with tree attention. That combination — self-drafting plus batched candidate verification — turned out to be a sweet spot of simplicity and speed, and it's why Medusa became a widely-used, easy-to-adopt inference accelerator.

Work each one out before you open it — a number, a prediction or a reason, not a definition. Producing the answer is what makes it stick. Your progress is kept on this device.

Five papers about the bill. Each one buys throughput back from a different part of the stack: the weights, the memory hierarchy, the KV cache, the decode loop.

1. 1Knowledge Distillation2014
2. 2FlashAttention2022
3. 3PagedAttention2023
4. 4Speculative Decoding2023
5. 5Medusa2024

All 110 papers, in 19 ladders →