Speculative decoding: guess, then check

Leviathan, Kalman & Matias, 2022

Free speed with no change to the output, which sounds like a lie until you print the distribution yourself.

Download speculative_decoding_from_scratch.py — everything below runs from it. NumPy only.

The KV-cache page ended on an uncomfortable finding: decoding is memory-bound. Every token you generate drags all twelve-and-a-half gigabytes of weights across the memory bus to do about five billion multiply-accumulates, and the chip finishes that arithmetic while it is still waiting for the next chunk of weights. The arithmetic units idle. The compute is, for practical purposes, free.

This paper spends it. That's the entire idea. I ignored it for a year because "speculative" sounded like it meant "approximate", and I don't ship approximate. It doesn't. The output distribution is provably, exactly the big model's — and that guarantee is four lines of code you can verify by counting.

It's branch prediction

Your CPU does this on every conditional. It doesn't wait to know whether the branch is taken. It guesses, runs ahead speculatively, and when the real answer lands it either keeps the work or throws it away. You never get a wrong result. You only ever get wasted work. 

# what decoding does today: one expensive call, one token
for _ in range(n):
    token = big_model(context)     # 12.5 GiB of weights read, 1 token out
    context.append(token)

Now guess ahead. A small cheap model drafts k tokens. Then the big model checks all of them in one call, because a forward pass over k+1 positions costs about the same as a forward pass over one — you read the weights once either way. 

# speculative: k cheap calls, then ONE expensive call that checks them all
guesses = [small_model(context + guesses) for _ in range(k)]   # cheap, serial
p = big_model(context + guesses)      # k+1 distributions, one pass
keep = accept(guesses, p)             # a coin flip per token, then STOP
context.extend(keep)                  # 1 to k+1 tokens for one big call

The clever part is not the drafting. Anyone would think of the drafting. The clever part is accept, which is rigged so that the tokens coming out the other end are distributed exactly as if the big model had generated them alone. Not approximately. Exactly.

TermMeansProgrammer's analogy
Target modelThe big one whose answers you promised your usersThe real branch outcome
Draft modelA small cheap model, same tokeniserThe branch predictor
p, qTarget's and draft's probability for the same next tokenTwo dicts you compare key by key
AcceptKeep the guess, with a probability that corrects for the draft's biasPrediction hit — commit the work
RejectThrow away this guess and everything after itPipeline flush

The one formula

Two lines, and the second is the one everybody forgets.

accept x with probability min(1, p(x) / q(x))
on rejection, sample from p′(x) = norm( max(0, p(x) − q(x)) )

Read it as a pipeline:

StepCostWhat it does
draft k tokensk cheap serial callsGuess ahead. Each guess conditioned on the previous guess
one target call1 expensive callReturns k+1 distributions — one per prefix, all at once
min(1, p/q)a coin flipThe draft over-proposed x? Keep only p/q of them. Under-proposed? Keep all
norm(max(0, p−q))one sampleThe residual. Exactly the mass the accept step failed to deliver
stopfreeEverything after a rejection was conditioned on a token that no longer exists

Accept plus residual add back up to p. That's the proof — an identity you can check with a calculator, not a limit argument.

The rule on one token

My models here are not neural nets. They are 6×6 probability tables — a row per current token, the row being the next-token distribution. That's deliberate. The paper's claim is about distributions, and a table is the only model whose true distribution I can print and compare against. A 70B target and a 6×6 table obey the same acceptance rule.

