AI ENGINEERING: DEEP DIVE · MODULE 08 — SCALE & INFERENCE · 24 A hands-on series · math + code

Module 8.3

Faster, not different

A small, cheap "draft" model guesses several tokens ahead; the real model checks all of them in one parallel pass and accepts or corrects each one with exact math — never approximating what the big model would have said.

CONTINUING: "Your export failed because your ___" — DRAFT MODEL PROPOSES 4 TOKENS AHEAD
Step 1 of 5

The math

Accept, or correct — never approximate

Acceptance rule
accept token x with prob. min(1, p(x)/q(x))
p = target model's probability, q = draft model's
if p ≥ q, always accept — the draft undersold it
Correction on reject
p'(x) = max(0, p(x) − q(x)) / Z
resample from what's "left over" after removing the draft's mass
Z normalizes the leftover distribution to sum to 1
Why it's exact
P(final = x) = p_target(x), for every x
accept-or-correct is constructed to cancel the draft's bias exactly
same output distribution, fewer sequential target passes

This is the part that makes speculative decoding different from just "using a smaller, faster model": the output is provably identical in distribution to always running the big model alone.

The code

One round of speculative decoding

speculative_decode.py
def speculative_round(draft_tokens, q_probs, p_probs, p_dists, q_dists):
    accepted = []
    for x, q, p, p_dist, q_dist in zip(draft_tokens, q_probs, p_probs, p_dists, q_dists):
        accept_prob = min(1.0, p / q)
        if np.random.random() < accept_prob:
            accepted.append(x)                       # keep the draft's token
        else:
            leftover = np.maximum(0, p_dist - q_dist)
            corrected = leftover / leftover.sum()
            accepted.append(sample(corrected))       # exact correction, then stop
            break
    return accepted   # one target-model forward pass produced all of this

The target model scores all 4 drafted positions in a single batched forward pass — the expensive part barely gets more costly for checking 4 tokens instead of 1, which is exactly where the speedup comes from.

The Nimbus thread

Nimbus-mini drafts, Nimbus-large verifies

Nimbus could run a small, fast "Nimbus-mini" model to draft several tokens of a support reply, then have the full "Nimbus-large" model verify them all in one pass using exactly the accept/correct math above. The customer sees the same replies, with the same quality guarantees, as if Nimbus-large had generated every token alone — just fewer expensive sequential steps to get there.

So what

Speculative decoding buys speed with extra cheap compute — never by relaxing what the big model is allowed to say.