Lesson 13.10 · 23 min · video lesson
Medusa: Parallel Decoding via Multiple Prediction Heads
What if a language model could guess its next three words by itself, with no second model to help, and still never change what it would have said?
In short: Medusa speeds up text generation by bolting a few small extra "heads" onto an existing LLM so that, in one forward pass, the model guesses several future tokens at once. A special tree-shaped attention mask lets the model check many of those guesses in the very next pass, and every guess that matches what the model would have produced anyway is kept for free.
What you will learn
- What is Medusa
- Why text generation is slow
- A quick recap of speculative decoding
- The problem with needing a draft model
- The big idea: many heads on one model
- How tree attention checks many guesses at once
- The math behind the speedup with small numbers
- The results
- How Medusa lives on today
- Quick Summary
In this lesson
- What is Medusa?
- Why text generation is slow
- A quick recap of speculative decoding
- The problem with needing a draft model
- The big idea: many heads on one model
- How tree attention checks many guesses at once
- The math behind the speedup with small numbers
- The results
- How Medusa lives on today
- Worked example, step by step
- Practice: try it yourself
- Quick summary
Key terms in this lesson
- Medusa head
- A small feed-forward layer on top of an LLM's last hidden state that predicts a token several positions ahead.
- Speculative decoding
- Guessing several future tokens cheaply, then verifying them all with the main model in one forward pass.
- Tree attention
- An attention mask that lets many candidate continuations share one forward pass, each node seeing only its ancestors.
- Memory-bandwidth bound
- Limited by how fast data (here, model weights) can be read from memory rather than by arithmetic speed.
- Typical acceptance
- Medusa's relaxed rule that accepts a guess when the model gives it a high enough probability, trading exactness for speed.
- Self-distillation
- Training the heads on text the model itself generated, used when the original training data is unavailable.
This lesson is part of the Generative AI Engineering and Complete AI Engineer tracks. See plans and pricing to open the full lesson, its video, labs and quiz, or try the free lesson first (free account needed).
← 13.9 N-gram Speculation: Draft Tokens Without a Draft Model · 13.11 EAGLE: Feature-Level Drafting for Faster Inference →