Don't take my word for it. Here's the actual output:

  context token          : 'a'
  target  p(.|ctx)       : [0.0943 0.31   0.0908 0.0015 0.3821 0.1213]
  draft   q(.|ctx)       : [0.0823 0.2638 0.0777 0.0013 0.4479 0.127 ]

  draft proposes 'b'
    q[b] = 0.2638, p[b] = 0.3100
    accept with min(1, p/q) = min(1, 0.3100/0.2638) = 1.0000 ✓
    draft under-proposes it -> always accept

  draft proposes 'e'
    q[e] = 0.4479, p[e] = 0.3821
    accept with min(1, p/q) = min(1, 0.3821/0.4479) = 0.8532
    draft over-proposes it -> reject 14.7% of the time

  on rejection, resample from residual = norm(max(0, p - q)):
    p - q          : [ 0.0119  0.0461  0.0131  0.0002 -0.0657 -0.0057]
    residual       : [0.1673 0.6457 0.1838 0.0032 0.     0.    ]

Follow it by hand. The draft says "e" too often, so 14.7% of the time we throw that guess away. But then we can't just resample from p — those rejections are not spread evenly across the vocabulary, they're concentrated on exactly the tokens the draft over-proposed. So we sample from what's missing: p minus q, negatives clipped to zero, renormalised. Notice "e" and "f" have probability 0 in the residual. The draft already delivered enough of those.

The correctness guarantee, counted

This is the section that matters. Everything else is performance engineering; this is the reason you're allowed to turn it on in production. I run the whole scheme 60,000 times from a fixed context, keep the first two emitted tokens, and compare the empirical joint against the target model's true joint. Then I break the rule three ways and measure how far the distribution drifts.

The measure is total variation distance: half the sum of absolute differences. It's the largest probability two distributions can disagree about on any event. Zero means identical. 

  60000 rounds, k=4 drafted tokens per round, first two tokens kept

  first-token marginal, target vs speculative decoding:
    token      target   speculative    diff
      a       0.0943      0.0905      -0.0038
      b       0.3100      0.3155      +0.0056
      c       0.0908      0.0920      +0.0012
      d       0.0015      0.0015      -0.0001
      e       0.3821      0.3801      -0.0020
      f       0.1213      0.1204      -0.0010

  total variation distance on the two-token joint:
    correct rule                     : 0.0079 ✓
    BROKEN: accept on argmax agreement     : 0.2735
    BROKEN: resample from p, not residual  : 0.0349
    BROKEN: keep drafts after a rejection   : 0.0333

  sampling-error floor (target model sampled directly): 0.0059

Read the last line first. If I sample 60,000 pairs from the target model itself and compare that to the target model, I get 0.0059 — that's pure counting noise, the price of a finite sample. Speculative decoding with the correct rule sits at 0.0079. Same floor. It is not "close to" the target model; it is the target model, and the residue is my sample size.

The three broken rules sit four to thirty times above that floor. That is a different model answering your users, and no eval you own will tell you which one it was.

The greedy shortcut is the tempting one. "Just accept the draft whenever the target agrees on the argmax" — it's faster, it needs no random number, and it feels safe. It scores 0.2735. You have silently replaced sampling with something close to greedy decoding, which is a different product.

Acceptance rate is the whole game

Two numbers decide whether this pays: how often the draft is accepted, and how much the draft costs. Nothing else. I sweep the draft from a perfect copy of the target down to junk, with the draft priced at 10% of a target call per token, drafting 4 tokens a round: 

  k=4 drafted tokens per round, draft cost c=0.10 of a target call

  draft quality      accept rate   tokens/target call   speedup
  alpha=1.00 (perfect)   1.0000        5.0000            3.571x ✓
  alpha=0.95             0.9204        4.6814            3.344x
  alpha=0.85             0.7769        4.1076            2.934x
  alpha=0.60             0.4964        2.9856            2.133x
  alpha=0.30             0.2698        2.0792            1.485x
  alpha=0.00 (junk)      0.1316        1.5265            1.090x

  temperature mismatch (same model, wrong sampling temperature):
    draft T=1.0 vs target T=1.0 : accept 1.0000, tokens/call 5.0000, 3.571x
    draft T=1.3 vs target T=1.0 : accept 0.8625, tokens/call 4.4501, 3.179x
    draft T=2.0 vs target T=1.0 : accept 0.6658, tokens/call 3.6631, 2.617x
    draft T=0.5 vs target T=1.0 : accept 0.6743, tokens/call 3.6970, 2.641x

