Speculative Decoding Is a Maximal Coupling
11 min read
A large language model can be slow because generating each new token requires another pass through the model. Speculative decoding lets a small draft model guess several tokens ahead, then asks the large target model to check those guesses together [LKM, ICML ’23], [CBILSJ, ’23]. When the guesses are accepted, one large-model pass produces several tokens.
The first time I heard this described, I assumed the speed came at some cost in accuracy. The verification rule preserves exactly the distribution we would get by sampling from the target model itself. It accepts some draft tokens and, on rejection, samples a replacement to supply the target probabilities still missing. For a single proposed token, it achieves the greatest acceptance probability compatible with those exact output probabilities. That combination of exactness and the largest possible agreement is what makes the rule a maximal coupling.
Free Verification
Before the accuracy question, why is checking \(\gamma\) guesses any cheaper than making them? Drafting only helps if checking a guess costs far less than making one. 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. 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 produces its next-token distribution at each of the \(\gamma + 1\) positions, and the guesses are checked against those distributions left to right. The check itself gives up no accuracy at all: the rule that keeps or replaces each guess is a maximal coupling of the draft and target distributions, so the output is exactly the target model’s distribution, and no exact rule could keep the draft’s guesses more often.
The Rule
Which guesses do we keep, and what replaces a rejected one? Keeping every draft token would sample from the draft model rather than the target. The verification rule must decide which guesses to keep and how to replace the rest while leaving the target distribution unchanged.
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 move the trimmed mass back to \(a\) and \(b\).
The draft proposes a token \(x \sim q\), and the acceptance probability thins any excess mass:
\[ \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 kept with mass \(q(x) \cdot p(x)/q(x) = p(x)\), so acceptance contributes \(\min(p(x), q(x))\) at token \(x\). 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. \]
The accepted mass at \(c\) is therefore \(0.2\), 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, so \(r\) puts mass only on tokens the draft under-serves and in proportion to that shortfall. 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.
Exactness of the Output
Why is the output distributed exactly as \(p\), whatever the draft proposed? 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.
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:
\[ \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))_+\). The rejection probability is 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. When \(p(x)>q(x)\) the two terms sum to \(q(x)+p(x)-q(x)=p(x)\). Thus the output is distributed as \(p\); in 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.
For a draft covering \(\gamma\) positions, accept its tokens from 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 the target’s conditional distribution given the prefix, so by induction the whole generated text is an exact sample from the target model for any draft; a poor draft is rejected more often.
The Best Possible Acceptance Rate
Could a cleverer exact rule keep more of the draft’s tokens? The acceptance probability is the accepted mass summed over the vocabulary, \(\sum_x \min(p(x), q(x))\), shown in teal above. This overlap 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.
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 the overlap,
\[ \Pr(X = Y) \le 1 - \mathrm{TV}(p, q), \]
and the rule above achieves the bound.
Proof
The joint distribution 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. The textbook construction lays down the shared mass \(\min(p, q)\) and completes each marginal with its normalized leftover (Chapter 4 of Levin & Peres [LP, ’17]). Speculative decoding runs this construction causally: the draft proposes, and the target corrects the tokens it under-served. Its acceptance rate is therefore the overlap \(1 - \mathrm{TV}(p, q)\), leaving the draft model and draft length \(\gamma\) as the throughput choices.
The Throughput Accounting
How many tokens does one target pass buy, and how long should the draft be? Throughput depends on both the acceptance rate and the number of tokens returned per target pass.
Approximate the drafted positions as independent, each accepted with probability \(\beta\). In practice, \(\beta = 1 - \mathrm{TV}(p, q)\) shifts with the prefix; here \(\beta\) denotes 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 keep the accept/reject rule and supply better drafts. Tree-shaped methods propose several continuations so one target pass can check a branching set, while self-speculation has the target model draft from a subset of its own layers. Because maximal coupling preserves the target distribution and maximizes acceptance, a rejected draft costs time but does not change the generated distribution.