Speculative decoding: a small draft model guesses, the big target model verifies. It works, but it has a problem. The draft model is a second model to train, align, and deploy. It's operational baggage.
Medusa (arxiv.org/abs/2401.10774, Cai et al. 2024) asks: what if the draft model was already inside the target?
The idea
The target model's last hidden state already contains information about what comes next. Medusa just adds tiny MLP heads on top of that hidden state, one per future position. Head 1 predicts token t+1, head 2 predicts t+2, and so on. Each head is a single-layer MLP, a few million parameters on a 70B model, nothing.
No separate draft model. No alignment. The heads are fine-tuned while the backbone is frozen (Medusa-1, lossless) or jointly (Medusa-2, faster).
Tree-based attention
Here's the clever part. Each head predicts its top-k tokens, and the combinations form a tree of candidate sequences. Medusa verifies all of them in a single forward pass using a specially constructed attention mask that lets each candidate attend to its own prefix.
This is what makes Medusa fast: it's not one guess per step, it's a whole tree of guesses verified at once. The more candidates, the more likely one is right, but the bigger the tree and the more compute per pass. It's a knob you tune.
Combined with a typical acceptance scheme (accept tokens whose probability clears a threshold derived from the model's entropy), Medusa hits 2.2 to 3.6x speedup. And the self-distillation extension solves the cold-start problem when you don't have domain data to train the heads on.
Speculative decoding is like hiring a junior dev to draft code and a senior dev to review it. Medusa is realizing the senior dev already knows what the junior would write, so you just add a few sticky notes to their monitor.
Medusa vs EAGLE
EAGLE drafts at the feature level rather than the token level. Medusa drafts at the token level but with multiple heads and a tree. EAGLE tends to win on acceptance rate; Medusa wins on simplicity. Both beat plain speculative decoding. The frontier is combining them.
The takeaway
Every optimization is a bet on where the waste is. Speculative decoding said the waste was idle compute. Medusa said the waste was the draft model itself. Both were right, one was simpler.
Tomorrow: MoE routing, where only part of the model wakes up.