Medusa — Multi-Head Speculative Decoding Math

The bottleneck Medusa attacks

Autoregressive decoding is memory-bandwidth-bound: to emit one token the model streams every weight (and the KV cache) from memory, does a tiny amount of arithmetic, and repeats. The GPU’s compute units sit almost idle — a single-token step’s arithmetic intensity is far below the hardware’s FLOPs-per-byte ratio. So generating n tokens costs roughly n full weight reloads, and latency scales linearly with output length no matter how much spare compute exists.

Every speculative method exploits the same slack: a forward pass that verifies k tokens at once costs barely more than a pass that emits one, because both are dominated by loading the same weights — the extra tokens ride along nearly free on unused compute. If you can cheaply guess the next few tokens and verify them in parallel, you amortize one weight reload over several accepted tokens. Medusa’s contribution is a guessing mechanism that needs no separate model: extra heads on the model you already have.

The Medusa head: math and shapes

Let h_t ∈ R^d be the backbone’s final hidden state at position t (the vector its own LM head would turn into the next-token distribution). Medusa attaches K heads. Head k is a single residual block feeding an unembedding:

p_t^(k) = softmax( W2_k · ( SiLU(W1_k · h_t) + h_t ) )

  h_t   : [d]           backbone final hidden state
  W1_k  : [d, d]        residual projection (init ≈ 0)
  W2_k  : [V, d]        unembedding (init = base LM head)
  p_t^(k): [V]          distribution over the (k+1)-th future token

Two initializations make this work. W1_k starts near zero, so initially SiLU(0)+h_t ≈ h_t and the head behaves like the original LM head; W2_k copies the base model’s unembedding, reusing its token geometry. Only the heads train — the backbone stays frozen, so training is cheap (hours, one GPU) and cannot regress the base model. Head k predicts the token at t+k+1; with the backbone’s own t+1 prediction, one pass drafts up to K+1 positions. Crucially the heads are conditionally independent given h_t — head 2 never sees head 1’s guess — Medusa’s main quality limitation, addressed structurally by the tree.

From heads to a candidate tree

A point prediction per head would be brittle: if head 1 is wrong, everything downstream is wasted. Instead Medusa keeps the top-s_k tokens from each head and expands their combinations into a tree. With per-head widths (s_1, …, s_K) the Cartesian product has ∏_k s_k leaf continuations — e.g. (4,3,2,2) gives 48 candidate paths from a single hidden state.

The tree shares prefixes: continuations that begin with the same head-1 token share that node, those agreeing on heads 1–2 share those two, and so on, so the tree has far fewer nodes than leaves×depth — each token appears once. In practice you skip the full product and keep only the highest-probability nodes (ranked by the product of head confidences) up to a fixed budget, since low-probability branches are almost never accepted and only cost compute. The shape is fixed ahead of time so the attention mask can be precomputed.

Tree attention: the mask math

All tree nodes are flattened into one input sequence and run through the backbone in a single forward pass. The trick is the attention mask M: a node attends to the prompt and to its own ancestors along its path from the root — and to nothing on sibling branches. Each node’s hidden state is then computed as if only its own candidate prefix existed, so many mutually-exclusive continuations are scored at once without contaminating each other.

M[i, j] = 0   if node j is an ancestor of node i (or i itself)
M[i, j] = -∞  otherwise

attn_i = softmax( (q_i K^T + M[i]) / sqrt(d_k) ) V

Ordinary causal decoding uses a lower-triangular mask — a linear chain where every position attends to all earlier ones. Tree attention generalizes that to a partial order: still ‘attend to your ancestors,’ but ancestry follows tree edges, not sequence position, so two cousins on different branches sit adjacent yet cannot see each other. Because the shape is fixed, M and the positional indices (each node’s position id = prompt length + its depth) are precomputed once. A tree of m nodes thus verifies ∏_k s_k candidate paths at the cost of one m-token forward pass — the core efficiency of Medusa.

Typical acceptance, not exact match

Classic speculative decoding accepts draft tokens by a rejection rule that provably reproduces the target’s exact sampling distribution — elegant but conservative, and it needs the draft’s probabilities. Medusa instead uses typical acceptance, a threshold on the backbone’s own probability for the proposed token given the accepted prefix:

accept x  iff  p_base(x | prefix) > min( ε,  δ · exp(-H(p_base)) )

  H(p) = -Σ_x p(x) log p(x)     entropy of the base next-token dist
  ε, δ  : tunable constants (e.g. ε=0.09, δ=0.3)

