Speculative Decoding From Scratch
Quick Answer / TL;DR
Generating text one token at a time is slow because each token needs its own forward pass, and a forward pass is dominated by reading the model's weights from memory, not compute. So the pass is mostly idle. Speculative decoding leverages that idle time: a smaller and faster model (called draft model) proposes several tokens, the larger and slower model (called target model) checks all of them in a single pass, and a rejection-sampling step decides how many to keep. The remarkable part is that the tokens it keeps are drawn from exactly the target model's distribution, not an approximation of it, so the output is unchanged. We demonstrate the core aspects of this idea and provide from scratch implementation in numpy.
LLM decoding basics
Let us start with the single step every language model runs at inference. Given the tokens produced so far, the model does one forward pass over all of its weights and outputs a probability distribution over the next token; we sample one token from that distribution, append it, and repeat. Each new token needs its own forward pass, and the passes are strictly sequential, because the input to the next pass includes the token we just produced. Generating a hundred tokens means a hundred passes, one after another.
It is natural to assume each pass is slow because the model does a lot of arithmetic, and at very large batch sizes that is true. But for a single request the bottleneck is elsewhere. Every pass has to read the model's parameters out of memory into the compute units, and for a large model that memory traffic dwarfs the arithmetic. The compute units spend most of their time idle, waiting for weights to arrive.
If a pass is mostly idle compute, then scoring several candidate tokens in one pass costs almost the same as scoring a single token, since the weights are moved only once. The catch is that we do not have candidate tokens yet, and generating them is the very sequential process we are trying to avoid. Speculative decoding introduces a division of labor: a small, cheap draft model guesses the next few tokens, and the large target model, instead of generating them one at a time, spends a single pass checking all of them at once.
What is speculative decoding
The setup uses two models over the same vocabulary. The target model, whose distribution we write, is the one whose output we actually want. The draft model,, is a smaller, cheaper model that only approximates . A round has two phases. First the draft phase: generates a short continuation of tokens the ordinary autoregressive way small, sequential steps, cheap because is small. Then the verification phase: the target scores all of those positions in a single pass, and a decision rule walks left to right deciding which to keep. The figure below contrasts this with standard decoding and marks where the speedup comes from.
It is worth being precise about what the draft produces and why the target can check it in one pass, because this is the core of the method. The draft does not produce several independent candidates. It produces one sequence: tokens generated one after another, each conditioned on the prompt and the draft tokens before it, exactly as ordinary generation would. Call the prompt and the drafted tokens.
Because the draft tokens are already fixed, the target does not need to produce them, only to score them, and scoring every position of a known sequence is exactly what one transformer forward pass does. Feed the target the prompt followed by all draft tokens, and thanks to causal attention that single pass returns the next-token distribution at every position at once:
That is distributions from one pass: one to check each drafted token against, plus one for the position after the last token, which is why a fully accepted round can emit an extra “bonus” token. Generation is serial only because you cannot condition on a token you have not produced yet; verification sidesteps that, because here the tokens are given. What remains is the decision rule that compares each drafted token's target probability with its draft probability and decides to keep it or replace it.
If that rule were simply “keep the draft tokens the target also finds likely,” the output would drift away from , because we would be filtering the target's distribution through the draft's guesses. The contribution of Leviathan et al. and, independently, Chen et al. in 2023 was a decision rule that does not distort anything: a modified rejection-sampling scheme whose output is provably distributed exactly as . Before assembling the full loop, it is worth getting that primitive exactly right.
The core concept: accept or resample
Consider a single position. The draft proposes a token x sampled from its distribution q. The target gives us its own distribution p over the same vocabulary. We want a procedure that turns the draft sample into a sample from p. Accept with probability . If x is accepted, it is our output. If it is rejected, we discard it and instead sample a replacement from the residual distribution proportional to , renormalized to sum to one.
This is a variation on classical rejection sampling, where one draws from an easy proposal distribution and keeps each draw with a probability chosen so that the survivors follow a harder target. The twist is what happens on a rejection. Ordinary rejection sampling throws the draw away and tries again, but here we cannot afford to waste the expensive target pass, so we make the rejection productive: we resample once from the leftover mass , the part of the target the draft failed to cover. That single corrected draw is what lets one pass still return a valid token.
Where the draft proposes a token less often than the target would have (), the ratio exceeds one and we always accept: the draft is under-representing that token, so every proposal should be kept. Where the draft over-proposes a token (), we accept only a fraction of the time, trimming the excess. The mass that gets trimmed by rejection has to be redistributed somewhere, and it is redistributed precisely onto the tokens the draft under-proposed, which is what the residual captures. Here is the primitive in code, together with a Monte Carlo check that its output really is distributed as p. Interested readers may look into the original speculative sampling paper for a formal proof.
import numpy as np
def speculative_step(p, q, rng):
"""One token from target p, using draft q. Returns (token, accepted)."""
x = rng.choice(len(q), p=q) # draft proposes x ~ q
if rng.random() < min(1.0, p[x] / q[x]): # accept w.p. min(1, p/q)
return x, True
resid = np.maximum(p - q, 0.0) # else resample from (p - q)+
return rng.choice(len(p), p=resid / resid.sum()), False
rng = np.random.default_rng(0)
p = np.array([0.30, 0.25, 0.20, 0.15, 0.07, 0.03]) # target
q = np.array([0.10, 0.10, 0.30, 0.30, 0.15, 0.05]) # draft (deliberately off)
N = 2_000_000
counts = np.zeros(len(p)); accepts = 0
for _ in range(N):
tok, acc = speculative_step(p, q, rng)
counts[tok] += 1; accepts += acc
emp = counts / N
print("target p :", np.array2string(p, precision=4, floatmode='fixed'))
print("empirical :", np.array2string(emp, precision=4, floatmode='fixed'))
print("max error :", f"{np.max(np.abs(emp - p)):.5f} (sampling SE ~ 0.00032)")
print("accept rate:", f"{accepts/N:.4f}")
print("sum min(p,q):", f"{np.minimum(p,q).sum():.4f}")
print("1 - TV(p,q):", f"{1 - 0.5*np.abs(p-q).sum():.4f}")target p : [0.3000 0.2500 0.2000 0.1500 0.0700 0.0300]
empirical : [0.2998 0.2503 0.1998 0.1501 0.0699 0.0300]
max error : 0.00034 (sampling SE ~ 0.00032)
accept rate: 0.6499
sum min(p,q): 0.6500
1 - TV(p,q): 0.6500The draft here is badly matched to the target on purpose. Even so, the empirical distribution of the output matches the target to within the sampling error of two million draws. That is the sense in which speculative decoding is lossless: it is exact within numerical sampling errors.
Why it is exact
The formal proof is short and is given in the original paper. To get the intuition behind it, one may think about how much probability the whole procedure ends up placing on some token x. There are two ways x can come out: the draft proposed it and we kept it, or the draft proposed something else that was rejected and the resample landed on x. On the first path the draft offers x as often as q likes, but we only ever keep it up to the target's appetite, so the mass that survives is the smaller of the two, , the part of x that the draft and target already agree on.
The second path repairs the disagreement. Each rejection frees up some probability, and by construction the resample can only land on tokens the draft under-proposed, which are exactly the ones still short of the target. So the freed-up mass is poured back precisely where it was missing. Add the part they agreed on to the repaired shortfall and, token by token, you are left with the target's own probability:
Notice that the draft distribution q has dropped out of the result entirely. That is the whole point: a poor draft changes how often we accept, and therefore how fast we go, but it does not change what we output. Good draft or bad, the tokens you get are the target's.
The acceptance rate is a distance
If the draft does not affect correctness, it does however affects speed. The acceptance rate, which we call α, is the probability that a proposed token is accepted, averaged over what the draft proposes. That average is , and by the same simplification as before this collapses to , the overlapping area under the two distributions.
The overlap is exactly one minus the total-variation distance between the two distributions, since total variation is half the summed absolute difference and the overlap is what remains after you remove it. In the run above, the acceptance rate came out to 0.6499, matching both and at 0.6500 to within sampling noise. This is a useful way to think about the whole method: the acceptance rate is a similarity score between the draft and the target. A draft that agrees with the target eighty percent of the time in this distributional sense buys you far more than one that agrees sixty percent of the time.
Drafting several tokens at once
A single accepted token is not yet a speedup, because we spent a target pass to get it, the same as ordinary decoding. The win comes from drafting γ tokens in a row and verifying them together. The target scores all γ positions in one pass, and we walk left to right applying the accept test at each. As long as tokens keep being accepted, we advance for free. At the first rejection we resample that one position from its residual and stop the round, discarding any draft tokens past it, because they were conditioned on a token we just changed. And if all γ tokens are accepted, the target's parallel pass has already given us the distribution for the position after them, so we sample one bonus token from it directly. A round therefore yields somewhere between one and γ + 1 tokens, always for the price of a single target pass.
Under the simplifying assumption that each position is accepted independently with probability α, the number of accepted tokens before the first failure is geometric, and adding the guaranteed extra token gives a clean closed form for the expected yield of a round.
We can also check that formula through a simple simulation.
def expected_tokens(alpha, gamma):
return (1 - alpha**(gamma + 1)) / (1 - alpha)
def sim_tokens(alpha, gamma, rng, trials=300_000):
total = 0
for _ in range(trials):
acc = 0
for _ in range(gamma): # walk the gamma drafted tokens
if rng.random() < alpha: acc += 1 # accept (iid idealization)
else: break # first reject ends the round
total += acc + 1 # + resampled-or-bonus token
return total / trials
rng = np.random.default_rng(1)
print("alpha gamma formula simulated")
for alpha in [0.5, 0.7, 0.9]:
for gamma in [1, 4, 8]:
f = expected_tokens(alpha, gamma)
s = sim_tokens(alpha, gamma, rng)
print(f" {alpha:.1f} {gamma:2d} {f:.4f} {s:.4f}")alpha gamma formula simulated
0.5 1 1.5000 1.5015
0.5 4 1.9375 1.9371
0.5 8 1.9961 1.9950
0.7 1 1.7000 1.7002
0.7 4 2.7731 2.7743
0.7 8 3.1988 3.1979
0.9 1 1.9000 1.8999
0.9 4 4.0951 4.0938
0.9 8 6.1258 6.1319The formula and the simulation agree, and the table already shows the shape of the returns. At a low acceptance rate the yield saturates quickly: with α = 0.5, going from a draft length of four to eight barely moves the expected tokens from 1.94 to 2.00, because rejections almost always cut the round short long before the eighth token. At a high acceptance rate the story is different, and longer drafts keep paying off, with α = 0.9 reaching more than six tokens per round at γ = 8. As γ grows without bound the expected yield approaches , which is the ceiling that derives from the acceptance rate.
From tokens per round to wall-clock speedup
Tokens per round is not the same as speedup, because a round is not free: it costs γ cheap draft passes plus one target pass. To turn yield into wall-clock time we need the cost ratio between the two models, , the time for one draft forward pass as a fraction of one target forward pass. A round then costs target-equivalents and produces tokens, and standard decoding costs one target pass per token, so the speedup is the ratio of the two.
This is where the design tension becomes visible. A longer draft raises the numerator, capturing more tokens per round, but it also raises the denominator, because every drafted token costs a draft pass whether or not it survives, and tokens past a rejection are wasted. Beyond some point the extra draft passes cost more than the tokens they buy, and the speedup turns over. The optimal draft length is therefore finite, and it depends jointly on the acceptance rate and the cost ratio. The chart below plots the speedup against draft length for a range of acceptance rates, at a draft that costs a tenth of a target pass, with a dot marking each curve's optimum.
The chart says almost everything about when the method is worth it. The acceptance rate controls both how high the curve reaches and where its peak sits: a weak draft at α = 0.5 tops out below 1.5× and wants only a two-token draft, while a strong draft at α = 0.9 climbs past 3.4× and keeps rewarding drafts as long as ten tokens. Pushing the draft longer than the optimum is not neutral; it actively gives speedup back, and a long enough draft on a weak model slips below the parity line and runs slower than plain decoding.
def speedup(alpha, gamma, c):
# c = t_draft / t_target (cost of one draft pass vs one target pass)
return (1 - alpha**(gamma + 1)) / ((1 - alpha) * (gamma * c + 1))
c = 0.1
print(f"draft cost c = {c} (draft is 10% of a target pass)\n")
print("alpha best gamma speedup ceiling 1/(1-alpha)")
for alpha in [0.5, 0.6, 0.7, 0.8, 0.9]:
gs = np.arange(1, 21)
vals = [speedup(alpha, g, c) for g in gs]
gstar = gs[int(np.argmax(vals))]
print(f" {alpha:.1f} {gstar:2d} {max(vals):.2f}x {1/(1-alpha):.1f}x")draft cost c = 0.1 (draft is 10% of a target pass)
alpha best gamma speedup ceiling 1/(1-alpha)
0.5 2 1.46x 2.0x
0.6 3 1.67x 2.5x
0.7 4 1.98x 3.3x
0.8 6 2.47x 5.0x
0.9 10 3.43x 10.0xThe gap between the realized speedup and the ceiling is the tax paid for drafting overhead and for the tokens wasted past each rejection. It narrows as the draft gets cheaper: a smaller c lets you afford longer drafts and pushes the optimum higher, which is why in practice a great deal of effort goes into making the draft model both cheap and well-aligned, since those are the two levers, c and α, that the whole formula rests on.
Putting it together: generating a sequence
The single-token primitive and the round structure are enough to build the full generation loop, as described in the original Speculative Decoding paper. To keep the code self-contained and easily tractable, we will stand in for the two models with explicit conditional distributions rather than trained networks: each “model” is a matrix whose row t gives the next-token distribution when the last token is t. The target is one such matrix, and the draft is a blurred, cheaper copy of it. This is a faithful stand-in because the speculative machinery only ever touches the models through their next-token distributions; a real transformer would supply the same rows from a forward pass. The test that matters is whether the sequences this loop produces are distributed exactly like sequences from the target alone.
# Two "models" as last-token conditional distributions (stochastic matrices).
# Row t of P is target P(next | last=t); Q is the cheaper draft.
V = 5
rng = np.random.default_rng(2)
def stochastic(temp):
e = np.exp(rng.normal(0, 1, (V, V)) / temp); return e / e.sum(1, keepdims=True)
P = stochastic(0.7)
Q = 0.6 * P + 0.4 * np.ones((V, V)) / V; Q /= Q.sum(1, keepdims=True) # blurred target
def resample(p, q):
r = np.maximum(p - q, 0.0); return rng.choice(V, p=r / r.sum())
def target_only(start, L):
seq = [start]
for _ in range(L): seq.append(rng.choice(V, p=P[seq[-1]]))
return seq[1:]
def speculative(start, L, gamma):
seq, calls = [start], 0
while len(seq) - 1 < L:
last = seq[-1]; ctx = last; draft = []
for _ in range(gamma): # draft gamma tokens from Q
nxt = rng.choice(V, p=Q[ctx]); draft.append(nxt); ctx = nxt
calls += 1; ctx = last # ONE parallel target pass
for j in range(gamma):
p, q = P[ctx], Q[ctx]; x = draft[j]
if rng.random() < min(1.0, p[x] / q[x]):
seq.append(x); ctx = x
if len(seq) - 1 >= L: break
else:
seq.append(resample(p, q)); break
else:
if len(seq) - 1 < L: seq.append(rng.choice(V, p=P[ctx])) # bonus token
return seq[1:L + 1], calls
L, gamma, N = 64, 4, 60_000
ft = np.zeros(V); fs = np.zeros(V); total_calls = 0
for _ in range(N):
for tok in target_only(0, L): ft[tok] += 1
s, c = speculative(0, L, gamma); total_calls += c
for tok in s: fs[tok] += 1
ft /= N * L; fs /= N * L
print("token freq, target-only :", np.array2string(ft, precision=4, floatmode='fixed'))
print("token freq, speculative :", np.array2string(fs, precision=4, floatmode='fixed'))
print("max abs diff :", f"{np.max(np.abs(ft - fs)):.5f} (SE ~ 0.00022)")
print(f"tokens per target call : {(N*L)/total_calls:.2f}")
print(f"target calls saved : {(1 - total_calls/(N*L))*100:.0f}%")token freq, target-only : [0.3022 0.1299 0.1285 0.1209 0.3185]
token freq, speculative : [0.3024 0.1298 0.1287 0.1207 0.3184]
max abs diff : 0.00029 (SE ~ 0.00022)
tokens per target call : 3.54
target calls saved : 72%Over sixty thousand generated sequences the two token distributions are identical to within sampling error, so the losslessness we proved for a single position survives all the way through a full autoregressive loop with resampling and bonus tokens threaded through it. And it does so while calling the expensive target model on average once for every 3.54 tokens produced, a cut of nearly three quarters in target passes. That number is a little below what the single-rate formula would predict at this draft's average acceptance rate because the acceptance rate here varies from one context to the next, and a round is only as long as its weakest link, so the harder contexts drag the realized yield down. It is a good reminder that the α in the formulas is an average standing in for a distribution, and the tail of that distribution matters.
What moves the acceptance rate in practice
Since the acceptance rate is the total-variation overlap between draft and target, improving it means making the draft agree with the target on the next-token distribution. The most common approach is to pair a large target with a small model from the same family trained on similar data, so that the two agree on the many easy, predictable tokens, the whitespace, the closing brackets, the obvious continuations, which is where most of the acceptance comes from. A related line of work removes the separate draft model entirely and lets the target predict several tokens ahead of itself through added lightweight heads, as in the Medusa and EAGLE families, which sidesteps the alignment problem by construction because the drafts come from the target's own features. In all of these the goal is the same quantity we derived: raise the overlap, and the speedup ceiling rises with it.
Because acceptance depends on the draft's per-context agreement, speculative decoding helps most on the predictable stretches of a generation and least on the genuinely uncertain ones, where the draft and target disagree and rejections are frequent. The wall-clock numbers in the original papers, roughly two to three times faster on real models, reflect that mix. On a laptop with two toy models the arithmetic advantage that makes verification nearly free does not exist, so a from-scratch implementation like this one is only the place to confirm the losslessness and the acceptance-rate maths.
Common questions
Does speculative decoding change the model's output or lower its quality?
No. The accept-and-resample rule is constructed so that the token it emits is drawn from exactly the target model's distribution, which we both proved in one line and confirmed by simulation. A weaker draft only lowers the acceptance rate and therefore the speed; it cannot shift the output distribution. In that sense the method is not an approximation with a quality knob, it is an exact reformulation of sampling from the target, up to floating-point precision.
What acceptance rate do I need for it to be worth it?
Enough that the peak of the speedup curve clears 1× by a comfortable margin, which in practice means the draft has to be both reasonably aligned and genuinely cheap. At a draft costing a tenth of the target, an acceptance rate around 0.6 already yields a real speedup, and the returns grow quickly above that because the ceiling steepens. If the draft is expensive relative to the target, the same acceptance rate buys much less, since the cost ratio c enters the denominator directly.
How long should the draft be?
There is a finite optimum, and it grows with the acceptance rate and shrinks as the draft gets more expensive. The closed form makes it a one-line search, as in the table above, where the best draft length runs from two tokens at α = 0.5 up to ten at α = 0.9 for a draft costing a tenth of the target. Guessing longer than the optimum is the more common mistake, because each extra drafted token is paid for whether or not it is accepted, and everything past the first rejection is discarded.
References
- Leviathan, Kalman & Matias, “Fast Inference from Transformers via Speculative Decoding” (ICML 2023) — introduces the method, the acceptance-rate analysis, and the expected-tokens and speedup formulas used here.
- Chen, Borgeaud, Irving, Lespiau, Sifre & Jumper, “Accelerating Large Language Model Decoding with Speculative Sampling” (2023) — the concurrent DeepMind formulation, with the modified rejection-sampling scheme that preserves the target distribution.
- Cai et al., “Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads” (2024) — drafting from added heads on the target itself.
- Li et al., “EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty” (2024) — feature-level drafting from the target's own representations.
Related Cookbooks
RAG vs SFT: When to Use Which | SR Cookbooks
A technical breakdown of when to use Retrieval-Augmented Generation (knowledge) versus Supervised Fine-Tuning (behavior) in enterprise AI pipelines.
Inspecting What a Tiny Transformer Actually Learned | SR Cookbooks
A technical guide to probing a character-level PyTorch Transformer. Learn how to measure rule acquisition, test generalization, and ablate attention heads.
Implementing a KV Cache From Scratch: Pure PyTorch | SR Cookbooks
Learn the mechanics of autoregressive LLM optimization. Implement a transformer Key-Value (KV) cache from scratch in PyTorch to massively speed up decoding.