Speculative Decoding Math: Algorithms & Speedup Limits › The Acceptance Criterion

The basic challenge in speculative decoding is accepting draft tokens in a way that produces the exact same distribution as sampling directly from the target model. If we simply accept tokens when the draft and target models agree, we would bias the output toward the intersection of their distributions. Instead, we need a principled rejection sampling approach.

The acceptance criterion is the heart of speculative decoding. Everything else, the expected speedup, the optimal draft length, the effect of model quality, all flows from this single decision rule. Getting this criterion right means speculative decoding is as correct as running only the target model. Getting it wrong means the output distribution changes subtly, potentially introducing biases that are hard to detect but real in their effect on generated text quality.

To understand why this problem is non-trivial, consider a small vocabulary example. Suppose our vocabulary has just four tokens and the target model assigns probabilities $(0.5, 0.3, 0.15, 0.05)$ while the draft model assigns $(0.6, 0.2, 0.1, 0.1)$. The models agree that the first token is most likely, but they disagree on the others. If we simply accept the draft token whenever it happens to match what the target would have chosen, we need to think carefully about what distribution that produces. The answer depends on complex interactions between the two distributions that are not immediately obvious without the mathematical framework we are about to develop.

Rejection sampling provides exactly the right tool for this problem. The classical version of rejection sampling, developed by John von Neumann in the 1950s, allows you to sample from a target distribution by filtering samples from a proposal distribution. The core idea is to accept proposal samples with a probability that is proportional to how much the target distribution "wants" that sample relative to how often the proposal generates it. Applied to token acceptance, this gives us a precise mathematical rule.

To understand the acceptance criterion, we must first appreciate the subtle problem it solves. When a draft model proposes a token, we face a basic question: how do we decide whether to keep that token while ensuring our final output looks exactly as if we had sampled directly from the target model? The decision must maintain a precise mathematical relationship between what we accept and what the target model would have produced on its own.

Think of the two distributions as two people independently ranking a list of candidates. The target model's ranking is authoritative. The draft model's ranking is a fast approximation. We want to use the fast approximation as much as possible, but we need to correct for cases where the approximation's preferences differ from the authoritative ranking. The acceptance criterion is precisely this correction mechanism.

The formal setup is clean. Let $p(x)$ denote the target model's probability distribution over the next token, and let $q(x)$ denote the draft model's distribution. These are distributions over the same vocabulary $V$, so $\sum_{x \in V} p(x) = \sum_{x \in V} q(x) = 1$. When the draft model proposes token $x$ (sampled from $q$), we need an acceptance probability $\alpha(x)$ such that the final output, considering both accepted draft tokens and resampled tokens on rejection, follows exactly $p(x)$.

The key insight is that rejection sampling provides an acceptance probability that automatically adjusts for the ratio of target to draft probabilities. If we accept with probability:

$$\alpha(x) = \min\left(1, \frac{p(x)}{q(x)}\right)$$

where:

• $\alpha(x)$: acceptance probability for token $x$
• $p(x)$: target model probability of token $x$
• $q(x)$: draft model probability of token $x$

then tokens where $p(x) \geq q(x)$ are always accepted, while tokens where $p(x) < q(x)$ are accepted with probability proportional to how much the target model favors them relative to the draft model.

To build intuition for this formula, consider what happens in two contrasting scenarios. In the first scenario, suppose the target model assigns probability 0.3 to token $x$ while the draft model only assigns probability 0.1. The ratio $p(x)/q(x) = 3$ exceeds 1, so we set $\alpha(x) = 1$ and always accept. This makes sense: the draft model is underproposing this token relative to what the target wants, so we should accept it whenever it appears. In the second scenario, suppose the target model assigns probability 0.1 to token $x$ while the draft model assigns probability 0.3. Now the ratio $p(x)/q(x) = 1/3$, so we accept with probability 1/3. The draft model is overproposing this token, so we need to reject some instances to bring its frequency down to what the target model expects.

The minimum function ensures acceptance probabilities stay in $[0, 1]$. When $p(x) \geq q(x)$, the ratio exceeds 1, but we cap it at 1 since probabilities cannot exceed 1. When $p(x) < q(x)$, the ratio is less than 1 and represents the correct acceptance probability. The symmetry of this rule is elegant: tokens are accepted or rejected based purely on whether the draft model over- or underrepresents them relative to the target.

