Accelerating Large Language Model Decoding with Speculative Sampling

The cost of one token is the cost of reading the weights

Generating a single token from a large model means streaming every weight from HBM into the compute units, multiplying once, and throwing the loaded weights away. The matmul for one token is trivial; the weight traffic is not. The arithmetic intensity of a single-token decode step is roughly

$$\frac{2N\ \text{FLOPs}}{2N\ \text{bytes}} \approx 1\ \text{FLOP/byte},$$

for a model with $N$ parameters in fp16 — each weight contributes two FLOPs (a multiply and an add) and costs two bytes to load. One FLOP per byte sits far below an A100's balance point, so for any batch size under roughly 64 the GPU is memory-bandwidth-bound: it spends its time moving weights, not computing. The compute units are mostly idle, and that idleness is the opportunity.

One pass can score many tokens

Here is the asymmetry the method exploits: a transformer can score $K$ candidate tokens in a single forward pass for almost the same cost as scoring one, because the weights are loaded once either way and the extra work is only the small matmul over $K$ positions. So instead of paying $K$ weight-loading passes to produce $K$ tokens, you guess the tokens cheaply and pay one expensive pass to check the whole guess at once. A small, cheap draft model proposes a run of tokens; the large target model verifies them in parallel.

The algorithm

The draft model generates $K$ tokens autoregressively, $\tilde{x}_1, \tilde{x}_2, \ldots, \tilde{x}_K$, each step cheap because the draft model is small. The target model then runs one forward pass over the concatenation $[\,\text{context}, \tilde{x}_1, \ldots, \tilde{x}_K\,]$, which yields its own distribution $p_{\text{target}}(\cdot \mid \cdot)$ at every one of those positions simultaneously. Now walk the draft tokens left to right and accept each $\tilde{x}_i$ with probability

$$\min\!\left(1,\ \frac{p_{\text{target}}(\tilde{x}_i)}{p_{\text{draft}}(\tilde{x}_i)}\right).$$

A token the draft over-confidently preferred relative to the target is accepted only fractionally. At the first rejection you stop, resample that position from the adjusted residual distribution $\max(0,\, p_{\text{target}} - p_{\text{draft}})$ renormalized, and discard the remaining drafted tokens. Every accepted prefix plus the one resampled token is the output of this round.

Why the output distribution is exactly the target's

The acceptance-then-resample rule is ordinary rejection sampling, arranged so that the probability of emitting any particular token equals the target model's probability of that token. For an accepted token, the chance it was drafted and survived is $p_{\text{draft}}(x) \cdot \min(1, p_{\text{target}}(x)/p_{\text{draft}}(x)) = \min(p_{\text{draft}}(x), p_{\text{target}}(x))$; the resampling branch contributes exactly the missing mass $\max(0, p_{\text{target}}(x) - p_{\text{draft}}(x))$, and the two sum to $p_{\text{target}}(x)$. The draft model never appears in the result distribution — only in how fast you reach it. Speculative decoding is lossless: the samples are drawn from precisely the distribution the target model would have produced on its own.

How much faster

Let $\alpha$ be the per-token acceptance probability and suppose the draft model is cheap enough that its cost is negligible against a target pass. The expected number of tokens accepted per target forward pass is

$$\mathbb{E}[\text{tokens per pass}] = \frac{1 - \alpha^{K+1}}{1 - \alpha},$$

the mean run length before the first rejection. At $\alpha = 0.8$ and $K = 4$ this is about 2.8 tokens emitted for every target pass — close to a 2.8× speedup when the draft is genuinely cheap, less once you charge for the draft model's own time.

Where it pays off

The lever is the acceptance rate $\alpha$, which is just how often the draft and target agree. It runs high on long, predictable continuations and at low sampling temperature, where the target distribution is peaked and easy to guess, and it climbs further when the draft is a smaller member of the same model family that has learned the same idioms. It collapses at high temperature, where the target is deliberately diffuse and any single guess is unlikely, and whenever the draft is poorly aligned with the target — then most proposals are rejected, you fall back to one token per expensive pass, and you have paid for the draft model for nothing.