Speculative Decoding

The problem

Generating with a large model is slow because every token needs a full pass through it, and a pass is limited by reading memory. But scoring several tokens in one pass costs almost the same as scoring one. Speculative decoding uses this: a small, fast draft model proposes $\gamma$ (gamma) tokens, and the large target model checks all of them in one pass. You will write the step that decides which proposals to keep, so that the result is distributed exactly as if the target had sampled alone (Leviathan et al., 2023).

The inputs:

• draft: the $\gamma$ proposed tokens;
• q, shape $(\gamma, V)$: the draft's distribution over the $V$ tokens at each proposed position;
• p, shape $(\gamma + 1, V)$: the target's distribution at the same positions, plus one more after the last proposal.

The rule

For $i = 0, \dots, \gamma - 1$, with $x$ = draft[i]:

• draw u = rng.random() and accept the token if u < p[i, x] / q[i, x]: if the target likes it at least as much as the draft ($p \ge q$) it is always accepted; otherwise it is accepted with probability $p/q$;
• if rejected, sample a replacement from the residual distribution $\mathrm{norm}(\max(0, p_i - q_i))$, where $\mathrm{norm}$ divides by the total so that it sums to 1, and stop;
• if all $\gamma$ proposals were accepted, sample one more token from p[γ], which costs nothing extra.

Return the list of new tokens: the accepted ones, followed by the replacement or the extra token. It always has between 1 and $\gamma + 1$ tokens.

Example. $V = 3$, $\gamma = 2$. Target $p$ = [[0.2, 0.5, 0.3], [0.1, 0.1, 0.8], [0.3, 0.3, 0.4]], draft $q$ = [[0.5, 0.3, 0.2], [0.2, 0.2, 0.6]].

All accepted: draft = [1, 2] and the random numbers come out $0.512, 0.950, 0.144$.

• Position 0, token 1: $p/q = 0.5/0.3 = 1.67$; $0.512 < 1.67$: accepted.
• Position 1, token 2: $p/q = 0.8/0.6 = 1.33$; $0.950 < 1.33$: accepted.
• All accepted: one more token from p[2] = [0.3, 0.3, 0.4] with the next random number $0.144$. The cumulative sums are $0.3, 0.6, 1.0$, and $0.144$ falls in the first: token 0. The result is [1, 2, 0]: 3 tokens from one target pass.

Rejection: draft = [0, 2], same random numbers.

• Position 0, token 0: $p/q = 0.2/0.5 = 0.4$; $0.512 \ge 0.4$: rejected.
• Residual: $\max(0, p_0 - q_0) = (0, 0.2, 0.1)$, normalized $(0, 0.667, 0.333)$. The next random number, $0.950$, falls in the third interval (cumulative sums $0, 0.667, 1.0$): token 2.
• Stop. The result is [2]: the second proposal is never looked at.

Why the result is exactly $p$: the replacement comes only from tokens the target prefers to the draft, which are exactly the ones the draft under-proposes. The theory page checks it with numbers.

What to implement

speculative_step(p, q, draft, rng): return the list of new tokens. Sample from a distribution with sample(probs, rng), which is given: the inverse CDF (cumulative distribution function) method with one rng.random() call.

Requirements

• Use the random numbers exactly in this order (one per accept test, one per sample), so the tests can reproduce your choices.
• u < p/q is the same as accepting with probability $\min(1, p/q)$.