A perfect draft never gets rejected, so you emit all 4 drafted tokens plus the free bonus token the target already computed at the last position: 5 tokens for one target call. A junk draft emits 1.5 and, after paying for the drafting, buys you 9%. These are my toy model's numbers, not the paper's; the shape is what transfers, not the magnitudes.

The temperature rows are the ones I'd put on an incident postmortem. Nothing is wrong with either model. The draft is the target, exactly, just sampled at a different temperature — and acceptance falls from 1.00 to 0.67. Acceptance rate is a property of the pair, not of either model. It's also why the draft must share the target's tokeniser: without that, p and q aren't distributions over the same set and the ratio p/q is meaningless.

How far ahead to guess

Draft one token and you've barely used the free compute. Draft sixteen and you're paying for fifteen guesses that get flushed the moment the first one is rejected. There is an optimum, and it moves with acceptance rate: 

  middling draft (alpha=0.55), draft cost c=0.10

   k   accept rate   tokens/target call   cost 1+k*c   speedup
   1      0.7078        1.7078           1.10        1.553x
   2      0.5997        2.1995           1.20        1.833x
   3      0.5194        2.5583           1.30        1.968x
   4      0.4494        2.7974           1.40        1.998x
   5      0.4002        3.0009           1.50        2.001x ✓
   6      0.3490        3.0937           1.60        1.934x
   8      0.2829        3.2634           1.80        1.813x
  12      0.1987        3.3845           2.20        1.538x
  16      0.1486        3.3774           2.60        1.299x

  best k here: 5 at 2.001x

Watch the two middle columns pull against each other. Tokens per call saturates — going from k=8 to k=16 buys 0.11 of a token, because the run almost never survives eight acceptances. Meanwhile the cost column keeps climbing linearly, 0.10 per extra draft, whether the token survives or not. That's the whole shape of the curve: a saturating numerator over a linear denominator.

The interview question is "why not draft 50?" Because acceptance compounds. If each token is accepted with probability r, the chance all k survive is rk. At r = 0.7 that's 3% by k = 10. You are buying lottery tickets at a fixed price.

When the draft model makes you slower

The cost model is one line. A round costs one target call plus k draft calls, where c is the draft's cost per token as a fraction of a target call. So speedup = tokens per round ÷ (1 + k·c), and you break even when c = (tokens per round − 1) / k. 

  middling draft, k=4, tokens per target call = 2.8064
  break-even draft cost c* = (tokens - 1)/k = 0.4516
  i.e. the draft model must cost less than 45.2% of the target per token.

  draft cost c    1 + k*c    speedup
      0.0100      1.0400     2.698x
      0.0500      1.2000     2.339x
      0.1000      1.4000     2.005x ✓
      0.4516      2.8064     1.000x  <-- break-even
      0.6000      3.4000     0.825x  SLOWER than no speculation
      1.0000      5.0000     0.561x  SLOWER than no speculation

  a 'draft' that is a quarter of the target (c=0.25) and only
  as good as our poor draft:
    tokens/call 1.5523, speedup 0.776x

That last block is the failure mode people actually ship. You reach for a bigger draft model because it gets accepted more often — and a draft at a quarter of the target's cost, with a poor acceptance rate, lands at 0.776x. You made inference slower and the output distribution is still perfectly correct, so nothing alerts. The only symptom is the p50 latency graph.

In practice c is a ratio of parameter counts, roughly: a 7B drafting for a 70B is around 0.1. A model drafting for itself is c = 1 and can never win.

Where the free compute comes from

