Speculative decoding: Visual walkthrough — AI-ML deep dive

Prerequisite ideas we lean on: 3.1.2-Transformer-architecture (what "the model outputs a probability for the next word" means) and 5.3.2-Knowledge-distillation (how the draft model is made to agree with the target).

Step 1 — A probability distribution is a set of bar heights

WHAT. When a language model reads "The capital of France is", it does not output one word. It outputs a number for every possible next word — how likely that word is. Those numbers are all positive and add up to $1$ . We call this collection a distribution.

We write $p (x)$ for the height of the bar over word $x$ in the target (big, correct) model. We write $q (x)$ for the same word's bar height in the draft (small, fast) model.

WHY these letters. $p$ = the distribution we are allowed to call correct (the expensive model we actually want). $q$ = the distribution we can sample from cheaply. The whole game is: sample from $q$ , but end up with samples that look like they came from $p$ .

PICTURE. Look at the two rows of bars. Same words on the x-axis, different heights. The target (red) is what truth looks like; the draft (black) is our cheap guess. They disagree — that disagreement is the entire problem.

Step 2 — Sampling from $q$ overshoots and undershoots

WHAT. If we just keep every word the draft proposes, our output distribution is $q$ , not $p$ . Compare the two heights word by word and one of two things is true:

$p (x) \leq q (x)$ : the draft proposes this word too often (its bar is taller than truth).
$p (x) > q (x)$ : the draft proposes this word too rarely (its bar is shorter than truth).

WHY this split matters. These two cases need opposite fixes. Where the draft is too eager we must throw some proposals away. Where the draft is too shy we must manufacture extra proposals. Every later step is just one of these two repairs.

PICTURE. The red curve is $p$ , the black curve is $q$ . Shade where black pokes above red (draft too eager — overshoot) and where red pokes above black (draft too shy — undershoot). Those two shaded regions must have exactly equal area, because both curves sum to $1$ .

Step 3 — The accept rule fixes the overshoot

WHAT. After the draft hands us a word $\tilde{x}$ , we accept it with probability $α (\tilde{x}) = min (1, \frac{p ( x ~ )}{q ( x ~ )}) .$

WHY this exact formula. We want the kept proposals of word $x$ to have height exactly $p (x)$ . The draft delivers word $x$ with height $q (x)$ . So we must keep a fraction $p (x) / q (x)$ of them: $how often draft offers x q (x) \times fraction we keep \frac{p ( x )}{q ( x )} = p (x) .$ The $min (1, \cdot)$ is only there because a probability cannot exceed $1$ : when $p (x) > q (x)$ (undershoot) the ratio $p / q$ is bigger than $1$ , so we clip it to $1$ — meaning keep all of them, we can never keep more than the draft gave us. That leftover deficit is Step 4's job.

PICTURE. Take the black draft bars. Over each word, chop them down to the red line if the red line is lower (overshoot fixed). Where red is above black, keep the whole black bar — you still fall short of red. The kept (shaded) region now hugs $min (p, q)$ , the lower of the two curves.

Step 4 — Rejection buys us a second draw from the deficit

WHAT. Accepting gives us the shaded lower curve $min (p, q)$ . Truth is $p$ . The gap left over, word by word, is $p (x) - min (p (x), q (x)) = max (0, p (x) - q (x)) .$ Whenever we reject, we throw the proposed word away and instead sample a fresh word from exactly this leftover shape: $p_{resample} (x) = \frac{m a x ( 0 , p ( x ) - q ( x ))}{\sum _{y} m a x ( 0 , p ( y ) - q ( y ))} .$

WHY this shape. The numerator $max (0, p - q)$ is precisely the red undershoot region from Step 2 — the words truth wanted but the draft under-supplied. The denominator is just its total area, dividing so the new bars sum to $1$ (a legal distribution). We only reach this branch on a rejection, and rejections happen exactly often enough to fill the gap — that is what Step 5 verifies.

PICTURE. The red curve minus the already-accepted grey $min (p, q)$ region leaves a sliver — the accent-red patch sitting over the undershoot words. That sliver, rescaled to area $1$ , is the distribution we resample from. Notice it lives only where $p > q$ : rejection never resamples a word the draft was already over-supplying.

