Speculative Decoding Deep Dive I: Math Foundations

A cheap draft proposes several tokens, the target verifies them in one pass, and rejection sampling keeps the output probability distribution equal to the target.

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}\mathrm{accurate}q+\mathrm{verify}\mathrm{many}\mathrm{tokens}\mathrm{in}\mathrm{one}\mathrm{pass}$$

Core. Exact speculative sampling keeps the target probability distribution. 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^{\prime }$. 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.

Why it is fast

Key: a Transformer can parallel-verify many positions.

Low-batch decode is memory-bound: each ordinary pass reloads weights for one token. Verifying $\gamma$ draft tokens is still one target forward — load weights once, then spend idle GPU compute on $\gamma +1$ positions at once. In that regime, $T_{\mathrm{verify}}$ is often close to a normal decode step.

• 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.

The autoregressive ceiling

If the drafter is itself autoregressive, $T_{\mathrm{draft}}$ is the bottleneck. A KV cache avoids recompute, but not the $\gamma$ serial one-token passes:

$$T_{\mathrm{draft}}=\gamma \cdot t_{\mathrm{step}}$$

Accepted draft tokens (the $\tau$ formula above already adds the bonus/correction) saturate:

$$E[N]=\sum _{i=1}^{\gamma }{\alpha }^i=\frac{\alpha (1-{\alpha }^{\gamma })}{1-\alpha }\underset{\gamma \to \infty }{\to }\frac{\alpha }{1-\alpha }$$

Cost grows linearly in $\gamma$; extra draft tokens buy almost nothing. So an AR drafter must stay shallow, which caps $\alpha$. Parallel drafting (DFlash) makes $T_{\mathrm{draft}}$ independent of $\gamma$.

Why output follows target distribution

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))$$

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{array}{rl}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{array}$$

• $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)$.

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^{\prime }(x)=r(x)/Z$.

Accept probability $A$ is the total overlap. Since $\sum p=1$,

$$A=\sum _{x}q(x)\alpha (x)=\sum _x\min (p(x),q(x))$$ $$P(\mathrm{reject})=1-A=Z$$

Correction path:

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

Add the two paths:

$$\begin{array}{rl}{P}_{\mathrm{SD}}(x)& =\min (p(x),q(x))+\max (0,p(x)-q(x))\\ & =p(x)\end{array}$$

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

Multiple draft tokens

Draft $d_1,\dots ,d_{\gamma }$. After $d_1,\dots ,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),p_i(x)=p(x\mid c_i),{\alpha }_i=\min \left(1,\frac{p_i(d_i)}{q_i(d_i)}\right)$$

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

All $\gamma$ accepted. The same pass has the next-token distribution after the whole draft. Sample $x_{\mathrm{bonus}}\sim p_{\gamma +1}$. Emit $d_1,\dots ,d_{\gamma },x_{\mathrm{bonus}}$.

$$p_{\gamma +1}(x)=p(x\mid \mathrm{prefix}+d_1+\cdots +d_{\gamma })$$

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

$$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.

References

• Leviathan et al., 2023. Fast Inference from Transformers via Speculative Decoding. arXiv 2211.17192
• Chen et al., 2023. Accelerating Large Language Model Decoding with Speculative Sampling. arXiv 2302.01318
• Cai et al., 2024. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv 2401.10774