Let's verify that this acceptance criterion produces the correct distribution. The verification proceeds through careful probability calculations that track what happens when we combine sampling from the draft model with our acceptance decision. When we sample $x \sim q(x)$ and accept with probability $\alpha(x)$, the probability of accepting a specific token $x$ is:

$$\begin{aligned} P(\text{accept } x) &= q(x) \cdot \alpha(x) && \text{(definition of joint probability)} \\ &= q(x) \cdot \min\left(1, \frac{p(x)}{q(x)}\right) && \text{(substitute } \alpha(x) \text{)} \\ &= \min\left(q(x), q(x) \cdot \frac{p(x)}{q(x)}\right) && \text{(distribute } q(x) \text{)} \\ &= \min(q(x), p(x)) && \text{(simplify)} \end{aligned}$$

where:

• $P(\text{accept } x)$: probability of generating and accepting token $x$
• $q(x)$: probability of drafting token $x$
• $\alpha(x)$: probability of accepting drafted token $x$
• $p(x)$: probability of token $x$ under the target distribution

This derivation reveals that the probability of accepting any particular token $x$ is the minimum of the two distributions at that point. This creates a kind of "clipping" effect where we keep only the overlapping probability mass between the draft and target distributions. Think of it as finding the common ground between two people's preferences: we only commit to choices where both parties have at least some agreement.

The total acceptance probability across all tokens is:

$$P(\text{accept}) = \sum_x \min(q(x), p(x))$$

where:

• $P(\text{accept})$: probability that a drafted token is accepted
• $\sum_x$: sum over the entire vocabulary
• $q(x)$: draft model probability of token $x$
• $p(x)$: target model probability of token $x$

This sum has a beautiful geometric interpretation. If you plot both distributions as histograms over the vocabulary, $P(\text{accept})$ equals the total area of overlap between the two histograms. When the distributions are identical, every token is accepted (total overlap equals 1). When they share no common support, no tokens are accepted (overlap equals 0). The more similar the two models, the more frequently we can reuse draft tokens without needing to resample.

Conditional on acceptance, the distribution over accepted tokens is:

