Fast Inference from Transformers via Speculative Decoding

Each step of a large model is slow, and steps run serially: decoding \(K\) tokens takes \(K\) runs of the model. Yet inference is often bottlenecked on memory bandwidth and communication rather than arithmetic, so extra compute is frequently available.

(1) Hard language-modelling tasks contain easier steps that a much smaller model approximates well; (2) with speculative execution and a new sampling method, the large model can check several of the small model’s guesses in parallel, producing several tokens per run without changing the output distribution.

A processor optimization (e.g. branch prediction) that performs a task in parallel with verifying whether it is needed. Leviathan et al. extend it to the stochastic setting, where a task is needed only with some probability, and use samples from a small model as speculative prefixes for the large one.

Those change the architecture or training, need new or retrained models, and usually change the outputs. Speculative decoding works with off-the-shelf models, needs no retraining, and guarantees exactly the same output distribution as the target model.

(1) The approximation model \(M_q\) generates \(\gamma\) guesses autoregressively; (2) the target model \(M_p\) scores all of them in one parallel run; (3) accept the longest prefix of guesses allowed by speculative sampling, then sample one more token from an adjusted distribution to replace the first rejected guess (or to extend the sequence if all were accepted).

Every parallel run of \(M_p\) yields at least one new token (the correction or the extra sample), and up to \(\gamma + 1\) tokens if \(M_q\) guesses well.

Each is cast as plain sampling from an adjusted distribution (e.g. argmax = zero out all but the maximum and renormalize). \(p(x)\) and \(q(x)\) denote the two models’ distributions after the same adjustment, so one sampling procedure covers all methods.

Sample \(x \sim q(x)\). Keep it if \(q(x) \le p(x)\); if \(q(x) > p(x)\), reject it with probability \(1 - p(x)/q(x)\) and instead sample from \(p'(x) = \mathrm{norm}(\max(0, p(x) - q(x)))\).

A token \(x'\) is guessed and accepted with probability \(q(x')\min(1, p(x')/q(x')) = \min(p(x'), q(x'))\). Rejection happens with probability \(1 - \beta\), and the residual \(p'\) gives \(x'\) with probability \((p(x') - \min(p(x'), q(x')))/(1-\beta)\); the two terms sum to \(p(x')\).

A non-iterative rejection sampler would accept \(x\) with probability \(p(x)/(M q(x))\), where \(M = \max_x p(x)/q(x)\), and otherwise fall back to sampling \(p\) directly; its acceptance rate is much lower than speculative sampling’s \(\sum_x \min(p(x), q(x))\).

\(\beta\) is the probability that a guess \(x \sim q\) is accepted, given the prefix: \(\beta = \sum_x \min(p(x), q(x)) = 1 - D_{LK}(p, q)\), where \(D_{LK}\) is a symmetric divergence in \([0,1]\) (0 iff \(p = q\), 1 iff disjoint supports). \(\alpha = E(\beta)\) measures how well \(M_q\) approximates \(M_p\).

Assuming i.i.d. \(\beta\)s, it is a capped geometric variable: \(E(\#\text{tokens}) = \frac{1 - \alpha^{\gamma+1}}{1 - \alpha}\). For \(\alpha = 0.8\) and \(\gamma = 5\) this is about 3.7 tokens per run of \(M_p\).

With cost coefficient \(c\) (time of one \(M_q\) run over one \(M_p\) run), the expected speedup is \(\frac{1 - \alpha^{\gamma+1}}{(1 - \alpha)(\gamma c + 1)}\), assuming the \(\gamma + 1\) parallel runs of \(M_p\) don’t increase walltime. Some \(\gamma\) gives a speedup whenever \(\alpha > c\).

More arithmetic: rejected guesses waste the parallel computation of \(M_p\), so total operations grow (e.g. 1.6× for a 3.7× speedup at \(\alpha = 0.8, \gamma = 5\)). Memory accesses for the target model’s weights and KV cache, however, go down, since they are read once per step rather than once per token.

When no spare compute is available: it trades more arithmetic operations for fewer serial steps, so it pays off mainly when decoding is memory-bandwidth bound.

Pick the integer maximizing the walltime improvement for the given \(\alpha\) and \(c\); it grows with \(\alpha\) and shrinks as \(c\) grows. Varying \(\gamma\) on the fly by predicting \(\beta\) could help further, which they leave for future work.

Any \(M_q\) preserves correctness. Off-the-shelf smaller Transformers about two orders of magnitude smaller than \(M_p\) worked best; negligible-cost options include n-gram models, copying from the context (for summarization or iterative editing), non-autoregressive models, and even random tokens.

\(\alpha \approx 0.2\), which with \(c \approx 0\) and \(\gamma = 3\) gives about a 1.25× speedup: a trivial table-lookup model already helps, because its cost is negligible.

With T5-small (77M) as approximation model on a TPU-v4 at batch size 1: 3.4× (argmax) and 2.6× (sampling) on English–German translation, 3.1× and 2.3× on CNN/DM summarization, with identical outputs; larger approximation models gave higher \(\alpha\) but less speedup.

Approximation models a couple of orders of magnitude smaller than the target gave \(\alpha\) between 0.5 and 0.9 (e.g. LaMDA 137B with LaMDA 8B: 0.75); \(\alpha\) grew with the size of \(M_q\) and was higher for sharper distributions (argmax); even unigram and bigram models gave non-negligible \(\alpha\).

Blockwise parallel decoding also produces several tokens per step, but supports only greedy decoding, requires training a custom model, and aims to preserve task quality rather than to guarantee identical outputs; speculative decoding handles general stochastic sampling with unchanged models and an exact guarantee.