Speculative decoding: watch the rejection sampling

Everyone cites the speedup number. Here is the part nobody shows: the per-token accept/reject math that makes it lossless.

GitHub Live GPU-free Space

1 · The verification step, live

A cheap draft model q proposes a token; the expensive target model p verifies it. Below is one vocabulary position with toy 8-token distributions. Drag the bars to reshape either distribution, or pick a preset.

draft q (proposes) target p (verifies) — drag either series to edit; both stay normalized
α = Σ min(p,q) · expected acceptance total variation ½Σ|p−q| = 1−α empirical acceptance (0 trials)

One verification, step by step

Press “Propose & verify one token”.

Why it’s lossless

Tokens that come out of accept-or-resample are distributed exactly like the target’s p — the whole point of the scheme. Run trials and watch the output histogram converge onto the p markers.

output via acceptance output via resample target p (goal)

2 · From acceptance rate to speedup

With acceptance rate α and γ drafted tokens per round, one target forward pass yields E[tokens] = (1 − αγ+1) / (1 − α). The curve below tracks the α of your distributions above — reshape them and watch the ceiling move. High α → longer accepted runs → more tokens per expensive call.

3 · The math, in four lines

1. draft samples x ~ q(·)          (cheap, autoregressive, γ at a time)
2. accept x with prob min(1, p(x)/q(x))   (target scores all γ in ONE pass)
3. if rejected: resample x′ ~ norm(max(p − q, 0))
4. P(output = x) = q(x)·min(1, p(x)/q(x)) + P(reject)·residual(x) = p(x) ✓

Line 4 is the theorem (Leviathan et al. 2023, arXiv:2211.17192): acceptance keeps exactly the probability mass where q overlaps p, and the residual puts back precisely the mass q under-covered. Nothing is approximated — the output distribution is the target’s.

4 · Real models, real traces

The toy above is exact math on toy distributions. The next step is the same algorithm on actual model logits: distilgpt2 drafting for gpt2, quantized ONNX, executed entirely in your browser with transformers.js — no server, no GPU quota, no cold start. The repo also ships a Gradio app (KV-cached, with attention-trace diffing and an optional 70B cloud cross-check) to run locally.

or open the Hugging Face Space ↗

MIT licensed · source · built as a visual artifact for the speculative-decoding verification step.