Speculative Decoding

Verifying $\gamma$ tokens costs one forward pass of the target model, because the pass is memory-bound and processing $\gamma$ positions instead of one is nearly free. The saving is real only if the draft is usually right.

Modified rejection sampling makes the scheme exact. Accept a drafted token with probability $\min(1, p/q)$; on rejection, resample from the normalised positive part of the difference.

With per-token acceptance rate $\alpha$ the expected number of tokens produced per target pass is a geometric sum, so the speedup saturates: past $\alpha \approx 0.8$ there is little to gain from drafting deeper.

import torch
from torch import Tensor


def verify(draft_tokens: Tensor, p: Tensor, q: Tensor) -> tuple[Tensor, int]:
    """p, q: [gamma, V] target and draft distributions over the draft.
    Returns accepted tokens and the count accepted."""
    accepted = []
    for i, tok in enumerate(draft_tokens):
        ratio = (p[i, tok] / q[i, tok]).clamp(max=1.0)
        if torch.rand(()) < ratio:
            accepted.append(tok)
            continue
        # rejection: resample from the residual distribution
        residual = (p[i] - q[i]).clamp(min=0)
        residual = residual / residual.sum()
        return torch.tensor(accepted + [residual.multinomial(1).item()]), i

    # all accepted: take a free bonus token from the target
    return torch.cat([draft_tokens, p[-1].multinomial(1)]), len(draft_tokens)

The bonus token on full acceptance is what makes the expected-tokens formula reach $\gamma + 1$, and dropping it costs roughly one token of throughput per accepted block. The residual must be clamped before normalising; normalising a signed vector produces negative probabilities and a silent corruption of the output distribution.