A rejection rule keeps speculative decoding's output exact

A forward pass scores a whole prefix for the price of one token

Autoregressive decoding calls the target model once per output token. Each call is a full forward pass: every layer's weights, read from memory, produce one next-token distribution, then get discarded. The same forward pass, run over a longer prefix, scores every position in that prefix for roughly the same memory traffic one position costs alone. A transformer's attention and feed-forward layers apply to a batch of positions at once, not one at a time. Hand the target model five extra draft tokens appended to a prompt, and one pass returns five next-token distributions, not one.

Speculative decoding spends that asymmetry directly: a small, cheap draft model guesses several tokens ahead, and the large target model checks the whole guess in one batched pass instead of one serial call per token. Leviathan, Kalman, and Matias reported "2X-3X acceleration compared to the standard T5X implementation, with identical outputs" running exactly this scheme on an 11-billion-parameter T5-XXL model.¹ "Identical outputs" is the part worth rebuilding by hand. The speedup does not come from trusting the draft's guesses. It comes from a specific rejection-sampling rule applied to every guess. That rule is provable, and small enough to run to completion in a few hundred lines of Python.

A target and a draft small enough to print in full

Checking an exactness proof needs distributions a reader can see in full, not a 70-billion-parameter black box. The target here, call it p, is an order-1 character bigram model: p(x \mid \text{previous character}). It is trained on an 80-character corpus, so every next-character distribution is a table of eighteen numbers that sums to one. The draft, q_w, is a mixture of that same bigram model and a context-free unigram model: q_w(x \mid \text{ctx}) = w \cdot p(x \mid \text{ctx}) + (1-w) \cdot q_0(x), where q_0 ignores context entirely. At w=0 the draft has thrown away the previous character: the cheapest, worst-agreeing model available. As w rises toward 1, the draft converges on the target. That single knob stands in for what a real draft-target pair varies by parameter count and training data: how much of the target's context the draft actually models. Every probability in this piece stays exactly enumerable regardless.

This uses Leviathan et al.'s letters throughout: p is the target, q is the draft.¹ Chen et al. derived the identical rule independently at DeepMind, with the letters swapped: q as target, p as draft.² Any formula quoted from their paper needs relabeling before it sits next to code written in the other convention.

CORPUS = (
    "to be or not to be that is the question whether tis nobler in the mind "
    "to suffer"
)
VOCAB = sorted(set(CORPUS))          # tokens the models can emit
BOS = "^"                            # start-of-sequence context, never emitted
CONTEXTS = [BOS] + VOCAB             # possible "previous token" values

def normalize(counts):
    total = sum(counts.values())
    return {k: v / total for k, v in counts.items()}

# Order-1 bigram counts (target model), Laplace-smoothed so every
# (context, token) pair has positive probability.
bigram_counts = {ctx: {tok: 1 for tok in VOCAB} for ctx in CONTEXTS}
prev = BOS
for ch in CORPUS:
    bigram_counts[prev][ch] += 1
    prev = ch
P_TABLE = {ctx: normalize(bigram_counts[ctx]) for ctx in CONTEXTS}

# Order-0 unigram counts (the cheap, context-free half of the draft).
unigram_counts = {tok: 1 for tok in VOCAB}
for ch in CORPUS:
    unigram_counts[ch] += 1
Q0_TABLE = normalize(unigram_counts)

def p_dist(ctx):
    return P_TABLE[ctx]

def q_dist(ctx, w):
    p = P_TABLE[ctx]
    return {tok: w * p[tok] + (1 - w) * Q0_TABLE[tok] for tok in VOCAB}

def sample_categorical(dist, rng=random):
    r = rng.random()
    acc = 0.0
    for tok, prob in dist.items():
        acc += prob
        if r <= acc:
            return tok
    return next(reversed(list(dist)))  # floating-point fallback

def autoregressive_sample(n_tokens, ctx=BOS, rng=random):
    out = []
    for _ in range(n_tokens):
        tok = sample_categorical(p_dist(ctx), rng)
        out.append(tok)
        ctx = tok
    return out

Running autoregressive_sample two hundred thousand times from the context 't', and tallying the outcomes against P_TABLE['t'], checks nothing about speculative decoding yet. It is a sanity check on sample_categorical itself, before anything else gets built on top of it:

The largest gap over the full 18-token vocabulary is 0.00076, well inside the noise a 200,000-draw Monte Carlo estimate of an 0.1-scale probability should show. sample_categorical and P_TABLE are trustworthy ground truth for everything that follows.

What the accept/reject rule has to balance