Back to the KV-cache finding, but spent instead of merely observed. The weights you drag off memory for a forward pass are the same weights whether that pass covers one token or sixteen. Only the arithmetic scales. 

  Llama-2-7B-shaped, 512 tokens of context, fp16 weights
  weights dragged off memory per forward pass: 13,480,000,000 bytes (12.55 GiB)

   tokens in the pass        MACs        MACs per byte moved
         1             5,167,382,528              0.38
         2            10,334,765,056              0.77
         4            20,669,530,112              1.53
         8            41,339,060,224              3.07
        16            82,678,120,448              6.13 ✓

An accelerator needs roughly 100–300 MACs per byte of memory traffic to keep its arithmetic units busy. One-token decoding sits at 0.38 — off by nearly three orders of magnitude. Verifying sixteen tokens in the same pass moves the identical bytes, does sixteen times the arithmetic, and still doesn't saturate the chip. The k extra tokens are close to free.

This is why the trick exists at all. If decoding were compute-bound, verifying k tokens would cost k times as much and the whole scheme would collapse to zero gain. Speculative decoding is not a clever algorithm that happens to be fast; it is a direct exploitation of one hardware imbalance.

What came after

Chen et al. published the same acceptance rule independently a couple of months later (arXiv 2302.01318), applied to Chinchilla — worth knowing because half the blog posts you'll read cite one and half cite the other.

The obvious annoyance with all of this is that you now have to serve two models, keep them on the same tokeniser, and keep them in sync through every fine-tune. Later work removes the second model: bolt extra prediction heads onto the target model itself so it drafts its own continuations, Medusa being the well-known example. Same acceptance machinery, one set of weights. Different paper, different page.

The code

Pure NumPy. No torch, no autograd, no neural network at all — the models are explicit probability tables, so the distribution the paper promises is a thing you can print and subtract. Every function under 20 lines. Below is the acceptance rule and the verify loop; the drafting loop, the cost model and the six demos are in the file.

Download the full script — about 400 lines plus 6 runnable demos. Then: 

python3 speculative_decoding_from_scratch.py
# ---------------------------------------------------------------------------
# STAGE 2 -- THE WHOLE PAPER, part 1: the acceptance rule.
#
# The draft model proposed token x, and it proposed it with probability q[x].
# The target model would have said p[x]. Then:
#
#     accept x with probability   min(1, p[x] / q[x])
#
# q[x] > p[x] means the draft over-proposes x, so we keep only p/q of those.
# q[x] <= p[x] means the draft under-proposes x, so we keep every one.
# ---------------------------------------------------------------------------
def accept_prob(p, q, x):
    return min(1.0, p[x] / q[x])


# ---------------------------------------------------------------------------
# STAGE 3 -- THE WHOLE PAPER, part 2: the residual distribution.
#
# Accepting alone under-samples the tokens the draft was too shy about. On a
# rejection we do NOT resample from p -- that would double-count. We resample
# from the positive part of (p - q), normalised. That is exactly the mass the
# accept step failed to deliver, so the two steps sum back to p. This one line
# is the difference between "provably identical output" and "usually fine".
# ---------------------------------------------------------------------------
def residual(p, q):
    r = np.clip(p - q, 0.0, None)
    s = r.sum()
    if s <= 0:                     # degenerate: p is fully covered by q
        return p / p.sum()
    return r / s


