Era 6 · Generative models and systems · 2023

57 Speculative Decoding

Fast Inference from Transformers via Speculative Decoding · Leviathan, Kalman, Matias · Google · ICML · and Accelerating LLM Decoding with Speculative Sampling · Chen et al. · DeepMind · 2023
🟧 read selectively~1.5 horiginal ↗
The gist in 20 seconds. A small draft model cheaply proposes K tokens in a row; the big target model checks all of them in ONE parallel pass. Each draft token is accepted with probability min(1, p/q); on the first rejection — resample from the "remainder". The output comes out exactly from the target's distribution (lossless), but with fewer sequential passes of the big model → a ×2–3 speed-up.

Context

Autoregression decodes one token at a time, and every token is a full pass of the big model. And decoding is memory-bound: the time goes on loading the weights, not on the arithmetic, so checking K tokens at once costs almost the same as checking one. That is where the slack hides.

The idea and the mechanism

Let q be a cheap draft model (small, or a distilled one) and p the target. Draft phase: q generates K candidates autoregressively. Verify phase: in a single pass p computes its own probabilities for all K positions. Acceptance goes by rejection sampling: a draft token x is accepted with probability min(1, p(x)/q(x)); at the first rejection the token is resampled from the normalized remainder max(0, p − q) and the rest of the draft is thrown away. One pass of the target moves us several tokens forward.

probability · algorithms Acceptance, and why the distribution is NOT distorted

The draft proposed a token x ∼ q. We accept it with probability

\[ \alpha(x) = \min\!\left(1,\ \frac{p(x)}{q(x)}\right) \]

If we reject it, we take a token from the residual distribution (what the target wants and the draft under-delivered):

\[ x \sim \mathrm{norm}\big(\max(0,\ p(\cdot) - q(\cdot))\big) \]

The trick is that the total probability of emitting x matches p(x) exactly: "accepted from q" plus "resampled from the remainder" add up to the target distribution —

\[ \Pr[\text{return } x] \;=\; q(x)\,\alpha(x) \;+\; (\text{rejection mass})\cdot p_{\text{res}}(x) \;=\; p(x) \]

That is why speculative decoding is lossless: the result is indistinguishable from honest sampling out of p. The speed-up comes from accepting several tokens, on average, per expensive target pass; the closer q is to p, the higher the acceptance rate.

Python One speculative decoding loop 
def spec_step(draft, target, ctx, K):
    xs, qs = [], []
    for _ in range(K):                       # the draft proposes K tokens
        q = draft.probs(ctx + xs); x = sample(q)
        xs.append(x); qs.append(q[x])
    P = target.probs_parallel(ctx, xs)        # the target checks all K in one pass
    out = []
    for j, x in enumerate(xs):
        if random() < min(1, P[j][x] / qs[j]):
            out.append(x)                     # accepted
        else:
            out.append(sample(norm(relu(P[j] - draft_dist[j])))); break  # resample from the remainder, stop
    return out                                 # distribution = exactly the target
draft qcheap K candidates target p1 parallel pass ✓ accepted (3)✗ → resample, stop several tokens per target pass, and the same distribution
The draft proposes K tokens and the target checks them in one pass: the correct prefix is accepted, and the first position the target dislikes is rewritten from the residual distribution. Faster, but the quality is the target model's.
Analogy. A junior (the draft) quickly sketches the next few words. The boss (the target) skims the whole sketch at once: wherever it matches what he would have written himself, he accepts it; the first place it diverges he fixes himself, and stops reading there. The final text reads as if the boss wrote it, but it took less time, because checking a batch is faster than writing word by word.

Why it matters

A ×2–3 inference speed-up without any loss of quality is a rare combination; it became a standard serving trick (vLLM and descendants: Medusa, EAGLE, self-speculation). The deeper lesson: since decoding is memory-bound, parallel verification is nearly free — and you can trade it against sequential generation.

Connections

← speeds up32. Transformer

It attacks exactly the autoregressive nature of #32 — "one pass per token". Without changing the model or its distribution, speculative decoding cuts the number of expensive sequential passes of the big model.

← builds on28. Knowledge Distillation

The gain grows the closer the draft is to the target in distribution — which is why the draft is often produced by distilling the target model (#28). A good "small twin" = a high acceptance rate = a large speed-up.

↔ a different axis of speedup48. FlashAttention

Both speed up inference, but differently: FlashAttention makes one attention computation cheaper, speculative decoding cuts the number of model passes. They are orthogonal and stack in production.

Questions worth asking

Why does acceptance preserve the target distribution EXACTLY rather than approximately?

It is an exact consequence of rejection sampling: acceptance with probability min(1, p/q) plus a resample from the remainder max(0, p−q) together give every token an emission probability = p(x). Mathematically it is a coupling that reproduces the target distribution even though the proposals come from q. So the output is statistically indistinguishable from honest sampling out of the target — not "almost", but exactly.

If the target computes all K tokens anyway — where does the speed-up come from?

Because decoding is memory-bound: what dominates is loading the model weights, not the arithmetic. One target pass over K positions loads the weights once and costs almost as much as a pass over a single one. Accepting several tokens on average per such pass cuts the number of sequential expensive steps — hence the win in wall-clock time.

What if the draft model is bad?

Then the acceptance rate is low: the target rejects proposals often and the gain melts away (in the limit — ordinary generation plus the overhead of running the draft). The key is a draft close to the target in distribution. Hence the various recipes: distillation, self-speculation (the same model predicting ahead), Medusa heads, EAGLE — all of them about getting proposals accepted more often.

What to read in the original

Read the essentials — the acceptance/resampling algorithm itself and the proof that the distribution is preserved (that is the heart of the work). Leviathan et al. (arXiv 2211.17192) give the formal statement; Chen et al. (2302.01318) give the equivalent "speculative sampling" formulation and results on large models.