$$P(x \mid \text{accept}) = \frac{\min(q(x), p(x))}{\sum_{x'} \min(q(x'), p(x'))}$$

where:

• $P(x \mid \text{accept})$: probability of token $x$ given it was accepted
• $\sum_{x'}$: normalization sum over all possible tokens $x'$
• $q(x)$: draft model probability
• $p(x)$: target model probability

This does not yet equal $p(x)$. The distribution is corrected by handling rejections appropriately.

When we reject a draft token, we don't simply redraft. Instead, we sample from a carefully constructed residual distribution that "fills in" the probability mass that rejection sampling missed. This residual distribution is the mathematical key that makes speculative decoding exact rather than approximate.

To understand why we need a residual distribution, consider what happens with pure rejection sampling. When we accept with probability $\min(q(x), p(x))$, we capture all the probability mass where the draft distribution meets or exceeds the target. But what about tokens where $p(x) > q(x)$? For these tokens, the draft model underestimates the target probability, and our rejection sampling only captures the portion up to $q(x)$. If we simply re-ran the draft model on rejection, we would keep trying to sample from $q$ and would never correctly generate the tokens that $p$ favors more than $q$.

Think of the residual distribution as a "correction fund." The acceptance mechanism captures the portion of $p$ that $q$ already covers. The residual distribution covers what $q$ missed. Together, they add up to the full target distribution $p$. The residual samples only from tokens where $p(x) > q(x)$, specifically in proportion to how large that gap is, which is exactly what we need to account for the "missing mass."

Define the residual distribution as:

$$p_{\text{resid}}(x) = \frac{\max(0, p(x) - q(x))}{Z}$$

where:

• $p_{\text{resid}}(x)$: residual distribution probability for token $x$
• $p(x)$: target model probability of token $x$
• $q(x)$: draft model probability of token $x$
• $\max(0, p(x) - q(x))$: the positive part of the difference between target and draft probabilities (zero when draft exceeds target)
• $Z$: normalizing constant, $Z = \sum_x \max(0, p(x) - q(x))$

The residual distribution captures the probability mass where $p(x) > q(x)$, which is exactly the mass that acceptance sampling misses. Geometrically, if you imagine the target distribution as a histogram and the draft distribution as another histogram overlaid on top, the residual distribution represents the portions of the target histogram that "stick out above" the draft histogram. By sampling from $p_{\text{resid}}$ on rejection, we ensure the overall distribution matches $p(x)$.

Notice that the normalizing constant $Z$ equals $1 - P(\text{accept})$, which is the total rejection probability. To see this, note that $\sum_x \max(0, p(x) - q(x)) = \sum_x (p(x) - \min(p(x), q(x))) = 1 - \sum_x \min(p(x), q(x)) = 1 - P(\text{accept})$. The probability mass captured by the residual distribution exactly equals the probability mass that rejection sampling misses. This keeps everything adds up correctly and the final distribution is properly normalized.

The full acceptance procedure for a single token position works as follows. This algorithm combines the acceptance criterion with the residual distribution to guarantee exact sampling from the target model:

1. Sample draft token $x \sim q(x)$
2. Sample uniform $u \sim \text{Uniform}(0, 1)$
3. If $u < \frac{p(x)}{q(x)}$, accept $x$
4. Otherwise, sample $x \sim p_{\text{resid}}(x)$

This procedure guarantees that the output follows distribution $p(x)$ exactly. To see why, consider that there are two mutually exclusive ways to generate a token $x$. The first path is through acceptance: we draft $x$ with probability $q(x)$ and accept it with probability $\min(1, p(x)/q(x))$, contributing $\min(q(x), p(x))$ to the total probability. The second path is through rejection and resampling: with probability $1 - P(\text{accept})$ we reject and then sample $x$ from the residual with probability $\frac{\max(0, p(x) - q(x))}{Z}$, contributing $\max(0, p(x) - q(x))$ to the total probability. Adding these two contributions gives $\min(q(x), p(x)) + \max(0, p(x) - q(x)) = p(x)$, exactly as desired.

The key insight is that the two paths together cover the entire target distribution without double-counting. The acceptance path handles the overlap region (where both models assign positive probability), while the residual path handles the regions where the target model exceeds the draft model. The algebraic identity $\min(a,b) + \max(0, a-b) = a$ is the mathematical guarantee that these two contributions sum to the target probability exactly.

In practice, the draft model proposes $k$ tokens at once. We verify them sequentially from left to right, accepting tokens until the first rejection. This sequential verification is needed because language models produce conditional distributions: the probability of each token depends on all preceding tokens.

The sequential structure of verification mirrors the sequential nature of language itself. Each word in a sentence depends on everything that came before it. When the draft model proposes "the cat sat on the mat," the probability of "sat" depends on having seen "the cat" first, and the probability of "on" depends on having seen "the cat sat" first. If we reject "cat" and replace it with something else, then all subsequent positions, "sat," "on," "the," "mat," become invalid because they were conditioned on "cat." This is why we must verify left to right and stop at the first rejection.

Let $x_1, x_2, \ldots, x_k$ be the draft sequence. For position $i$, conditioning on the prefix $x_1, \ldots, x_{i-1}$:

$$\alpha_i = \min\left(1, \frac{p(x_i \mid x_1, \ldots, x_{i-1})}{q(x_i \mid x_1, \ldots, x_{i-1})}\right)$$

where:

• $\alpha_i$: acceptance probability for the $i$-th token
• $x_i$: the $i$-th token in the draft sequence
• $p(x_i \mid \dots)$: target model conditional probability given all preceding tokens
• $q(x_i \mid \dots)$: draft model conditional probability given all preceding tokens

If position $i$ rejects, we sample from the residual distribution and discard positions $i+1, \ldots, k$. This maintains correctness because each accepted token was drawn from the correct conditional distribution. The key insight is that we cannot "skip" a rejection and continue verifying later tokens, because those tokens were conditioned on the rejected token. Once we sample a different token from the residual distribution, the entire subsequent sequence becomes invalid and must be regenerated.

This left-to-right verification creates an important efficiency consideration. Even though we verify all $k$ positions in a single parallel forward pass of the target model, we can only use tokens up to the first rejection. However, because we always sample from the residual distribution upon rejection, we are guaranteed to produce at least one valid token per iteration, even if all $k$ draft tokens are rejected. In fact, when all tokens are accepted, we get $k + 1$ tokens, since the target model's forward pass also provides the distribution for the next position beyond the draft sequence. This bonus token is a small but consistent benefit: we always get at least one more token than we pay for in draft generation.