Vol. XXIV · Making Inference Cheap Plate 4 of 10speculative decoding, draft model, verification, acceptance rate, rejection sampling, latencyAE

Speculative Decoding.

Guess cheaply, check in parallel, keep what the large model agrees with.

A decode step costs one read of the weights whether it processes one token or several (plate XXIII.1). Speculative decoding (Leviathan et al., 2023; Chen et al., 2023) exploits that idle arithmetic: a small, fast draft model proposes γ tokens; the target model runs once over all of them, computing its own distribution at each position; each draft token x is accepted with probability min(1, p(x)/q(x)), where p and q are the target’s and draft’s probabilities, and at the first rejection a replacement is sampled from the normalised max(0, p − q). If every draft token is accepted, the target’s distribution at the end gives one more token free.

The rule guarantees that the output is distributed exactly as the target’s own samples — speculation changes the speed, never the text’s distribution. Its gain depends on the acceptance rate α and the draft’s cost. The course model micro (1.8 million parameters, four layers; plate I.9) is quantised, pruned and sped up in the page, and its loss is measured on held-out course text it never trained on. Here the draft is micro itself quantised to 4 or 3 bits (all matrices, groups of 32), costing a quarter or three-sixteenths of a target pass to read; the workshop decodes 40 tokens after “The attention”.

Below: tokens produced per target pass as the number of drafted tokens grows, measured, beside the formula using the measured acceptance rate; readouts give the acceptance rate and the estimated speed-up.

What you will learn
Workshop
Target: micro. Draft: micro with every matrix, embedding included, quantised by round-to-nearest in groups of 32. Prompt “The attention”, 40 tokens, seed 1. Acceptance rate: accepted draft tokens over draft tokens checked (those after a rejection are discarded unchecked). Speed-up: tokens per target pass divided by 1 + γc, where c is the draft’s bytes over the target’s; arithmetic assumed free, as in memory-bound decode.

Exact samples, fewer passes

With the 4-bit draft at temperature 1 and γ = 4, 72% of the checked draft tokens are accepted and each target pass yields 3.00 tokens; charging the draft 0.25 of a pass per token, the estimated speed-up is 1.50. Longer drafts yield more tokens per pass but cost more drafting and waste more on rejections, so the speed-up peaks at a small γ. The 3-bit draft agrees with micro far less (α = 0.40) and speculation with it is slower than none (0.94). At temperature 0 the output is identical to micro’s own greedy text — checked: 1 means yes.

Production systems pair a large target with a much smaller model of the same family — Leviathan and colleagues reported two to three times faster T5-XXL decoding with unchanged outputs — and the same arithmetic governs them: a draft must be both cheap and agreeable. Speculation helps most when the GPU has spare arithmetic, at small batches; at large batches decode is already closer to the compute limit (plate XXIII.5) and the extra verification work competes with other users’ tokens.

The guarantee is exact, not approximate: the accept-or-resample rule is a form of rejection sampling whose output has precisely the target’s distribution, whatever the draft.

Terms

Draft model
A cheap model that proposes tokens.
Verification
The target scoring all draft tokens in one pass.
Acceptance rate
The share of checked draft tokens the target keeps.
Rejection sampling
Accept-or-resample rule that preserves the target’s distribution.
Speculative decoding spends idle arithmetic to skip passes, with outputs exactly the target’s; it pays when the draft is cheap and usually agrees.

Sources

  1. Leviathan, Kalman and Matias, “Fast Inference from Transformers via Speculative Decoding”, ICML 2023.
  2. Chen et al., “Accelerating Large Language Model Decoding with Speculative Sampling”, arXiv:2302.01318, 2023.