A draft token cannot simply be kept whenever the draft model produced it. If every draft token were kept, the output would follow q, not p. A cheaper model would quietly be steering the sequence. But a rule that rejects too liberally throws away work the draft model already paid for, and the whole appeal of drafting collapses back into calling the target model one token at a time. The rule needs three things at once. It has to keep every draft token the target model would have produced at least as often as the draft did. It has to discard the excess when the draft over-produced a token relative to the target. And it has to route what got discarded toward whatever the draft under-produced. That is exactly what a ratio test does.

Draw a candidate x \sim q from the draft, then draw a fresh uniform r \sim U(0,1). Accept x when r falls under the ratio of the two models' probabilities for it, capped at 1. When the draft under-weighted x relative to the target (q(x) \le p(x)), the ratio is at least 1 and x is always kept. When the draft over-weighted it (q(x) > p(x)), the ratio is p(x)/q(x) < 1 and x survives only that often. On rejection, resample from what is left of the target once the draft's own mass is subtracted out at that token. That remainder is the positive part of p - q, renormalized to sum to 1, so it is a real distribution again.

Leviathan et al. state this exact rule in Section 2.3: sample x \sim q(x), "keeping it if q(x) \le p(x)." In case q(x) > p(x), they "reject the sample with probability 1 - \frac{p(x)}{q(x)} and sample x again from an adjusted distribution p'(x) = \operatorname{norm}(\max(0, p(x) - q(x))) instead."¹ Chen et al.'s "Modified Rejection Sampling" states the identical construction, target-model probability over draft-model probability, with their swapped letters.²

def residual_dist(ctx, w):
    """norm(max(0, p(x|ctx) - q_w(x|ctx)))."""
    p = p_dist(ctx)
    q = q_dist(ctx, w)
    raw = {tok: max(0.0, p[tok] - q[tok]) for tok in VOCAB}
    z = sum(raw.values())
    return {tok: v / z for tok, v in raw.items()}

def speculative_step(ctx, gamma, w, rng=random):
    draft, q_dists = [], []
    c = ctx
    for _ in range(gamma):
        qi = q_dist(c, w)
        xi = sample_categorical(qi, rng)
        draft.append(xi)
        q_dists.append(qi)
        c = xi

    # Batched verification: p at every position along the drafted
    # prefix, computed from the original context forward.
    p_dists = []
    c = ctx
    for i in range(gamma):
        p_dists.append(p_dist(c))
        c = draft[i]
    p_dists.append(p_dist(c))

    n = 0
    n_tests = 0
    for i in range(gamma):
        x = draft[i]
        ratio = min(1.0, p_dists[i][x] / q_dists[i][x])
        n_tests += 1
        if rng.random() <= ratio:
            n += 1
        else:
            break

    if n == gamma:
        bonus = sample_categorical(p_dists[gamma], rng)
        return draft[:n] + [bonus], n, n_tests
    resample_ctx = ctx if n == 0 else draft[n - 1]
    t = sample_categorical(residual_dist(resample_ctx, w), rng)
    return draft[:n] + [t], n, n_tests

Fixing gamma=1 collapses this to a single-token test, the smallest case the exactness claim has to hold for. Running it two hundred thousand times from the same context 't' used above, with the draft at w=0.3, gives the table below. At that weight the draft keeps three tenths of the target's context and throws the rest into the unigram average:

The largest disagreement over all 18 tokens is 0.00146. That is the same order as the plain autoregressive sampler's own Monte Carlo noise measured above, not a systematic gap. That is what the rule in Figure 2 predicts, not a coincidence. For any fixed x', the chance of producing it via direct acceptance is q(x') \cdot \min(1, p(x')/q(x')) = \min(p(x'), q(x')). The chance of reaching it via the residual path is the chance of rejection, times the chance the residual draw lands on x'. The rejection probability's 1-\beta cancels exactly against the residual's own normalizing constant, leaving \max(0, p(x')-q(x')). Summing the two paths:

When p(x') \ge q(x') this is q(x') + (p(x')-q(x')) = p(x'); when p(x') < q(x') it is p(x') + 0 = p(x'). Either way the two paths add up to exactly the target's own probability at x', for any x', regardless of how p and q compare there. Leviathan et al. prove this in Appendix A.1; Chen et al. prove it independently as their Theorem 1, with the letters swapped and q the recovered target.¹² The table above is that proof, run.

Extending the rule from one token to a drafted run

The payoff from speculative decoding comes from drafting several tokens at once and verifying all of them in one batched pass, not from any single token. speculative_step above already does this for any gamma. It drafts \gamma tokens from q, computes p at all \gamma+1 positions along that draft in one batch, then walks the accept test left to right and stops at the first rejection. If every drafted token survives, a bonus token is sampled from the target's own distribution at the position after the draft. That distribution is already computed for free in the same batch, since the pass covers one position more than there are draft tokens. Chaining these steps, each one picking up where the last left off, produces a full sequence:

def speculative_generate(n_tokens, gamma, w, ctx=BOS, rng=random):
    out, steps = [], []
    c = ctx
    while len(out) < n_tokens:
        tokens, n, n_tests = speculative_step(c, gamma, w, rng)
        steps.append((len(tokens), n, n_tests))
        out.extend(tokens)
        c = tokens[-1]
    return out[:n_tokens], steps

Neither paper states a separate theorem for the joint distribution of a whole chained sequence. The single-position result above, applied position by position with each step conditioned on the accepted prefix so far, is what a real decoding loop actually relies on. It is worth checking directly rather than taking it on faith. Generating two tokens from the start-of-sequence context two hundred thousand times, with \gamma=2 and w=0.3, gives one way to check that composition. Compare the joint frequency of each two-character pair against the target chain's own exact joint, p(x_1)\cdot p(x_2 \mid x_1), computed directly from the same tables:

Every one of the 324 possible pairs occurred in both the exact and the empirical distribution. Summed over that full support, the total variation distance between the two joint distributions is 0.01523. That is the same kind of sampling noise as the single-token checks above, spread over a larger outcome space, not a bias that compounds as the draft runs longer. The proof composes the way a real decoding loop needs it to.

What draft accuracy buys, measured in tokens per call

Exactness alone does not explain why speculative decoding is faster. A draft that resamples from the residual on almost every token is exact and useless. What makes it worth building is how often a draft token survives the test, since a survived token is one output token produced without a serial call to the target model. Leviathan et al. name the per-position survival probability \beta, and its average over a generated sequence \alpha := E[\beta]:

\beta is exactly the overlap between p and q at a given context: the mass the two distributions already agree on. It can be computed directly from the tables above, with no sampling at all. Sum the smaller of the two probabilities at every token, then average over the target chain's own long-run distribution of contexts:

def beta_exact(ctx, w):
    p, q = p_dist(ctx), q_dist(ctx, w)
    return sum(min(p[t], q[t]) for t in VOCAB)

def alpha_exact(w):
    # PI_P: the target chain's stationary context distribution,
    # found by power-iterating P_TABLE to convergence.
    return sum(PI_P[ctx] * beta_exact(ctx, w) for ctx in VOCAB)

def expected_tokens(alpha, gamma):
    if alpha >= 1.0:
        return gamma + 1
    return (1 - alpha ** (gamma + 1)) / (1 - alpha)

With \gamma=4 draft tokens per step, each iteration produces a geometric count of tokens capped at \gamma+1, in expectation E[\#\text{tokens}] = \frac{1-\alpha^{\gamma+1}}{1-\alpha} (Leviathan et al., Equation 1).¹ Running speculative_generate for 200,000 tokens at each draft weight and counting the accepted fraction directly checks that formula against the run, not just against itself:

The measured acceptance rate tracks the exact one within four thousandths at every weight, and the measured tokens per step tracks the formula within 0.01 throughout. At w=0, the draft keeps no context at all, yet 72.6% of draft tokens still survive. A plain unigram model already recovers most of a bigram model's mass in a corpus this repetitive, and that alone yields 2.92 tokens per target call instead of 1. At w=1 the draft is the target, every token survives, and each call produces the full \gamma+1=5.

Real acceptance rates span a comparable range. Leviathan et al. measured \alpha \approx 0.20 for a bigram draft against an 11-billion-parameter T5-XXL target, barely better than this toy's worst case. The same paper measured \alpha = 0.88–0.89 for a 6-million-parameter draft against a 97-million-parameter GPT-like target, and 0.74–0.75 for an 8-billion-parameter LaMDA draft against the 137-billion-parameter version.¹ What this section measures, agreement between two explicit distributions setting how many tokens one target call buys, is the same relationship producing those numbers on models eleven orders of magnitude larger.

What changes when the draft and target are real models

This build sidesteps three things by using dictionaries instead of weights, and all three matter once p and q are transformers. First, a KV cache. Recomputing p_dist(ctx) here costs a dictionary lookup regardless of how far into the sequence ctx is. A transformer's cost to score a position depends instead on reusing the key and value tensors computed for every earlier position. A real verification pass has to extend the cache from the last accepted token, not recompute the whole prefix, or the "one batched pass" in Figure 3 is not actually cheap. Second, the batch itself. The loop building p_dists above runs one Python call per position. A real target model's verification is worth doing because those \gamma+1 positions are one matrix multiply, not \gamma+1 sequential ones. Third, the draft model itself. This piece built one by blending the target with a weaker version of itself. A real system has to choose what plays that role.

Leviathan and Chen's original setups use a separate, smaller checkpoint of the same model family.¹² Later work changes what a draft even is. Medusa removes the second model outright. It adds extra decoding heads to the target's own weights that "predict multiple subsequent tokens in parallel," verified with a tree-structured attention pattern. Medusa-1 reaches over 2.2× speedup and Medusa-2 2.3–3.6×, both by the authors' own measurement.³ EAGLE drafts at the level of the target's own hidden features instead of output tokens, on the claim that "autoregression at the feature... level is more straightforward than at the token level." It still keeps the exact accept/reject construction, stating that it is "maintaining the distribution of the generated text," and reports a 2.7×–3.5× latency speedup on LLaMA2-Chat 70B.⁴ "Draft & Verify" skips this problem differently. It drafts by "selectively skipping certain intermediate layers" of the target model itself, needing no second model at all. It states plainly that "this process ensures the final output remains identical to that produced by the unaltered LLM," and reports up to 1.99× measured on LLaMA-2 variants.⁵ Lookahead decoding drops the draft model entirely, in favor of an "exact, parallel decoding algorithm" built on n-gram extraction against the target alone. It reports up to 1.8× on MT-bench and up to 4× on multi-GPU code completion; the 4× figure is specific to multi-GPU scaling.⁶ All four keep the same rejection test this piece proved; they disagree only about where the cheap candidate comes from.

Acceptance rate has to beat cost, and often doesn't

Tokens per call is not the same as wall-clock speed. A draft that is accurate but expensive to run can cost more time than it saves. Cost ratio c is defined as one draft call's time over one target call's time. Leviathan et al.'s Theorem 3.8 gives the expected walltime improvement as E[\#\text{tokens}] / (\gamma c + 1).¹ Corollary 3.9 states the condition plainly: a speedup exists only when \alpha > c. This is not measured wall-clock time, since dictionary lookups carry no meaningful cost asymmetry. It is a computed projection, using the paper's own theorem and this run's own measured \alpha at w=0.9, γ=4. That projection gives 3.381× at an illustrative c=0.1 and 2.367× at c=0.25. At w=0, the same formula gives 2.080× and 1.456× respectively: a worse draft still helps, just less.

Real systems show the same condition biting in both directions. The PyTorch team's own gpt-fast benchmark states it directly. "When running CodeLlama-34B + CodeLlama-7B, we're able to obtain a 2x boost in tokens/s for generating code," but "when using Llama-7B + TinyLlama-1B, we're only able to obtain about a 1.3x boost in tokens/s." Both numbers come from the same codebase; a weaker agreement between draft and target cuts the win by roughly a third.⁷ vLLM's own benchmark shows the sign of the effect flipping with load. At low query rate, draft-model speculative decoding gets "up to a 1.5x speedup" on ShareGPT, and n-gram matching gets "up to 2.8x" on CNN/DailyMail. At higher request rates, the same post reports that "the overhead of speculative decoding can outweigh its benefits, leading to reduced performance," roughly a 1.4–1.8x slowdown on the same hardware and models.⁸ Autoregressive decoding is normally memory-bandwidth-bound. The GPU spends most of a token's time waiting on weights to arrive from memory, not on arithmetic, which is why a batched verification pass over several positions costs barely more than one. Under enough concurrent load, the GPU runs out of that idle compute; every draft attempt then competes for cycles the system would rather have spent serving other requests, and a rejected draft has bought nothing.

None of this is a flaw in the proof. Every run in this piece, real or toy, produced exactly p's own distribution, checked directly against the tables that define it. What acceptance rate and cost ratio govern is only how much of the memory-bandwidth idle time a given draft-target pair can actually spend. At w=0.9 agreement and γ=4 draft tokens, this build's own run produced 4.7336 tokens per target call, and a computed 3.381× walltime speedup at c=0.1. That sits just above the 2X-3X range Leviathan et al. measured on a real 11-billion-parameter model. Their own Corollary 3.9 names the reason: the win needs a draft cheap and accurate enough that \alpha clears c. Nothing about a bigger target model changes that arithmetic.