All posts

Decoding Medusa

/AI/4 min read

Extra prediction heads on the same model replace a separate draft model — 134 MB instead of 2 GB. The acceptance chain multiplies, so the fourth head adds 0.044 tokens and the sixth adds 0.001.

Generating a token requires one pass over the whole model, and that pass is dominated by moving weights out of memory rather than by arithmetic. The GPU does a small amount of work and waits.

Which means: if you could check several proposed tokens in the same pass, the extra checking would be nearly free. That is the idea behind every method in this family. Medusa's contribution is where the proposals come from.

The problem with a separate drafter

The established approach uses a small draft model to propose tokens and the large model to verify them. It works, and it costs:

memory
a 1B draft model beside a 7B target2.0 GB
four extra heads on the target itself134 MB

15× less — and memory is the binding constraint on how many requests a serving GPU can hold.

The other costs are not measured in bytes. A separate model has to be chosen, trained, versioned and deployed alongside the target, and it has to stay behaviourally close to it — a drafter that drifts proposes tokens that get rejected, and a rejected proposal is wasted work.

Heads instead

Medusa attaches small heads to the frozen base model, reading the same final hidden state.

The original head predicts the next token, as always. Head 1 predicts the token after that. Head 2 predicts the one after that. All of them read the same hidden state produced by the same forward pass, so the proposals cost almost nothing beyond a small matrix multiply each.

Each head is one transform at hidden size — reusing the base model's vocabulary projection, which is why it is 33 MB rather than 262 MB.

Nothing about the base model changes. The heads are trained on top of it, and the tokens they propose still have to be approved by the base model's own distribution, so the output is what the base model would have produced.

The acceptance chain, and why four heads

Here is the arithmetic that decides how many heads are worth adding.

A proposal is accepted only if every proposal before it was. So the contributions multiply, and each head's realistic accept rate is lower than the last because it is predicting further ahead:

accept ratechainedrunning total
base headalways1.0001.000
head 10.720.7201.720
head 20.550.3962.116
head 30.400.1582.274
head 40.280.0442.319
head 50.190.0082.327
head 60.120.0012.328

Four heads yield 2.32 tokens per forward pass — a 2.32× speedup, which matches the published range of roughly 2 to 3×.

And the reason head counts stop around four or five is right there in the column. The fourth head adds 0.044 tokens. The sixth adds 0.001. Each additional head costs memory and training and contributes the product of every accept rate before it, which collapses quickly.

EACH HEAD CONTRIBUTES THE PRODUCT OF ALL BEFORE IT base 1.000 head 1 0.720 head 2 0.396 head 3 0.158 head 4 0.044 head 5 0.008 head 6 0.001 four heads: 2.32 tokens per forward pass a fifth adds 0.008 and costs memory, training and a deeper tree which is why the count stops where it does
Nothing is accepted unless everything before it was, so the contributions are a running product.

Checking many guesses in one pass

Each head produces a distribution, not a single token. Taking the top two from each gives 2⁴ = 16 possible four-token continuations.

Verifying those separately would be sixteen forward passes, which defeats the purpose. Instead the candidates are arranged as a tree — shared prefixes share nodes — and the whole tree goes through the model in one pass, with an attention mask that lets each node see only its own ancestors.

distinct four-token paths16
positions in the tree~30
forward passes to verify them all1

Thirty positions in one pass, against one position per pass in ordinary generation. Since the pass was memory-bound anyway, processing thirty positions instead of one costs very little extra — and one of those sixteen paths is likely to be partly correct.

The mask is what makes this work. Without it, the candidates would attend to each other and the verification would be meaningless.

What it left behind

Medusa's specific arrangement was superseded quickly, and the ideas were not.

Proposals should come from the model already running. A second model is memory and drift; extra heads are neither.

Verification should be a tree, not a line. Checking one guessed sequence wastes the pass if the first token is wrong; checking a branching set of them does not.

Predicting several tokens ahead is a training objective, not only an inference trick. Training a model to predict more than one position turns out to help the model itself, not just the decoder attached to it — which is why the idea now shows up during pretraining rather than as an attachment afterwards.

What to take away

The speedup is one number: average tokens accepted per forward pass. Everything in the method exists to raise it, and everything about the method's limits comes from the fact that it is a product of accept rates rather than a sum.

Which is why a fourth head is worth adding, a sixth is not, and why getting each head's accept rate up matters more than adding another one.