Step 5 — Add the two branches: the heights become exactly $p$

WHAT. For any word $x$ , there are exactly two ways it can end up as our output:

The draft proposed $x$ and we accepted it — contributes $min (p (x), q (x))$ (Step 3).
We rejected whatever was proposed, then resampled $x$ from the sliver — contributes $Z \cdot p_{resample} (x)$ .

Add them: $P_{final} (x) = accepted min (p (x), q (x)) + resampled Z \cdot \frac{max ( 0 , p ( x ) - q ( x ))}{Z} = min (p, q) + max (0, p - q) .$

WHY it collapses. Look at the identity $min (a, b) + max (0, a - b)$ for any two numbers:

If $a \leq b$ : $min = a$ , $max (0, a - b) = 0$ , sum $= a$ .
If $a > b$ : $min = b$ , $max (0, a - b) = a - b$ , sum $= b + (a - b) = a$ .

Either way the sum is $a$ . With $a = p (x), b = q (x)$ this is $p (x)$ for every word, in both the overshoot and undershoot case. The draft cancels out completely.

PICTURE. Stack the grey accepted region and the red resampled sliver on top of each other. Overshoot words: fully covered by grey alone. Undershoot words: grey (up to $q$ ) plus red sliver (the rest up to $p$ ). The two colours together trace the red target curve exactly — no gaps, no overhang.

Step 6 — Degenerate cases, so nothing surprises you

WHAT / WHY / PICTURE, three edge scenarios in one figure:

Perfect draft, $q = p$ . Overshoot and undershoot vanish. $α = min (1, p / p) = 1$ everywhere — every token accepted, zero rejections. Maximum speed, and Step 5 trivially gives $p$ .
Draft never proposes the truth's favourite word ( $q (x) = 0$ where $p (x) > 0$ ). That word can only ever arrive through the resample branch: $max (0, p - q) = p (x) > 0$ keeps it in the sliver. The rule still reaches it — just always via rejection. Nothing truth wants is unreachable.
Draft is confident garbage ( $q$ tall where $p \approx 0$ ). Then $p / q \approx 0$ , so $α \approx 0$ — almost always rejected. We fall back to resampling from $p$ 's sliver, which here is essentially $p$ itself. Correct output, no speedup — exactly the "worst case = normal decoding" promise.

The one-picture summary

Everything above compressed into a single stacked bar over one word: the black draft bar $q$ , chopped to $min (p, q)$ by the accept test (kept, grey), topped by the red resample sliver $max (0, p - q)$ — the two pieces summing to the red target line $p$ .

Recall Feynman retelling — say it back in plain words

The big model is slow but right; the small model is fast but sloppy. We let the small model blurt out a word. Then we hold its guess up against the big model's own opinion for that word. If the big model likes the word at least as much as the small model did, we keep it. If the big model likes it less, we flip a biased coin — keep it only with probability (big's liking ÷ small's liking) — and if the coin says no, we throw the word out and instead pick a fresh word from a special leftover pile. That pile contains exactly the words the big model wanted more than the small model offered, and nothing else. Add up "words we kept" plus "words we pulled from the leftover pile," and by a tiny algebra fact — $min (a, b) + max (0, a - b) = a$ — you land exactly on the big model's own distribution. The small model has cancelled out. So we get the big model's answers, but most of the time for the price of the small model's speed. A hopeless small model doesn't corrupt the answer — it just makes us throw more words away, so we run slower, never wronger.

Recall Check yourself

Where does a word with $q (x) = 0$ but $p (x) = 0.3$ come from? ::: Only from the resample branch — the accept branch needs the draft to propose it, which it never does; the sliver $max (0, p - q) = 0.3$ carries it. Why is there a $min (1, \cdot)$ instead of just $p / q$ ? ::: Because $p / q > 1$ in undershoot regions, and you cannot keep a larger fraction of proposals than you were given; you cap at "keep all of them." What does a perfect draft $q = p$ do to the acceptance rate? ::: Makes it $1$ everywhere — nothing is ever rejected, maximum speedup.