Speculative Decoding Is a Maximal Coupling
11 min read
Speculative decoding is the trick behind much of fast language model serving: a small draft model guesses a few tokens ahead, and the big target model checks all the guesses in a single forward pass. It was found twice, nearly simultaneously: by Leviathan, Kalman & Matias [LKM, ICML ’23] at Google and by Chen, Borgeaud, Irving, Lespiau, Sifre & Jumper [CBILSJ, ’23] at DeepMind, with measured wall-clock speedups of 2–3x on T5-XXL and 2–2.5x on the 70B-parameter Chinchilla. The first time I heard it described, I assumed it was a lossy approximation: the small model’s guesses standing in for the big model’s judgment, at some cost in accuracy. There is no such cost: the verification rule is a construction probabilists call a maximal coupling, and that makes the trick exact and optimal at once.
Free Verification
Drafting only helps if checking a guess costs far less than making one, so start with where a forward pass spends its time. Autoregressive decoding produces one token per pass of the target model, and for a large transformer at small batch size that pass is bottlenecked on memory bandwidth, not arithmetic: every parameter is streamed from memory to emit a single token. Model size divided by memory bandwidth is a hard ceiling on tokens per second. But the same streamed weights can score many positions at once, so one pass can evaluate a whole candidate continuation in roughly the time it takes to generate a single token. Checking is nearly free.
So let a cheap draft model propose, one token at a time, a guess at the next \(\gamma\) tokens, where the draft length \(\gamma\) is small, say \(4\) or \(8\) (in practice the draft is 10–100x smaller: an 11B target drafted by a 77M model, a 70B target by a 4B one). One target pass on the prefix plus all \(\gamma\) guesses then hands us its next-token distribution at each of the \(\gamma + 1\) positions, and the guesses are checked against those distributions left to right. Everything now rides on the verification rule at a single position.
The Rule
The draft’s tokens come from the wrong distribution. The obvious move is to keep them anyway, but then the generated text is a sample from the draft model, not the target: precisely the lossy approximation I had assumed. Trusting a guess outright needs the two models to agree on the next token’s distribution, and they do not. So we need a rule for which guesses to keep and how to replace the rest, one that leaves the output distribution untouched.
Fix one position and condition on everything before it; write \(p\) for the target model’s next-token distribution there and \(q\) for the draft model’s, both over the same vocabulary. To see what the rule must do, take a 3-token vocabulary \(\{a, b, c\}\) with target \(p = (0.5, 0.3, 0.2)\) and draft \(q = (0.3, 0.1, 0.6)\). The draft under-serves \(a\) and \(b\) and over-serves \(c\): too many of its guesses come out \(c\), too few come out \(a\) or \(b\). A correcting rule should therefore keep every drawn \(a\) and \(b\), thin the drawn \(c\)’s from the \(0.6\) the draft supplies to the \(0.2\) the target wants, and hand the trimmed mass back to \(a\) and \(b\).
The acceptance rule does the thinning. The draft proposes a token \(x \sim q\), and we accept it with probability
\[ \min\!\Big(1,\; \frac{p(x)}{q(x)}\Big). \]
In words: if the target puts at least as much mass on \(x\) as the draft did, \(p(x) \ge q(x)\), we keep it outright, and otherwise we keep it with probability \(p(x)/q(x)\). The ratio is the largest acceptance probability that does not overshoot \(p\): a token the draft over-proposes is then kept with mass \(q(x) \cdot p(x)/q(x) = p(x)\), exactly what \(p\) assigns it. So acceptance delivers \(\min(p(x), q(x))\) at every token and never more, and covering the shortfall is the job of the rejection branch. In the example, drawn \(a\)’s and \(b\)’s are always accepted, while a drawn \(c\) survives with probability
\[ \frac{p(c)}{q(c)} = \frac{0.2}{0.6} = \frac13. \]
That trims \(c\)’s share of the output to \(0.2\), exactly the target’s mass.
On rejection, don’t ask the draft again; sample once from the residual distribution:
\[ r(x) = \frac{(p(x) - q(x))_+}{\sum_z (p(z) - q(z))_+}, \]
where \((\cdot)_+\) is the positive part and the denominator normalizes. The numerator is the shortfall itself, so \(r\) puts mass only on the tokens the draft under-serves, in proportion to how badly. Accepted or resampled, that token is the output. In the example, rejection happens with probability
\[ q(c)\,\Big(1 - \frac{p(c)}{q(c)}\Big) = 0.6 \cdot \tfrac{2}{3} = 0.4, \]
and the residual puts all of that mass back on the under-served tokens:
\[ r = \frac{(0.2,\; 0.2,\; 0)}{0.4} = \big(\tfrac12,\; \tfrac12,\; 0\big). \]
A rejected \(c\) is converted into a fair coin flip between \(a\) and \(b\). In the plot, the two bars at each token are the draft mass \(q\) and the target mass \(p\), both standing on the shared mass \(\min(p, q)\) in teal, with the draft’s excess in gray and the residual in amber.
The two branches add back to \(p\) here. Whether they always do, for any \(p\) and \(q\), is what the acceptance probability was built for.
Why the Output Is Exact
There are two ways the output can equal \(x\): the draft proposed \(x\) and we accepted, or we rejected and the residual resampled \(x\). The two are exclusive, so the output mass at \(x\) is their sum.
Take acceptance first. The draft proposes \(x\) with probability \(q(x)\) and we keep it with probability \(\min(1, p(x)/q(x))\), so the accepted mass at \(x\) is the product:
\[ \begin{aligned} q(x) \cdot \min\!\Big(1,\; \frac{p(x)}{q(x)}\Big) &= \min\!\Big(q(x) \cdot 1,\; q(x) \cdot \frac{p(x)}{q(x)}\Big) \\ &= \min\big(q(x),\; p(x)\big), \end{aligned} \]
where the first line uses \(q(x) \ge 0\) to multiply both arguments of the minimum, since a nonnegative factor cannot change which is smaller, and the second cancels \(q(x)\). (If \(q(x) = 0\) the draft never proposes \(x\), both sides are zero, and the ratio is never evaluated.)
Rejection is the complement, so its probability is one minus the accepted mass summed over the vocabulary:
\[ \begin{aligned} 1 - \sum_x \min(p(x), q(x)) &= \sum_x p(x) - \sum_x \min(p(x), q(x)) \\ &= \sum_x \big(p(x) - \min(p(x), q(x))\big) \\ &= \sum_x (p(x) - q(x))_+ , \end{aligned} \]
where the first line replaces \(1\) with \(\sum_x p(x)\), the second collects the two sums into one, and the third uses that \(p(x) - \min(p(x), q(x))\) is the positive part \((p(x) - q(x))_+\). That total is exactly the residual’s normalizing constant, so multiplying it by \(r(x)\) cancels the denominator:
\[ \begin{aligned} \Pr(\text{reject}) \cdot r(x) &= \Big(\textstyle\sum_z (p(z) - q(z))_+\Big) \cdot \frac{(p(x) - q(x))_+}{\sum_z (p(z) - q(z))_+} \\ &= (p(x) - q(x))_+ . \end{aligned} \]
Adding the two branches gives the output mass at \(x\):
\[ \min(p(x), q(x)) + (p(x) - q(x))_+ = p(x) . \]
When \(p(x) \le q(x)\) the first term is \(p(x)\) and the second is zero; can you check the other case in your head? So the output is distributed exactly as \(p\), and on the example token \(a\) collects \(0.3\) from acceptance plus \(0.4 \cdot \tfrac12 = 0.2\) from resampling, for \(p(a) = 0.5\). This is a sharper form of rejection sampling: the classical version scales \(q\) up by the worst-case ratio \(\max_x p(x)/q(x)\) and retries until a draw is accepted, while here one resample settles the position.
That settles one position, and a draft covers \(\gamma\). Accept the drafted tokens left to right; at the first rejection, resample from that position’s residual and throw the rest of the draft away. The tail has to go because those guesses were drawn conditioned on the rejected token, so they come from the wrong conditionals. Each surviving position’s output is exactly the target’s conditional distribution given the prefix, so by induction the whole generated text is an exact sample from the target model.
Exactness holds for any draft model, however bad, since a hopeless draft is just rejected more often. What the draft model sets is the acceptance rate, and that rate has a ceiling.
The Best Possible Acceptance Rate
How often does the rule accept, and could a different exact rule accept more often? Summing the accepted mass over the vocabulary answers the first: the acceptance probability is \(\sum_x \min(p(x), q(x))\), the teal mass in the plot above. That quantity has a name: it is one minus the total variation distance,
\[ \mathrm{TV}(p, q) = \tfrac12 \sum_x |p(x) - q(x)| . \]
To see the connection, split \(p - q\) into its positive and negative parts. The two have equal total mass, because the full difference sums to zero:
\[ \sum_x (p(x) - q(x))_+ - \sum_x (q(x) - p(x))_+ = \sum_x \big(p(x) - q(x)\big) = 0 . \]
Their sum, meanwhile, is the total absolute difference:
\[ \sum_x (p(x) - q(x))_+ + \sum_x (q(x) - p(x))_+ = \sum_x |p(x) - q(x)| . \]
Two equal quantities that add to \(\sum_x |p - q|\) are each half of it, so the rejection probability is exactly the total variation distance:
\[ \sum_x (p(x) - q(x))_+ = \tfrac12 \sum_x |p(x) - q(x)| = \mathrm{TV}(p, q) . \]
In the example, \(\mathrm{TV}(p, q) = \tfrac12 (0.2 + 0.2 + 0.4) = 0.4\), the rejection probability we computed by hand.
That the acceptance rate lands on such a clean quantity is no accident; it is forced by an optimality argument. A coupling of \(p\) and \(q\) is any joint distribution for a pair \((X, Y)\) whose marginals are \(X \sim p\) and \(Y \sim q\). The speculative step couples the output with the draft: on acceptance, the two are equal.
Claim: any scheme that sees the draft’s token \(Y \sim q\) and outputs a token \(X\) distributed exactly as \(p\) keeps the draft’s token with probability at most \(1 - \mathrm{TV}(p, q)\), and the rule above achieves the bound.
In words, two tokens can agree no more often than their distributions overlap.
Proof
The joint law of \((X, Y)\) is a coupling of \(p\) and \(q\), whatever the scheme’s internal randomness. For any token \(x\), the event \(\{X = Y = x\}\) sits inside both \(\{X = x\}\) and \(\{Y = x\}\), so its probability is at most the smaller marginal:
\[ \Pr(X = Y = x) \le \min(p(x), q(x)) . \]
Summing over the vocabulary gives the coupling inequality:
\[ \Pr(X = Y) = \sum_x \Pr(X = Y = x) \le \sum_x \min(p(x), q(x)) = 1 - \mathrm{TV}(p, q) . \]
Keeping the draft’s token is an event on which \(X = Y\), so its probability obeys the same bound. The speculative rule accepts with total probability
\[ \sum_x q(x)\, \min\!\big(1, p(x)/q(x)\big) = \sum_x \min(p(x), q(x)) , \]
by the accepted-mass computation above, so it meets the bound with equality.
A coupling that meets the bound is a maximal coupling, and building one is a textbook exercise (Chapter 4 of Levin & Peres [LP, ’17]): lay down the shared mass \(\min(p, q)\), then complete each marginal with its normalized leftover. Speculative decoding is exactly that construction run causally: the draft proposes, and the target corrects only the tokens it under-served. So the acceptance rate is fixed by the two models themselves: it is the overlap \(1 - \mathrm{TV}(p, q)\) of their next-token distributions, and no cleverer threshold can push it higher. The only levers left are a better draft model and the draft length \(\gamma\).
The Throughput Accounting
An acceptance rate is not yet a speedup. What matters is tokens per target pass, and the draft length moves that as much as the acceptance rate does.
Say each drafted token is accepted with probability \(\beta\), independently across positions. This is a simplification: really \(\beta = 1 - \mathrm{TV}(p, q)\) shifts with the prefix, and we take \(\beta\) to be its average. A single target pass delivers the accepted prefix plus one more token: at a rejection the residual supplies it, and if all \(\gamma\) drafts survive, the same pass has already computed the distribution at position \(\gamma + 1\).
The number of tokens from a pass runs from \(1\) to \(\gamma + 1\), and it exceeds \(k\) exactly when the first \(k\) drafts are all accepted, an event of probability \(\beta^k\). For such a count the expectation is the sum of the tail probabilities:
\[ \begin{aligned} \mathbb{E}[\text{tokens per pass}] &= \sum_{k=0}^{\gamma} \Pr(\text{tokens} > k) \\ &= \sum_{k=0}^{\gamma} \beta^k \\ &= \frac{1 - \beta^{\gamma+1}}{1 - \beta} , \end{aligned} \]
a geometric series truncated by the cap at \(\gamma + 1\).
In the plot, each curve bends toward its ceiling \(1/(1 - \beta)\): lengthening the draft from \(k - 1\) to \(k\) adds only \(\beta^k\) to the expectation. Drafting is not free either: each round runs \(\gamma\) draft passes on top of the target pass. If a draft pass takes a fraction \(c\) of a target pass, one round costs \(\gamma c + 1\) target passes, so the speedup over plain decoding is
\[ \frac{1 - \beta^{\gamma+1}}{(1 - \beta)(\gamma c + 1)} . \]
The numerator saturates in \(\gamma\) while the denominator grows linearly, so the best draft length is finite and modest. At \(\beta = 0.8\) and \(\gamma = 5\), the expectation is about \(3.7\) tokens per target pass, and the measured 2–3x wall-clock speedups are that figure less the drafting overhead.
Speculative decoding now ships in the major production serving stacks. Its descendants leave the accept/reject core untouched, because the coupling argument says that part cannot be improved, only fed better drafts: tree-shaped drafts propose several continuations at once so one target pass can check a whole branching set, and self-speculation has the target model draft for itself from a subset of its own layers. Speculation risks only wasted time, never correctness, and the coupling shows even that risk is as small as any exact scheme allows. The portable lesson: an accept/reject rule that is a maximal coupling is exact by construction, and all that is left to tune is how often it resamples.