Speculative Decoding Deep Dive

This session explains speculative decoding: a cheap draft proposes several tokens, the target verifies them in one pass, and rejection sampling keeps the output law equal to the target.

1. Overview

Autoregressive decoding: one token per target pass. Latency scales with length.

Speculative decoding (Chen et al. 2023; Leviathan et al. 2023): draft several tokens cheaply, verify them in one target pass. Methods differ in how they build $q$.

$$\mathrm{cheaper}\ q +\ \mathrm{more\ accurate}\ q +\ \mathrm{verify\ many\ tokens\ in\ one\ pass}$$

Core. Exact speculative sampling keeps the target law. A worse $q$ only lowers the accept rate.

$$p_{\mathrm{SD}}=p$$

Prefix: Under the pale moon

• Case 1: Reject. Draft a cat crossed the garden, keep a cat, drop the suffix, resample slept from $p'$. Emit Under the pale moon a cat slept.
• Case 2: All accept. Prefix Beside the silent lake, draft birds began to sing, bonus softly. Emit Beside the silent lake birds began to sing softly.

2. Why it is fast

Low-batch decode is memory-bound: each target pass reloads weights for one token. Speculative decoding spends idle compute — the target checks $\gamma+1$ positions in one pass.

• Acceptance rate $\alpha$ — mean keep-probability over positions/contexts
• Accepted length $\tau$ — expected tokens emitted per cycle (accepted drafts plus one correction or bonus)

$$\alpha=\mathbb{E}_{x\sim q}\left[\min\left(1,\frac{p(x)}{q(x)}\right)\right]=\sum_x\min(p(x),q(x))$$

If accepts are i.i.d. with mean $\alpha$ and draft length $\gamma$:

$$\tau=\sum_{i=0}^{\gamma}\alpha^{i}=\frac{1-\alpha^{\gamma+1}}{1-\alpha}$$

If $\alpha=1$, then $\tau=\gamma+1$. The $i=0$ term is the token the target always emits, so $\tau\in[1,\gamma+1]$.

Per-token latency (Sadhukhan et al.):

$$L=\frac{T_{\mathrm{draft}}+T_{\mathrm{verify}}}{\tau}$$

Latency per token drops if (1) drafts match $p$ more often (larger $\tau$); (2) drafting is cheaper; (3) verification is cheaper.

3. Math

This section proves that speculative sampling outputs $x_{\mathrm{out}}\sim p$, the same as the target distribution.

Let $p$ be the target next-token distribution and $q$ the draft distribution, same context.

Split $p$ into overlap (accepted drafts) plus leftover (filled only after reject):

$$p(x)=\min(p(x),q(x))+\max(0,p(x)-q(x))$$

3.1 Acceptance

$\alpha(x)$ is the probability we accept the draft token, given that the draft proposed $x$:

$$\alpha(x)=P(\mathrm{accept}\mid \mathrm{draft}=x)$$

Draw $u\sim\mathrm{Unif}[0,1]$; keep $x$ iff $u\le\alpha(x)$. For $q(x)>0$,

$$\alpha(x)=\min\left(1,\frac{p(x)}{q(x)}\right)$$

Then

$$\begin{aligned} P(x,\mathrm{accept}) &= q(x)\,\alpha(x) \\ &= q(x)\min\left(1,\frac{p(x)}{q(x)}\right) \\ &= \min(p(x),q(x)) \end{aligned}$$

• $q(x)>p(x)$: over-propose; keep fraction $p/q$
• $q(x)\le p(x)$: under-propose; keep every such $x$

Acceptance stores exactly the overlap $\min(p,q)$.

3.2 Rejection correction

$r(x)$ is the uncovered part of $p(x)$: how much target mass is still missing after the accept path has taken $\min(p,q)$.

• $p(x)>q(x)$: accept only covered $q(x)$, so $r(x)=p(x)-q(x)$
• $p(x)\le q(x)$: accept already covered all of $p(x)$, so $r(x)=0$

$$r(x)=p(x)-\min(p(x),q(x))=\max(0,p(x)-q(x))$$

$Z$ is the total uncovered mass (over the vocabulary):

$$Z=\sum_y r(y)$$

If rejected and $Z>0$, sample the correction from $p'(x)=r(x)/Z$.

Accept probability $A=\sum_x q(x)\alpha(x)=\sum_x\min(p(x),q(x))$. Since $\sum p=1$,

$$P(\mathrm{reject})=1-A=Z$$

Correction path:

$$P(\mathrm{reject})\,p'(x)=Z\cdot\frac{r(x)}{Z}=\max(0,p(x)-q(x))$$

Add the two paths:

$$\begin{aligned} P_{\mathrm{SD}}(x) &= \min(p(x),q(x))+\max(0,p(x)-q(x)) \\ &= p(x) \end{aligned}$$

So $P_{\mathrm{SD}}=p$. If $p=q$ then $Z=0$: every proposal is kept and $p'$ is unused.

3.3 Multiple draft tokens

Draft $d_1,\ldots,d_\gamma$. After $d_1,\ldots,d_{i-1}$ are kept, both models share context $c_i=\mathrm{prefix}+d_1+\cdots+d_{i-1}$. Verify left to right:

$$q_i(x)=q(x \mid c_i), \qquad p_i(x)=p(x \mid c_i), \qquad \alpha_i=\min\left(1,\frac{p_i(d_i)}{q_i(d_i)}\right)$$

Reject at $k$. Keep $d_1,\ldots,d_{k-1}$. Drop $d_k,\ldots,d_\gamma$ (later drafts assumed $d_k$). Sample $x_{\mathrm{corr}}\sim p_k'$. Emit $d_1,\ldots,d_{k-1},x_{\mathrm{corr}}$.

All $\gamma$ accepted. The same pass has $p_{\gamma+1}(x)=p(x\mid \mathrm{prefix}+d_1+\cdots+d_\gamma)$. Sample $x_{\mathrm{bonus}}\sim p_{\gamma+1}$. Emit $d_1,\ldots,d_\gamma,x_{\mathrm{bonus}}$.

A cycle emits between $1$ and $\gamma+1$ tokens. Repeating the one-step rule at each valid context gives the target sequence law:

$$P_{\mathrm{SD}}(x_{1:n}) = \prod_{i=1}^{n}p(x_i \mid x_{1:i-1}) = P_{\mathrm{target}}(x_{1:n})$$

Speculative decoding changes how many target passes are needed, not the target output distribution.