E7.3 · LESSON 3 OF 6

Rejection correction

Trace draft-and-verify decoding, explain distribution-preserving acceptance, and calculate a simplified speedup.

THE IDEA

Why the correction after a rejection matters

If the drafted token is rejected, a replacement must be chosen. Drawing it from the plain target p would double-count the tokens the draft already covers. The right choice is the residual: the positive part of (p − q), normalised. With it, the total probability of emitting each token comes out exactly equal to p.

token	target p	draft q	residual max(0, p−q) normalised	output with correction	output if resampling from p
0	0.5	0.25	1	0.5	0.375
1	0.3	0.5	0	0.3	0.375
2	0.2	0.25	0	0.2	0.25

P(emit i) = q_i·min(1, p_i/q_i) + P(reject)·residual_i, P(reject) = 0.25 with the residual the result equals p exactly. Resampling from p instead over-weights tokens the draft already covered (red).

The residual is the part of the target that the draft under-covered. Drawing from it after a rejection is exactly what fills the gap.

Target p = [0.5, 0.3, 0.2], draft q = [0.25, 0.5, 0.25]. Compare the corrected output with the naive one.

Predict first. Can simply resampling from the original target after every rejection change the intended distribution?

Chapter context

A cheap draft proposes several tokens and a target model verifies them with parallel work. The opportunity comes from amortizing an expensive target step, but drafting, verification and rejection recovery all cost time.

TRY IT YOURSELF

Part 1. Do it hands-on

Compute the residual and the total probability by hand, then build the exact step and prove by simulation that it preserves p while the naive variant does not.

token	target p	draft q	residual max(0, p−q) normalised	output with correction	output if resampling from p
0	0.5	0.25	1	0.5	0.375
1	0.3	0.5	0	0.3	0.375
2	0.2	0.25	0	0.2	0.25

The residual is the part of the target that the draft under-covered. Drawing from it after a rejection is exactly what fills the gap.

residual distribution:

P(emit token 0) with the correction P(emit token 0) when resampling from p

To do: After a rejection the next token is drawn from the residual: positive part of (p − q), normalised. With p = [0.5, 0.3, 0.2] and q = [0.25, 0.5, 0.25], p − q = [0.25, −0.2, −0.05]. Type the residual distribution (3 numbers)
To do: Probability that token 0 is emitted overall: accepted directly (q₀ · min(1, p₀/q₀) = 0.25) plus rejection probability (0.25) times the residual's chance of token 0. Type the total
To do: A WRONG shortcut: after a rejection, resample from the original target p. Then the chance of token 0 is 0.25 + 0.25 × 0.5. Type it

Part 2. Do it in code

residual(p, q): normalise max(0, p − q). If it is all zero (p == q) return p.
rejection_probability(p, q): 1 − Σ min(p, q).
exact_distribution(p, q): the probability of emitting each token in one speculative step: q·min(1, p/q) + rejection_probability · residual. (It should equal p.) Treat tokens with q = 0 as contributing 0 via the first term.
naive_distribution(p, q): the WRONG variant that resamples from p itself after a rejection.
speculative_step(p, q, rng): one real step: draw a token from q, accept with min(1, p/q); on rejection draw from the residual. Return the emitted token.
empirical(p, q, trials, seed, method): frequencies of the emitted tokens over many trials; method is "exact" or "naive".

import numpy as np

def residual(p, q):
    pass

def rejection_probability(p, q):
    pass

def exact_distribution(p, q):
    pass

def naive_distribution(p, q):
    pass

def speculative_step(p, q, rng):
    pass

def empirical(p, q, trials, seed, method="exact"):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Can simply resampling from the original target after every rejection change the intended distribution?

ANoBYes

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

Under constant independent acceptance a and k draft proposals:

E[N]=\sum_{i=0}^{k}a^i=\frac{1-a^{k+1}}{1-a}\quad(a\ne1).

At a=1 the sum is k+1. A simple speed estimate is $E[N]/(t_{verify}+k t_{draft})$ . Exact samplers use conditional probabilities; this constant-a model is not their full correctness proof.

Worked example

For a=0.8,k=4, expected produced tokens are 1+0.8+0.64+0.512+0.4096=3.3616. If target-only generation takes 10 ms/token, verification costs 12 ms and four draft steps cost 1 ms each, speculative throughput is 3.3616/0.016=210.1 tokens/s versus 100 tokens/s baseline: an illustrative 2.101× speedup.

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import isclose
a,k=.8,4
expected=sum(a**i for i in range(k+1))
assert isclose(expected,3.3616)
assert isclose(expected/.016,210.1)
assert isclose((expected/.016)/(1/.010),2.101)
assert sum(1**i for i in range(k+1)) == 5

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Expected accepted draft tokens alone omit the final target token in this model. a=1 needs the finite-sum limit. A correctness-preserving acceptance rule is more than an agreement heuristic.