Inference Optimisation advanced 9 min read 4 flashcards

Speculative Decoding

Use a small draft model to propose tokens that a large verifier accepts or rejects in parallel, giving lossless 2-3x latency wins on autoregressive generation.

Decoder inference is sequential by definition: token t+1 depends on t. That makes a single decode step memory-bandwidth-bound (you read the whole model and KV cache to emit one token), so the GPU sits idle on tensor cores most of the time. Speculative decoding turns that idle compute into a speed win by letting a small cheap "draft" model guess several tokens ahead, then having the big "verifier" model check all of them in a single forward pass. Done right, the output distribution is identical to the verifier running alone.

The protocol

  1. A draft model q autoregressively proposes k candidate tokens (typically k = 4-8).
  2. The verifier model p runs a single forward pass over the prefix plus those k tokens, producing p-distributions at each position in parallel.
  3. For each proposed token x_i, accept it with probability min(1, p(x_i)/q(x_i)). On rejection, sample a corrected token from the adjusted distribution (p - q)_+ (Leviathan et al. 2023).
  4. Append the accepted tokens (plus the correction) and loop.

Because step 3 follows the importance-sampling rule, the marginal distribution of accepted sequences is exactly p - no quality loss, even though you ran the cheap model most of the time.

Why it is a win

The verifier processes k+1 positions in one forward pass at roughly the same cost as one position (decode is bandwidth-bound, and you read the weights and cache once either way). If the average acceptance length is L, your speedup over plain decoding is L / (1 + L * c) where c is the cost ratio of draft to verifier. Typical numbers:

Setup Draft Verifier Acceptance Speedup
Leviathan 2023 (T5-XXL) T5-small T5-XXL ~4.5 tokens 2-3x
Chen 2023 (Chinchilla 70B) 4B Chinchilla 70B Chinchilla ~3 tokens 2.0-2.5x
Medusa-2 on Vicuna-7B n/a (heads) Vicuna-7B ~2.5 tokens 2.2-3.6x

The acceptance rate, not the draft model's speed, is the dominant factor. A faster draft that disagrees more often is a worse choice than a slower one that agrees more.

Medusa: single-model variant

Training a separate draft model is awkward in production: you need to keep two models in sync across fine-tunes, and serving infrastructure must hold both. Medusa (Cai et al., 2024) sidesteps this by adding k extra prediction heads to the verifier itself. Each head predicts the token at position t + i + 1 directly from the hidden state at position t. You then verify the candidate set using tree attention in a single forward pass of the base model.

Medusa-1 trains only the heads; Medusa-2 jointly fine-tunes the backbone. The trade-off is straightforward: you get speculation without a second model, at the cost of slightly more parameters and a quality dip if the heads are under-trained.

EAGLE and lookahead

EAGLE (Li et al., 2024) goes further: speculate on the feature (penultimate-layer hidden state) rather than the token. A tiny one-layer transformer autoregresses over features, and the verifier's LM head decodes them. Acceptance lengths rise to 4-5 tokens; reported 2.7-3.5x speedups on LLaMA2-Chat 70B.

Lookahead decoding (LMSYS, 2023) avoids speculation entirely. It uses the Jacobi iteration trick: predict k tokens in parallel with the verifier itself, refine them across iterations, and harvest any n-grams that stabilise as accepted prefixes. No draft, no extra parameters, lossless. Wins are smaller (1.5-2x) but it works with any model out of the box.

The speedup ceiling

The theoretical ceiling is the average acceptance length, and the practical ceiling sits around 2-3x for most production workloads:

  • Easy domains (code completion, structured output, repetition-heavy chat) hit 3-4x because the next few tokens are highly predictable.
  • Hard domains (open-ended reasoning, novel content, code with unusual identifiers) drag acceptance below 2 tokens; speedup drops to 1.3-1.5x.
  • Greedy decoding (temperature = 0) accepts more aggressively than sampling, so deterministic generation usually sees the higher numbers reported in papers.

When it falls down

  • Small batch sizes only. Speculative decoding helps when decode is bandwidth-bound. At large batch sizes you are already saturating compute, and the verifier's "free" parallel verification is no longer free.
  • Distribution mismatch. A draft model fine-tuned on a different distribution from the verifier sees acceptance collapse. Re-train the draft after every verifier fine-tune.
  • Long-tail latency. The variance of acceptance length means p99 latency can be worse than plain decode even when mean latency is much better. If your SLO is on p99, measure carefully.

Further reading

Check yourself

4 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track