The threshold adapts to uncertainty. When the model is confident (low entropy) the bar is high, so only near-certain tokens pass; when it is genuinely uncertain (high entropy, many acceptable words) the bar drops and more candidates qualify. You walk each tree path from the root, accept the longest prefix whose every token clears its threshold, then take the model’s own next token as a free bonus at the first rejection. This does not match the base sampling distribution token-for-token — it is a quality-preserving heuristic that at temperature 0 collapses to exact greedy matching. Dropping the exactness guarantee is what lets Medusa accept longer prefixes and go faster.

Expected speedup, with a worked example

Let τ be the mean accepted length — tokens confirmed per verification pass (the accepted prefix plus the bonus token). Without Medusa, τ tokens need τ passes. With Medusa they need one pass, slightly costlier because it processes m tree tokens and runs K heads. With per-step overhead factor c ≥ 1:

speedup ≈ τ / c

  τ : accepted tokens per pass   (typically 2.3 – 3.6)
  c : overhead of wider pass + heads (typically 1.05 – 1.2)

Overhead stays small because decoding is memory-bound: a few dozen tree tokens and a few heads barely change a step dominated by streaming weights. Worked case: a one-token step takes 10 ms; add K=4 heads and a 40-node tree and the wider pass measures 11.5 ms, so c = 1.15. If the heads clear thresholds for τ = 2.8 tokens on average:

baseline for 2.8 tokens : 2.8 × 10 ms = 28.0 ms
Medusa per pass         :                 11.5 ms   → 2.8 tokens
speedup = 28.0 / 11.5   = 2.43×   (check: τ/c = 2.8/1.15 = 2.43×)

Growing the tree to 100 nodes might lift τ to 3.1 but c to 1.25, giving only 2.48× — the tree is crossing into the compute-bound regime where extra nodes stop being free. The sweet spot is the largest tree still riding the memory-bound slack, and it is workload-dependent: code and templated text accept far longer runs than open-ended prose.

Versus draft-model speculative decoding

Standard speculative decoding pairs the big target model with a separate small draft model. The draft autoregressively generates γ tokens (so its guesses are sequentially conditioned — token 2 sees token 1), the target verifies them in one pass, and a rejection-sampling rule accepts a prefix while provably preserving the target’s exact distribution. That correctness guarantee is its headline advantage.

The costs are practical. You must find or train a draft model small enough to be cheap yet aligned enough to be accepted often, and host its weights and KV cache alongside the target. Medusa deletes that model: its ‘draft’ is a few megabytes of heads sharing the backbone’s hidden state, trained in hours, guaranteed to share its vocabulary and tokenizer. The trade is guess quality: Medusa’s heads are conditionally independent, so a raw head chain is weaker than a draft model’s sequential output — exactly why Medusa leans on the tree to hedge across many parallel guesses and on typical acceptance to keep prefixes long.

Versus EAGLE

EAGLE also drops the separate draft model, but fixes the conditioning problem differently. Instead of predicting tokens in parallel, it runs a small autoregressive module at the feature level: it predicts the next hidden state, feeds that predicted feature (plus the sampled token’s embedding) back into itself, and rolls forward. Because each draft step is conditioned on the previous predicted feature, EAGLE’s guesses are sequentially coherent in a way Medusa’s independent heads are not, typically buying a higher mean accepted length τ.

The contrast is a clean design axis. Medusa drafts non-autoregressively — parallel heads, one shot — and recovers coherence structurally through the candidate tree. EAGLE drafts autoregressively in a cheap feature space and recovers coherence temporally, at the cost of a small sequential module. Both verify against the frozen backbone. Medusa’s bet is to keep drafting embarrassingly parallel and let a tree mask plus typical acceptance turn many independent cheap guesses into long runs.

Practical notes and pitfalls

Tree budget dominates. The biggest knob is the node budget and per-head widths. Too small and τ collapses; too large and you cross into the compute-bound regime where c grows faster than τ. The optimal tree is often sparse and skewed — wide at head 1, narrowing with depth, since deep agreement is rare. Acceptance is not free correctness. Typical acceptance changes the output distribution versus pure sampling; if you need exactness (reproducibility, certain evals) prefer rejection-based speculative decoding or run Medusa at temperature 0, where it is exactly greedy.

Workload and batch sensitivity. The quoted ‘2–3×’ are averages; predictable text (code, JSON, boilerplate) can exceed them while high-entropy creative generation lands lower, since τ tracks predictability. And as batch size grows, decoding drifts from memory-bound toward compute-bound, the free-token slack shrinks, and Medusa’s advantage narrows — it shines most at low-latency, small-batch serving. Measure end-to-end wall-clock, not just τ, before trusting the win.