Speculative Decoding: Sampling After Rejection
Speculative decoding lets a draft model propose tokens for a target model to verify. Preserving the target distribution requires more than having the larger model check the draft: both the acceptance probability and the rejection fallback must be correct. With the usual acceptance rule, drawing directly from the full target distribution after rejection introduces bias.
Contents
- Define what distribution preservation means
- Count accepted probability mass at a fixed prefix
- Three tokens expose an incorrect fallback
- Enumerate exactly before sampling repeatedly
- Handle boundary cases explicitly
- Keep only the consecutively accepted draft prefix
- Use the distribution after the actual sampling rules
- Matching the most likely token can still mean low acceptance
- Acceptance rate is not an end-to-end speed measurement
- Sources and backfill note
A three-token example makes the discrepancy exactly enumerable. 阅读中文版.
Define what distribution preservation means
This article examines the exact sampling method in Fast Inference from Transformers via Speculative Decoding, published at ICML 2023. It is a conference-published foundational method predating this slot, not a new announcement.
The guarantee concerns the final output distribution under the algorithm’s conditions. It does not require token-for-token identical text for a given seed across different execution paths, nor does it say the draft’s uncorrected distribution equals the target. The sampling mechanism is what must preserve the specified distribution.
Count accepted probability mass at a fixed prefix
Consider the next token at one fixed prefix. Let p be the normalized target distribution and q the normalized draft distribution on the same token space. Propose x from q, then accept with probability min(1,p(x)/q(x)). A proposed candidate necessarily has q(x)>0, so its ratio is defined.
The probability of emitting x through acceptance is q(x) times that acceptance probability: min(p(x),q(x)). Sum these masses to obtain total acceptance probability A. The target mass still missing at x is p(x)−min(p(x),q(x))=max(p(x)−q(x),0).
Rejection happens with probability 1−A. Divide the missing mass by that quantity to form the conditional fallback distribution. Acceptance covers the overlap and rejection fills the deficit. Their unconditional contributions then add to p.
Propose from q
Acceptance path
Keep min(p,q)
Rejection path
Fill max(p−q,0)
Add both contributions
Recover target mass p
Three tokens expose an incorrect fallback
Choose p=(0.2,0.5,0.3) and q=(0.5,0.3,0.2). Acceptance contributes (0.2,0.3,0.2), totaling 0.7. The remaining mass is (0,0.2,0.1), and rejection probability is 0.3. The normalized fallback must therefore be (0,2/3,1/3).
Instead falling back to the full p gives (0.2,0.3,0.2)+0.3×(0.2,0.5,0.3)=(0.26,0.45,0.29). This is not p. It duplicates mass already covered by acceptance while failing to fill the actual deficit. The numbers are constructed for explanation, not measured token frequencies.
Target p
0.20 / 0.50 / 0.30
Accepted mass
0.20 / 0.30 / 0.20
Correct added mass
0.00 / 0.20 / 0.10
Wrong final distribution
0.26 / 0.45 / 0.29
Enumerate exactly before sampling repeatedly
This small discrete case allows direct enumeration of all probability paths. Python Fraction represents the decimal examples as exact rational numbers, preventing rounding from being mistaken for distribution bias.
# Exact enumeration for one fixed prefix, not a decoder or benchmark.
from fractions import Fraction as F
def corrected(p,q):
if len(p) != len(q) or not p:
raise ValueError("matching nonempty vocabularies required")
if any(x < 0 for x in p+q) or sum(p)!=1 or sum(q)!=1:
raise ValueError("normalized nonnegative distributions required")
kept = [min(a,b) for a,b in zip(p,q)]
rejected = 1-sum(kept)
if rejected == 0:
return kept, None, kept
residual = [max(F(0), a-b)/rejected for a,b in zip(p,q)]
output = [a+rejected*r for a,r in zip(kept,residual)]
return kept, residual, output
p = [F(2,10),F(5,10),F(3,10)]
q = [F(5,10),F(3,10),F(2,10)]
kept, residual, output = corrected(p,q)
assert kept == [F(2,10),F(3,10),F(2,10)]
assert residual == [F(0),F(2,3),F(1,3)]
assert output == p
reject = 1-sum(kept)
wrong = [a+reject*b for a,b in zip(kept,p)]
assert wrong == [F(26,100),F(45,100),F(29,100)]
assert wrong != p
assert corrected(p,p)[1] is None
assert corrected(p,[F(1),F(0),F(0)])[2] == p
assert corrected([F(0),F(1)],[F(1),F(0)])[2] == [F(0),F(1)]
print("exact output:", [str(x) for x in output])
print("incorrect fallback:", [str(x) for x in wrong])
The code was executed locally. It checks equality with p, the biased full-p fallback, identical distributions, zero draft probability for some tokens, and disjoint support. It enumerates one fixed prefix; it does not run a language model or measure decoding speed.
Handle boundary cases explicitly
When p=q, A=1 and rejection never occurs. Normalizing the residual would divide by zero. Skip that branch rather than inserting an arbitrary epsilon and sampling anyway. The example returns None to express that no fallback distribution is needed.
If q assigns zero probability to a token with positive p, the draft never proposes it, but the residual can restore its missing mass. Constructing the final distribution only over proposed token identities can therefore be wrong. Any restricted implementation needs a correctness definition that accounts for its actual truncation.
Floating-point implementations also need normalization, non-negativity, and tiny-residual checks. Silently treating a small rejection probability as zero while continuing to claim exact preservation is not justified by the derivation. Approximation can be useful, but its changed guarantee should be stated.
Keep only the consecutively accepted draft prefix
For multiple draft tokens, each p and q is conditioned on the corresponding preceding candidate prefix. A target verification computation can evaluate several positions together, but those distributions are not independent objects that may be freely recombined.
After the first rejection and a replacement token, later draft tokens were generated under a prefix that is no longer the emitted prefix. Retain only the consecutively accepted portion, correct the first rejected position, and continue from the updated actual sequence. If every proposal is accepted, the original algorithm draws an additional token from the appropriate target conditional distribution.
Cached state must match the retained prefix too. Correcting text while leaving rejected-position state can make subsequent computation inconsistent with the current sequence. The one-step enumeration does not validate state management; integration tests must cover it separately.
Use the distribution after the actual sampling rules
If temperature, top-k, or top-p is applied, p must mean the distribution actually sampled after those rules, not unprocessed logits. Likewise q must describe the mechanism that truly proposed the candidate. Substituting probabilities from incompatible stages breaks the argument.
Token events must also align. Different tokenizers can encode the same text as different sequences. Equal array indices alone do not establish equal events. Methods supporting different tokenizers need additional alignment machinery; the shared-vocabulary proof here does not cover them.
Different text after fixing a seed does not by itself establish distribution error. Similar-looking output from a few runs does not prove correctness either. Check exact small cases first, investigate bias with adequate statistical experiments second, and evaluate task quality separately.
Matching the most likely token can still mean low acceptance
Take p=(0.5,0.3,0.2) and q=(0.9,0.05,0.05). Both have the same most likely token, yet acceptance under this stochastic rule is only 0.5+0.05+0.05=0.6. Comparing argmax agreement discards the rest of the probability mass.
This constructed fixed-prefix number is not a model-wide acceptance estimate. Evaluate draft acceptance on the relevant task and generation strategy, using the corresponding prefixes. Greedy decoding has its own target rule and should not be conflated with stochastic acceptance here.
Acceptance rate is not an end-to-end speed measurement
Draft generation and verification both cost time, and longer proposals can waste computation after rejection. Acceptance measures reuse opportunity, not wall-clock improvement by itself. Hardware, batching, output length, and cache overhead belong in any performance comparison.
Validate the one-step correction first, then multi-step prefix and cache behavior, then compare end-to-end time under the same target model and generation strategy. The site’s LLM backpressure article covers service concurrency. Single-request decoding gains, service throughput, and tail latency should be observed separately.
Correct probabilities, correct state, and actual performance are three different pieces of evidence. If this exact example helps you avoid a supposedly lossless sampling bug, voluntary support below is welcome.
Sources and backfill note
The ICML 2023 proceedings page and Algorithm 1 and Appendix A.1. Diagrams, the three-token values, and rational enumeration are original teaching content without model or GPU benchmarks. Backfilled on the evening of October 15; the article date retains the planned 19:00 Beijing slot.
Support
If the explanations or examples helped you solve a problem or save time, you’re welcome to support this site with a voluntary tip.
Click a QR code to enlarge. More options: support page。


