Accelerating Large Language Model Decoding with Speculative Sampling

In plain words

Imagine a senior writer who is accurate but slow, paired with a fast junior who drafts a few words ahead. The senior reads the draft in one glance, keeps the words they would have written anyway, and fixes the first one they disagree with. Because reading is much cheaper than writing for the senior, the pair finishes sooner without the result changing.

That is speculative sampling, the same idea as speculative decoding by Leviathan et al., found independently at DeepMind. A small draft model proposes a handful of tokens and the large model checks them together. A probabilistic acceptance rule guarantees the output follows the large model's distribution.

The paper's practical lesson is about large deployments where one model is split across many chips. There, the draft model should be wide and shallow rather than narrow and deep, because every extra layer adds communication between chips. With that design, Chinchilla 70B produced each token about twice as fast.

TL;DR

DeepMind's version of speculative decoding, developed concurrently with and independently of Leviathan et al.. A cheap autoregressive draft model samples K tokens. The target model scores the K+1 prefixes in a single forward pass. Draft tokens are accepted left to right with probability \min(1, q/p). On the first rejection, a replacement token is drawn from the normalized residual (q-p)_+. If all K are accepted, one bonus token is drawn from the target. The output distribution equals the target's, up to hardware numerics. The paper's engineering contribution is the distributed-serving view: the draft is a wide, shallow 4B model (8 layers) served on the same 16 TPU v4 topology as Chinchilla 70B. That yields 2-2.5x lower per-token latency at batch size 1. There is no official code.

Problem

Autoregressive sampling (ArS) at small batch size is memory-bandwidth bound: every token requires streaming all parameters, and the KV cache, through the accelerators. Model size divided by total memory bandwidth therefore sets a hard ceiling on tokens per second §3. Large models also need model parallelism, which adds inter-device communication on every call §1 §3. Prior parallel-decoding methods only work with greedy decoding, bias the output, or target other modalities. None had been scaled to distributed setups §2. Assumptions: latency-critical serving at small batch size, and a target model you cannot modify §1 §7.

Section numbering note: the compiled paper uses unnumbered sections (nonumbering class option). Citations here number them in document order: §1 Introduction, §2 Related Work, §3 Auto-regressive Sampling, §4 Speculative Sampling (§4.1 Conditional Scoring, §4.2 Modified Rejection Sampling), §5 Choice of Draft Models, §6 Results (§6.1 Evaluation on XSum and HumanEval, §6.2 Acceptance rate changes per domain, §6.3 Trade off between longer drafts and more frequent scoring), §7 Conclusion. The appendix material sits under the unnumbered "Supplementary Materials" heading (Hyperparams, Proofs). All equations are unnumbered. Algorithms (1-2), tables (1-2), figure 1 and Theorem 1 are numbered.

Method

Why it works (conditional scoring) §4.1. For Megatron-sharded transformers at small batch size, sampling time is dominated by three components. Scoring K tokens costs about the same as sampling one for each of them:

1. Linear layers are memory bound for a few embeddings, so a small K adds little time.
2. Attention is dominated by KV-cache reads, and the KV-cache size does not change with K.
3. All-reduces (one after every feed-forward and attention layer under Megatron) carry few activations, so they are latency bound rather than throughput bound.

The paper warns that positional encoding, the decoding method (e.g. a sort for nucleus sampling) and hardware limits can still make scoring slower than sampling §4.1.

Modified rejection sampling §4.2. Given context x_{1..n} and draft tokens \tilde x_{n+1..n+K} \sim p, accept \tilde x_{n+1} with probability

\min\left(1, \frac{q(\tilde x_{n+1}\mid x_{1..n})}{p(\tilde x_{n+1}\mid x_{1..n})}\right) §4.2

On acceptance, continue with \tilde x_{n+2}. On rejection, resample

x_{n+1} \sim \big(q(x\mid x_{1..n}) - p(x\mid x_{1..n})\big)_+,\qquad (f(x))_+ = \frac{\max(0, f(x))}{\sum_x \max(0, f(x))} §4.2

and stop. Each loop emits at least 1 token and at most K+1 tokens §4.2. Theorem 1 proves that for a single draft token P(X=x) = \min(p(x),q(x)) + \max(0, q(x)-p(x)) = q(x) §Proofs. Applying the rule sequentially recovers the target distribution §4.2.

Algorithm: Speculative Sampling (SpS)  [alg. 2]
Input: target q, draft p, prompt x_1..x_t, lookahead K, min target length T
n <- t
while n < T:
  1. for t = 1..K: sample x~_t ~ p(. | x_1..x_n, x~_1..x~_{t-1})        (K sequential draft calls)
  2. in parallel (one target call), compute K+1 target distributions
       q(.|x_1..x_n), q(.|x_1..x_n, x~_1), ..., q(.|x_1..x_n, x~_1..x~_K)
  3. for t = 1..K:
       r ~ U[0,1]
       if r < min(1, q(x~_t | x_1..x_{n+t-1}) / p(x~_t | x_1..x_{n+t-1})):
           x_{n+t} <- x~_t ; n <- n+1
       else:
           x_{n+t} ~ (q(.|x_1..x_{n+t-1}) - p(.|x_1..x_{n+t-1}))_+ ; break
  4. if all K accepted: sample bonus token from the last target distribution q(.|x_1..x_n, x~_1..x~_K); n <- n+1
