Medusa speeds up generation by bolting a handful of extra prediction heads onto a frozen transformer, guessing several future tokens at once and verifying them in a single pass — no second draft model to host. The companion article Medusa: the math derives the head shapes, the tree-attention mask, and the speedup arithmetic. This piece goes past that into the parts that decide whether Medusa actually works in practice: why independent heads get less reliable the further ahead they reach, how a sparse candidate tree is built and checked in one forward pass, why Medusa trades exact rejection sampling for a looser ‘typical acceptance’ test, and above all the two training recipes — Medusa-1 with a frozen backbone and Medusa-2 with joint fine-tuning — plus the self-distillation trick that makes the heads match the model they are grafted onto.
What this deep dive assumes
Start from the shape the base article establishes. A decoder produces a final hidden state h_t at each position, and the original language-model head maps h_t to a distribution over the next token, t+1. Medusa adds K extra heads: head i reads the same h_t and predicts the token at t+i+1, so from one forward pass you now have distributions for t+1 through t+K+1.
That offset detail is worth stating precisely: the base model still owns t+1; Medusa head i owns t+i+1. Everything that follows is about turning those parallel-but-independent guesses into correct tokens cheaply. We take the head shapes and tree-attention mask as given and spend our words on the training the base article does not cover.
The heads are residual blocks, not linear layers
A Medusa head is deliberately tiny but not trivial. Rather than a bare projection, each head is a single residual block: it takes h_t, passes it through one learned linear layer with a SiLU non-linearity, adds the result back to h_t, and only then applies a language-model projection to logits. In shorthand, head_i(h_t) = LMproj( h_t + SiLU(W_i · h_t) ).
The residual connection is the whole trick. It initialises the head close to the identity, so at the start of training a head behaves almost like the model’s own next-token predictor and then learns a small correction for its offset. That keeps each head cheap while giving it enough capacity to specialise. It also explains why the heads add little inference cost: they run in parallel off one hidden state, so the extra work is a few small matmuls, not another pass through the network.
The conditional-independence gap
Here is the flaw the deep dive has to confront. Every Medusa head conditions on the same h_t and nothing else. Head 2 predicting t+3 never sees what head 1 guessed for t+2. So the heads produce marginal predictions, not a joint one — they cannot model ‘if the next word is New, the one after is probably York.’
The consequence is a steady decay: head 1 is fairly accurate, head 2 less so, and by head 4 or 5 the top guess is often wrong, because predicting several tokens ahead from a single frozen context is harder and the head cannot see the intervening choices. This is precisely the gap EAGLE closes with an autoregressive feature recurrence, and it is why Medusa leans on a tree of several candidates per position rather than betting on one chain: breadth compensates for the heads’ inability to condition on each other.
From top-k heads to a sparse candidate tree
Because no single head is reliable, Medusa keeps the top few guesses from each. Take the top k_1 tokens from head 1, the top k_2 from head 2, and so on. The naive combination is a Cartesian product — every head-1 guess paired with every head-2 guess — which is k_1 × k_2 × … candidate continuations, exploding fast.
Medusa does not verify them all. It prunes to a fixed sparse tree chosen ahead of time: paths through more-confident, earlier heads get more branches; unlikely deep paths get dropped. The tree is laid out as a flat sequence of nodes with a matching attention mask and position ids, so a node attends only to its own ancestors, and the entire tree is pushed through the backbone in one forward pass. The base article details that mask; the point here is that tree shape is a tuned budget, trading more candidates (higher acceptance) against a heavier pass.
Typical acceptance, not exact rejection sampling
Classic draft-model speculative decoding uses rejection sampling: each drafted token is accepted or rejected with a probability crafted so the final output is distributed exactly as if you had sampled from the target model token by token. That guarantee is its selling point — speculation changes speed, never the distribution.
Medusa deliberately abandons that guarantee. Its heads do not define a clean proposal distribution you can do exact rejection against, so instead it uses typical acceptance: accept a proposed token if the original model’s probability for it clears a threshold of the form min(ε, δ · exp(−H)), where H is the entropy of the model’s distribution at that step. An absolute floor ε keeps confidently-predicted tokens flowing; the entropy term loosens the bar when the model itself is uncertain and tightens it when the model is peaked. The trade is explicit: you give up exact distribution matching for a higher acceptance rate and more speed. In practice it tracks greedy or low-temperature decoding closely, the regime most deployments run.
Medusa-1: training on a frozen backbone
The first and simplest recipe freezes the entire pretrained model and trains only the new heads. Each head i is supervised with cross-entropy against the true token at offset t+i+1, summed over positions, with a per-head weight that decays as the offset grows — roughly Σ_i λ^i · CE_i — because the far heads are intrinsically harder and should not dominate the gradient.
Freezing the backbone has two virtues. It is cheap: only a few small matrices train, so you can fit heads onto a large model on modest hardware, often with the backbone quantised. And it is safe — because the original weights never move, the model’s quality cannot regress; Medusa-1 either speeds things up or, at worst, does nothing. The ceiling is that frozen features were never shaped to be predicted several steps ahead, so the heads’ acceptance rate, and thus the speedup, is capped below what a co-adapted model could reach.
Medusa-2: joint fine-tuning without wrecking the model
Medusa-2 lifts that ceiling by training the heads and the backbone together, so the shared representation h_t becomes a little more predictive of the future and the heads get sharper. The obvious danger is equally large: let the head loss pull on the backbone freely and you degrade the model’s core next-token quality — a faster model that is now slightly worse is a bad trade.
The recipe is therefore careful. The objective combines the original language-model loss (to preserve backbone quality) with the weighted head losses, and it uses a warm-up that trains the heads first before the backbone is allowed to move, plus a smaller learning rate on the backbone than on the heads. The result is a model whose own outputs are essentially unchanged but whose hidden states carry more forward-looking signal, giving higher acceptance lengths than Medusa-1. The cost is a real fine-tune: more compute, more data, more tuning care.
Self-distillation: matching the model to itself
Both recipes need training targets, and the naive choice — some public instruction dataset — has a subtle problem. The heads’ job is to predict what this model will say next, but public text is what humans said. Training on a distribution the backbone would not itself generate teaches the heads the wrong thing, and it is worse still when you lack the model’s original fine-tuning data.
The fix is self-distillation: take a set of prompts, let the model generate its own completions, and train the heads (and, in Medusa-2, the backbone) on that self-produced text. Now the targets are drawn from the model’s own distribution, so a head learns to anticipate the model’s actual next tokens rather than a stranger’s. This closes the distribution gap that would otherwise cap acceptance and sidesteps the missing-data problem — the model is its own teacher. For Medusa-2 it also protects backbone quality, since fine-tuning on the model’s own outputs pulls it toward, not away from, its existing behaviour.
What it costs on a CPU-bound SLM
On a small model running on a CPU, generation is memory-bandwidth-bound: each token drags the whole weight set through cache, and that traffic, not arithmetic, sets the pace. Medusa fits this regime well because verifying a tree of candidate tokens reuses that same weight read for many positions at once — you amortise the expensive memory pass over several confirmed tokens instead of one. The figure to watch is expected acceptance length: the average tokens confirmed per pass, which is the multiplier on throughput. Per-head accuracy decays with offset, so K and tree size are tuned against that decay — bigger trees raise acceptance but make every pass heavier.
The honest caveats: the heads add parameters and a little matmul per step, the verification pass is wider than a plain decode step, and if acceptance is low you pay that width for little return. Medusa-1 is the pragmatic starting point — cheap to add, impossible to make the model worse — and self-distilled Medusa-2 is the upgrade when you can afford the fine-tune. Either way the appeal on a CPU is the same: a real speedup with no second model to load, quantise, and keep resident alongside the one you already cannot fit comfortably.
t+i+1 in parallel from one hidden state, but because they cannot condition on each other their accuracy decays with distance — which is why Medusa verifies a sparse tree of several candidates in a single pass rather than one chain. It swaps exact rejection sampling for typical acceptance, trading distribution-matching for a higher acceptance rate. The real engineering is the training: Medusa-1 freezes the backbone and trains only the heads (cheap and quality-safe, but capped), while Medusa-2 jointly fine-tunes for higher acceptance at the risk of degrading the model — a risk tamed by a combined loss, warm-up, and a gentler backbone learning rate. Self-distillation ties it together by training on the model’s own outputs so the heads learn the distribution they actually serve. On a CPU-bound SLM the win is concrete: more tokens per memory pass, and no second draft model to keep resident.