# ---------------------------------------------------------------------------
# STAGE 5 -- verify all k drafts in ONE target call, then emit.
#
# The target sees the prefix plus all k drafted tokens at once, so it returns
# k+1 next-token distributions for the price of one forward pass (the same
# trick as training-time teacher forcing). Then walk them left to right:
#
#   accept  -> emit the drafted token, move on
#   reject  -> emit one token from residual(p, q) and STOP. Everything after
#              the rejection was conditioned on a token that no longer exists.
#   all k accepted -> you also get a free bonus token from the last
#              distribution, because the target already computed it.
#
# `rule` selects deliberate breakages; "correct" is the paper.
# ---------------------------------------------------------------------------
def verify_round(P, Q, ctx, drafted, rng, rule="correct"):
    prefixes = [ctx] + drafted          # k+1 positions, one target call
    out, n_accepted = [], 0
    for i, x in enumerate(drafted):
        p, q = P[prefixes[i]], Q[prefixes[i]]
        if rule == "greedy":
            ok = (x == int(np.argmax(p)))          # BREAK: argmax agreement
        else:
            ok = rng.random() < accept_prob(p, q, x)
        if ok:
            out.append(x)
            n_accepted += 1
            continue
        if rule == "no_resample":
            out.append(sample_from(p, rng.random()))          # BREAK
        else:
            out.append(sample_from(residual(p, q), rng.random()))
        if rule == "continue":
            out.extend(drafted[i + 1:])                       # BREAK
        return out, n_accepted
    out.append(sample_from(P[prefixes[-1]], rng.random()))     # bonus token
    return out, n_accepted

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — reverse-engineering from broken behaviour is your native mode. Every row below is a rule already wired into the script: pass it as rule=, or edit the function and rerun demo 2.

Break thisWatch for
Accept greedily: ok = (x == int(np.argmax(p)))Faster, and the output distribution is no longer the target's. TV distance goes from 0.0079 to 0.2735 against a 0.0059 noise floor.
In verify_round, replace residual(p, q) with plain pSubtle bias — nothing crashes, nothing looks wrong. TV distance 0.0349, four times the noise floor. The tokens the draft over-proposed stay over-represented.
Keep the remaining drafts after a rejection: out.extend(drafted[i+1:])TV 0.0333. Every kept token was conditioned on a token you just threw away — a stale prefix, silently emitted.
Price the draft at c = 0.60 in demo_5_cost_modelSpeedup 0.825x. Past the break-even c* = 0.4516, speculation costs more than it saves and the only symptom is latency.
Draft with temper(TARGET, 2.0) — same model, wrong temperatureAcceptance falls 1.0000 → 0.6658, speedup 3.571x → 2.617x. A pure configuration bug that halves your gain.
Push k to 16 in demo_4_draft_lengthTokens per call rises to 3.3774 (from 3.0009 at k=5) but speedup falls to 1.299x. You paid for 16 drafts to keep 3.
Then read the paper — the method section, where the acceptance rule and its correctness proof live, and the section on the expected number of accepted tokens, which is the k-sweep above done analytically. Skip the hardware benchmark tables; your acceptance rate will be different anyway. You'll find the paper is just the comments you've already read.

On the maths anxiety

The word "provably" scared me off this paper for a long time. It shouldn't have. The proof needs one idea: two ways of producing a token, whose probabilities add up to the distribution you wanted. Accept gives you min(p, q) of the mass. The residual gives you the rest, max(0, p − q). Add them: p. That's it. Rejection sampling in one sentence, and you can confirm it by counting outcomes in a loop, which is what demo 2 does.

You need to be comfortable with two things and no more. A probability distribution is a vector that sums to 1 — an array with an invariant. Total variation distance is 0.5 * np.abs(a - b).sum(), a diff between two such arrays. Neither is research maths.

What you actually need for this in production is engineering: measuring your real acceptance rate on your real traffic, pricing the draft model, tuning k per workload, and knowing that a distribution bug here is invisible to every eval you own. The arithmetic is the easy half.

Where to go next

This page gave you one thing: the acceptance rule as code you can break, and the cost model that decides whether it's worth it. It skipped the proof written out properly, the tree-structured variants, and everything about running it in a real serving stack.

Order I'd take them in: Hugging Face first to see it run on real models, then Jay Mody for the proof, then the paper. Chen et al. only if the proof still hasn't clicked. And read the KV-cache page before any of them if "decode is memory-bound" wasn't already obvious — it's the premise the whole trick stands on. The GPT page covers the plain decoding loop this replaces.