Output: x_1..x_n (n >= T)

Decoding transforms §4.2. For nucleus, top-k or temperature sampling, apply the transform to both p and q first, then run the rejection scheme on the modified distributions. The authors report that the acceptance rate is robust to the exact parameters §4.2.

Draft model choice §5 §6. Any drafter that exposes logits works. The authors reject options that change the target model or require sequence-level distillation because of their cost §5. They use a smaller model with the same tokenizer and dataset. For distributed serving the draft should be wide and shallow, so it can run on the target's topology without communication overhead dominating. A chinchilla-optimal 7B is fastest on 4 TPU v4 (5 ms/token) and gets slower on 16 §6. Few layers means few all-reduces, and the paper cites the robustness of LM quality to aspect ratio §6.

Hyperparameters

Results

Author-reported. Chinchilla 70B target, 4B draft, 16 TPU v4, batch size 1, K=4. Mean token time comes from TPU profiles plus logged tokens per loop tab. 1 §6.1.

• Benchmark scores stay at parity, which is the paper's empirical check of losslessness. Exact output equality is not expected because seeds and numerics differ between ArS and SpS §6.1.
• For HumanEval and greedy XSum, SpS throughput exceeds the theoretical memory-bandwidth ceiling for ArS §6.1.
• Varying K fig. 1 §6.3: loop time grows roughly linearly with K. Accepted tokens divided by K+1 falls as K grows. Speedup plateaus or regresses at large K, and the variance of time-to-128-tokens grows with K. The curves are given only as a figure, with no numeric table.
• Code accepts more tokens than prose. The authors attribute this to common sub-sequences, shorter tokens and temperature sharpening, but this is a hypothesis, not a measurement §6.2.

Paper vs code

No official code.

Pitfalls

• Bonus-token typo in alg. 2. The last line of alg. 2 writes q(x\mid x_1,\dots,x_n,x_{n+K}), but by then n has already been incremented K times. The intended distribution is the (K{+}1)-th target output from step 2, conditioned on the full accepted draft §4.2. Do not run another target call for it.
• Indexing in alg. 2 is sloppy. The inner loop variable t shadows the prompt length t. The acceptance ratio is written q(x\mid\cdot)/p(x\mid\cdot) but must be evaluated at the drafted token \tilde x_t, as in §4.2. If you increment n inside the loop and also index x_{n+t}, you double-count. Use a separate accepted-count variable.
• The residual distribution needs care. (q-p)_+ is renormalized over the vocabulary. If q=p numerically, rejection has probability 0 but the denominator is 0. Guard against it, for example by falling back to q. The paper does not discuss this case §4.2 §Proofs.
• Use the same transformed distributions everywhere. Temperature, top-k and nucleus filtering must be applied to both p and q before computing the ratio and the residual §4.2. Draft tokens must be sampled from the transformed p. Mixing raw and filtered distributions breaks exactness.
• Greedy mode is not specified. The paper reports greedy results tab. 1 but does not spell out the acceptance rule for greedy decoding. Applying the general rule to one-hot distributions reduces to "accept iff the draft argmax equals the target argmax; otherwise emit the target argmax". That is an inference, not stated in the paper.
• The speedup depends on the deployment. The gain relies on scoring K+1 tokens costing about one sampling step, which holds at small batch and small K §4.1. The paper only tests batch size 1. Larger batches, or implementations where scoring is slower than sampling (e.g. a sort for nucleus sampling), reduce the gain §4.1.
• The draft topology matters as much as its size. A deep small model served on its own optimal topology gives only a modest speedup and leaves the target's hardware idle during drafting §6. The paper's 4B draft is a custom model, and its training recipe beyond the architecture is not stated.
• Tail latency. Larger K raises the variance of sequence time, which matters for P90/P99 latency targets §6.3. The best K depends on the domain: K=3 for XSum nucleus fig. 1.
• KV-cache rollback is unstated. The paper says nothing on discarding rejected draft positions from the draft and target KV caches, or on running the draft over the bonus or resampled token. These steps are needed in any real implementation.

Lineage

• Builds on: blockwise parallel decoding (Stern et al. 2018) and aggressive decoding, which the authors say only work with greedy decoding or bias the output §2. Megatron-style sharding §4.1.
• Concurrent with: Fast Inference from Transformers via Speculative Decoding, the same core idea. This paper puts more emphasis on distributed serving at 70B scale §2.
• Related later work on the same concept: 2401.10774, 2401.15077. See their cards for how they differ.