Lost in Latent Space

Speculative Decoding: A Walkthrough

The idea: instead of running the big model one token at a time, have a small model propose several tokens, then run the big model once on the whole proposed sequence in parallel. Use the big model’s outputs to decide which proposed tokens to keep. If everything checks out, you generated K tokens for the cost of one big-model forward pass.

Here’s the loop. Let p be the big (target) model and q the small (draft) model.

def speculative_decode(prompt, gamma=5):
    tokens = list(prompt)
    while not done(tokens):
        # 1. Draft model proposes gamma tokens autoregressively.
        draft = []
        for _ in range(gamma):
            q_dist = q(tokens + draft)        # distribution over next token
            draft.append(sample(q_dist))

        # 2. Target model runs ONCE on the whole speculative sequence,
        #    returning distributions p_0 ... p_gamma at each position.
        #    This is the expensive call, but it's a single batched forward pass.
        p_dists = p(tokens + draft)            # length gamma+1

        # 3. Walk through the draft, accepting or rejecting each token.
        accepted = []
        for i, t in enumerate(draft):
            r = uniform(0, 1)
            if r < min(1, p_dists[i][t] / q_dists[i][t]):
                accepted.append(t)             # accept
            else:
                # Reject. Resample from a corrected distribution.
                corrected = normalize(max(0, p_dists[i] - q_dists[i]))
                accepted.append(sample(corrected))
                break

        # 4. If all gamma drafts were accepted, we get a bonus token
        #    from p_dists[gamma] (the target's prediction after the last
        #    accepted draft token) — free, because it's in the same batch.
        if len(accepted) == gamma:
            accepted.append(sample(p_dists[gamma]))

        tokens.extend(accepted)
    return tokens

Why this preserves the target’s output distribution exactly. The rejection sampler in step 3 is the standard one for sampling from p given proposals from q. The acceptance probability min(1, p/q) is chosen so that, marginalizing over the random uniform, the distribution of the accepted token (or the resampled-on-reject token) equals p exactly. There is no approximation here — speculative decoding produces the same output distribution as plain target-model sampling. The only randomness it introduces is in how many drafts get accepted per round, which affects latency, not output quality.

Where the speed comes from. The target model is run once per round, not gamma times. Modern GPUs are bandwidth-bound at inference: loading the target model’s weights into SRAM dominates, and once they’re loaded, processing gamma extra positions is nearly free. If the draft model is, say, 30x smaller than the target, drafting gamma=5 tokens costs ~5/30 = 1/6 of a target forward pass, and if all 5 are accepted you get a 6-token output (5 drafts + 1 bonus) for ~1.17 target-equivalent passes. Net speedup ~5x.

Where it fails. Acceptance rates are highly task-dependent. On code completion where the draft model is well-aligned with the target, typical acceptance fractions hit 70-80%, and gamma can be set to 7-8. On creative writing with high-temperature sampling, the draft and target disagree more, acceptance drops below 50%, and gamma needs to shrink to keep the wasted compute under control.

Practical knobs.