Fast Inference from Transformers via Speculative Decoding

In plain words

A large language model writes text one token at a time, and each token needs a full pass through the model. On modern accelerators that pass is limited by how fast the weights can be read from memory, not by arithmetic, so most of the chip's compute sits idle while it waits.

Speculative decoding puts that idle compute to work. A much smaller model quickly guesses the next few tokens. The large model then checks all of the guesses at once, which costs about as much as producing a single token. Guesses the large model agrees with are kept; at the first disagreement the large model supplies its own token and the process repeats.

The clever part is the acceptance rule. It is designed so that the final text is distributed exactly as if the large model had written it alone. The small model only decides how fast you go, never what you get. In the paper, a T5-small drafter made T5-XXL 2.3 to 3.4 times faster, with output drawn from exactly the same distribution.

TL;DR

Autoregressive decoding needs one serial target-model call per token and is usually memory-bandwidth bound, so spare compute is available. Speculative decoding runs a cheap approximation model M_q for \gamma steps, then evaluates the target M_p on all \gamma+1 prefixes in one parallel call, accepts each draft token with probability \min(1, p(x)/q(x)), and on the first rejection resamples from norm(\max(0, p - q)). The output distribution is provably identical to sampling from M_p alone, for any M_q. Off-the-shelf T5-small drafting for T5-XXL gives 2.3X-3.4X walltime speedup at batch size 1 (author-reported). Relevant if you serve a large model at low batch size and have a smaller model with the same tokenizer.

Problem

Decoding K tokens takes K serial runs of a large model §1. Prior speedups (distillation, sparsification, quantization, architecture changes, adaptive computation / early exit) typically require architecture or training changes and do not preserve identical outputs §1 §5. The authors observe that large-model inference is usually limited by memory bandwidth and communication rather than arithmetic, so extra concurrent computation can be cheap §1. Prior speculative approaches (Blockwise Parallel Decoding, Shallow Aggressive Decoding) support only greedy decoding and need custom training or only copy the input §5. Assumption throughout: enough compute to run \gamma+1 evaluations of M_p in parallel without increasing walltime §3.3.

Method

Standardized sampling. Any sampling scheme (argmax, top-k, nucleus, temperature) is first cast as standard sampling from an adjusted distribution; e.g. argmax = zero out non-max entries and renormalize. Below, p(x) and q(x) always denote the distributions of M_p and M_q after this adjustment §2.2.

Speculative sampling (one token). To sample x \sim p: draw x \sim q; keep it if q(x) \le p(x); otherwise reject with probability 1 - p(x)/q(x) and on rejection resample from p'(x) = norm(\max(0, p(x) - q(x))) §2.3. Proof that the result is distributed exactly as p for any p, q is in app. A.1; the normalizer of p' equals 1-\beta app. A.1.

Speculative decoding step alg. 1 §2.1:

Algorithm: SpeculativeDecodingStep  [alg. 1]
Input: target M_p, approximation M_q, prefix
1. for i = 1..gamma:                       # autoregressive drafting with M_q
     q_i(x) <- M_q(prefix + [x_1..x_{i-1}])
     x_i ~ q_i(x)
2. p_1(x), ..., p_{gamma+1}(x) <- M_p(prefix), ..., M_p(prefix + [x_1..x_gamma])   # ONE parallel call
3. r_1..r_gamma ~ U(0,1)
   n <- min({ i-1 | 1 <= i <= gamma, r_i > p_i(x_i)/q_i(x_i) } U { gamma })   # accepted count
4. p'(x) <- p_{n+1}(x)
   if n < gamma: p'(x) <- norm(max(0, p_{n+1}(x) - q_{n+1}(x)))
5. t ~ p'(x)
Output: prefix + [x_1..x_n, t]             # between 1 and gamma+1 new tokens

If all \gamma drafts are accepted, the bonus token t is sampled from the unmodified p_{\gamma+1} alg. 1. Each iteration yields at least one token, so the number of serial M_p calls never exceeds standard decoding §2.1.

Analysis. Acceptance rate \beta = \sum_x \min(p(x), q(x)) = 1 - D_{LK}(p,q) where D_{LK}(p,q) = \sum_x |p(x)-q(x)|/2 (Lemma 3.3, Theorem 3.5) §3.2; \alpha = E(\beta) = E(\min(p,q)) (Corollary 3.6) §3.2. Assuming i.i.d. \betas, tokens per iteration are a capped geometric variable §3.1:

