Hello, I'm Shrijith Venkatramana, and I'm building LiveReview — a blast-radius aware AI code review built for your business-critical systems.
Star us to help devs discover the project, give it a try, and share your feedback to help improve the product.
A modern LLM can contain hundreds of billions of parameters, run on extremely expensive accelerators, and still spend most of its inference time doing something that looks embarrassingly sequential: That is the awkward part of autoregressive generation.
The model may process a whole prompt in parallel during the initial prefill, but once generation starts, the next token depends on the previous token.
So generating 100 tokens looks conceptually like running the model 100 times.
And for many serving workloads, that is exactly where the money goes.
A family of techniques tries to break this bottleneck by asking a deceptively simple question: What if the model could predict several future tokens at once, then verify them in parallel?
That idea leads to speculative decoding, Medusa-style multiple decoding heads, and the broader multi-token prediction approach used during training.
The interesting part is that these are not merely "optimization tricks." They change the computational structure of decoding.
This article develops that idea from first principles and then gets into the engineering details.
1.
The problem: your GPU is doing an expensive sequential loop Consider ordinary autoregressive decoding.
Given a prompt: the model predicts: Then it feeds the new sequence back through the model: and predicts the next token.
Then again: and so on.
Formally, the model factorizes the probability of a sequence as: That conditional dependence is what makes language modeling so useful.
It is also what makes decoding annoying.
Prefill is parallel; decode is sequential This distinction matters enormously in production.
Suppose a prompt has 2,000 tokens and we want 200 generated tokens.
During prefill, the transformer can process many positions concurrently: During decoding: The transformer itself is highly parallelizable.
The dependency graph of generation is not.
This creates an unusual hardware situation.
A giant model can spend much of its decoding time limited not by arithmetic throughput but by repeatedly moving a large set of model weights through memory.
This was one of the motivations behind the speculative decoding work by Yaniv Leviathan, Matan Kalman, and Yossi Matias, and later became a central observation in Medusa.
The core question is therefore not: How do we make one forward pass cheaper?
It is: How do we get more than one accepted token out of each expensive forward pass?
That is a much more interesting question.
2.
The first big idea: speculate, then verify Imagine that instead of asking the large model for one token, we had a tiny model that could cheaply guess several: The big model can then evaluate those candidate tokens in parallel.
Suppose the small model proposes: The large model might agree with all four: Great.
One expensive evaluation effectively produced four tokens.
But perhaps the draft is: and the large model says: Then we keep the accepted prefix and let the large model continue from the rejected position.
This is the basic idea of speculative decoding.
It is beautifully simple: Leviathan et al. showed that this could accelerate generation while preserving the output distribution exactly, rather than merely producing "approximately similar" text.
Their 2023 paper reported roughly 2x-3x acceleration on the models they evaluated.
The crucial insight was that autoregressive generation does not mean the expensive model must discover every token sequentially.
It only means that the final accepted sequence has to respect the autoregressive distribution.
Speculation lets us take advantage of the fact that many consecutive tokens are easy to predict.
For example: contains many stretches where the answer is nearly obvious: or or The large model is still being asked to perform the full computation, but we are trying to amortize that expensive computation across several tokens.
This idea naturally leads to the next question: Why maintain a second model just to make guesses?
That is where Medusa becomes interesting.
3.
Medusa: give the model several extra heads Medusa, introduced by Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason Lee, Deming Chen, and Tri Dao, takes a different route.
Instead of: Medusa adds several lightweight decoding heads to the existing model.
Conceptually: The ordinary language-model head predicts: A Medusa head can try to predict: another: and another: The heads are cheap compared with running the entire transformer again.
So instead of doing this: we try to do: The important subtlety is that these predictions are not independent in the final decoding procedure.
Medusa uses a tree of candidate continuations and then asks the full model to verify the candidates together.
This is why "multi-token prediction" can sound simpler than it actually is.
The prediction is parallel.
The verification structure is the clever part.
4.
Why a tree instead of one straight-line guess?
Suppose Medusa predicts three future positions, and we keep the top two candidates at each position.
Naively, we could have: But not every combination is meaningful.
The structure is really: This is a candidate tree.
Why?
Because the second token depends on what happened at the first token.
If the first candidate is , then predictions for later tokens are conditional on .
If the first candidate is , the continuation is different.
A tree therefore lets the system represent multiple possible future sequences without running the full model separately for every path.
This is the key engineering trick.
The transformer can process the tree-shaped set of candidate continuations with a specially constructed attention pattern.
Conceptually: The model verifies many of these positions in one batched computation.
This converts part of the problem from: into: That is exactly the kind of workload modern GPUs are good at.
A useful mental model Think of ordinary decoding as exploring one path through a tree: At every step you pay for another expensive model evaluation.
Medusa says: Spend one expensive evaluation exploring a small local subtree, then keep the path that survives verification.
That is the entire game.
5.
The math: speed comes from accepted tokens per step The cleanest way to reason about these methods is with one quantity: Ordinary decoding has approximately: because one model evaluation produces one token.
Suppose a speculative or Medusa-style method gets: accepted tokens per verification step.
Then generating 100 tokens requires roughly: large-model evaluations instead of: That gives a first-order speedup of: But this is only the idealized calculation.
There is extra work: So a more realistic model is: and therefore: This equation is worth remembering.
It explains why a method that achieves does not necessarily deliver a literal 3x wall-clock improvement.
For example, suppose: Then: The accelerator is still doing extra work.
It is just doing substantially less serial expensive work.
Acceptance probability matters Suppose we try to predict four future tokens, and the probability each prediction is accepted is roughly: A crude approximation for the probability of getting all four accepted is: So only about 41% of branches would survive all four positions.
But we do not actually need all four to succeed.
Getting: is still useful.
The expected number of consecutive accepted tokens is approximately related to: for a K-token proposal horizon.
With: that gives: So even though all-four acceptance happens only about 41% of the time, we can still average roughly 2.36 accepted positions before considering the stop.
This is why improving head quality can be extremely valuable.
A relatively small increase in acceptance probability compounds across the sequence.
6.
Multi-token prediction is also a training idea