E(\#\text{generated tokens}) = \frac{1 - \alpha^{\gamma+1}}{1 - \alpha} eq. 1

With cost coefficient c = time of one M_q run / time of one M_p run (Definition 3.7), the expected walltime improvement (Theorem 3.8) is §3.3:

\frac{1 - \alpha^{\gamma+1}}{(1-\alpha)(\gamma c + 1)}

If \alpha > c some \gamma gives an improvement of at least (1+\alpha)/(1+c) (Corollary 3.9) §3.3. With \hat c = ratio of arithmetic ops per token of M_q to M_p, the total-ops increase factor (Theorem 3.11) is \frac{(1-\alpha)(\gamma\hat c + \gamma + 1)}{1 - \alpha^{\gamma+1}} §3.4. Memory reads of target weights and KV cache shrink by the factor of eq. 1, since they are read once per iteration §3.4. The optimal \gamma maximizes the walltime formula and is found numerically over integers §3.5 fig. 3.

Approximation model choices §3.6: any M_q preserves correctness. Off-the-shelf same-architecture models about two orders of magnitude smaller than M_p usually performed best. Negligible-cost models (c \approx 0; unigram/bigram tables, copying from context) bound the speedup by 1/(1-\alpha). Non-autoregressive drafters can replace the drafting loop with one call.

Extensions (appendix, not used in main results): beam search with M_q beam width u \ge w, accepting while top_w(M_p) \subseteq top_u(M_q) app. A.4; a lenience parameter l \in [0,1] that multiplies q(x) by l before comparison, which breaks exactness but guarantees no token is sampled with probability above p(x)/l app. A.5.

Hyperparameters

Results

Author-reported. Setup: T5-XXL 11B target, batch size 1, single TPU-v4, speedup in walltime vs the T5X implementation §4.1.

Theoretical predictions from Theorem 3.8 mostly match measurements (e.g. EnDe T5-small temp 0: expected 3.2X, measured 3.4X; EnDe T5-large temp 0: expected 2.5X, measured 1.7X) tab. 4. Idealized trade-off with c=\hat c=0: \alpha=0.8, \gamma=5 gives 3.69X speed at 1.63X total operations tab. 1. Outputs are identical to the baseline (exactness guarantee) §4.1 app. A.1; no quality metrics are reported because none change.

Paper vs code

No official code. The T5X-based implementation used for walltime measurements is not linked in the paper.

Pitfalls

• Standardize both distributions identically before the accept test. Temperature, top-k, nucleus or argmax must be applied to both p and q as explicit probability distributions; the ratio p_i(x_i)/q_i(x_i) and the residual \max(0, p-q) are computed on the adjusted distributions §2.2 §6. LaMDA results also had a top-40 filter on outputs tab. 3.
• Bonus token. When all \gamma drafts are accepted, sample t from the unmodified p_{\gamma+1}, not from a residual; when n<\gamma, sample from the residual at position n+1 and discard everything after alg. 1.
• Residual normalization can be numerically tiny. norm(\max(0,p-q)) has normalizer 1-\beta app. A.1; when q \approx p this is near zero. Not discussed in the paper; guard the division. With argmax standardization, p and q are one-hot, so acceptance is just "draft argmax equals target argmax" (follows from §2.2 and the accept rule).
• Theory assumes i.i.d. \beta and free parallelism. eq. 1 and Theorem 3.8 assume i.i.d. acceptance and that \gamma+1 parallel M_p evaluations cost the same walltime as one §3.1 §3.3; the paper attributes the gaps in tab. 4 (up to 2.5X predicted vs 1.7X measured) to this assumption and implementation differences app. A.3. Theorem 3.8 also assumes long generations §3.3.
• Hidden compute cost. Total arithmetic operations rise (e.g. 1.53X-1.63X in tab. 1); the method does not help when spare compute is unavailable, e.g. large batches §3.4 §6. All walltime results are at batch size 1 §4.1.
• c is hardware-dependent, unlike \alpha §3.3; re-measure it and re-pick \gamma on your stack. Bigger drafters raise \alpha but lose on c (T5-large is the slowest drafter at every setting) tab. 2.
• Tokenizer/vocabulary must match so p and q are over the same support; the paper only tested same-architecture drafters with the same standardization §3.6 and does not discuss mismatched vocabularies.
• KV caches. Rejected draft positions leave stale entries in both models' caches that must be rolled back; cache management is not described in the paper (only that target weights and KV cache are read once per iteration §3.4).
• Lenience breaks exactness. Results with l<1 app. A.5 tab. 5 are not distribution-preserving; all other results use strict l=1 app. A.5. For temp 0, lenience is applied differently (accept if p(x) \le l \cdot \max(p) as written in app. A.5; the inequality direction as printed looks inverted relative to the intended "accept near-max tokens", so verify before implementing).
• Beam search variant app. A.4 is only sketched, with no analysis or experiments.

Lineage

• Builds on: speculative execution / branch prediction from processors §1; Blockwise Parallel Decoding (Stern et al., 2018), which is greedy-only and needs custom training §5; Shallow Aggressive Decoding (arXiv 2106.04970), copy-only and greedy-only §5.
• Compared against: standard T5X autoregressive decoding of T5-XXL §4.1.
• Concurrent / independent: Accelerating LLM Decoding with Speculative Sampling, reported 2X-2.5X on Chinchilla 70B with the same idea §5.
• Follow-ups: Medusa and EAGLE change the drafter (cards compiled separately).