Speculative Decoding

Setup and Notation

Let $p_\theta(\cdot \mid x_{1:t})$ denote the target model's conditional distribution over vocabulary $\mathcal{V}$ at position $t+1$ given prefix $x_{1:t}$, and let $q_\phi(\cdot \mid x_{1:t})$ denote the corresponding draft model distribution. At each speculative step the draft model autoregressively generates $K$ candidate tokens $\tilde{x}_{t+1}, \ldots, \tilde{x}_{t+K}$, after which the target model performs a single forward pass over the extended sequence $x_{1:t}, \tilde{x}_{t+1}, \ldots, \tilde{x}_{t+K}$ to obtain all $K+1$ conditional distributions simultaneously.

Acceptance Probability

For each candidate token $\tilde{x}_{t+k}$ drawn from the draft distribution $q_\phi$, define the acceptance probability:

Explanation: Each draft token is accepted with probability equal to the ratio of target to draft probability, clamped to 1, so tokens the target model finds at least as likely as the draft are always accepted. Here $\alpha_k \in [0,1]$; when $q_\phi$ assigns higher probability than $p_\theta$ to a token, $\alpha_k < 1$ and rejection can occur. A Bernoulli trial with probability $\alpha_k$ decides acceptance. If token $k$ is accepted the algorithm proceeds to token $k+1$; otherwise a correction sample is drawn.

Corrected Distribution on Rejection

When candidate $\tilde{x}_{t+k}$ is rejected, the algorithm samples a bonus token from the residual distribution:

Explanation: The corrected distribution redistributes probability mass to tokens where the target exceeds the draft, ensuring the marginal distribution over accepted-or-corrected tokens is exactly $p_\theta$. The $\text{norm}(\cdot)$ operator normalizes the non-negative residual vector to sum to 1. This is the key identity: combining the acceptance step (Eq. [1]) with the correction step (Eq. [2]) yields outputs with marginal distribution $p_\theta$ regardless of $q_\phi$.

Distribution-Preservation Theorem

Theorem (Leviathan et al., 2023). For any draft distribution $q_\phi$ and any draft token $\tilde{x}$ drawn from $q_\phi$, the token emitted by the accept/correct procedure has distribution $p_\theta$. Proof sketch. Let $x \in \mathcal{V}$ be arbitrary. The probability of emitting $x$ is:

Explanation: The first term accounts for accepting the draft token $x$; the second accounts for rejecting whatever draft token was drawn and then sampling $x$ from the corrected distribution. Substituting Eq. [1] for $\alpha(x)$ and Eq. [2] for $p_{\text{corr}}(x)$ and simplifying via the identity $\sum_x \max(0, p-q) = \sum_x \max(0, q-p)$ (both equal $\frac{1}{2}\|p-q\|_1$) yields $\Pr[\text{emit}\;x] = p_\theta(x)$ for all $x$. $\square$

Expected Tokens per Step

The key throughput quantity is the expected number of tokens produced per target-model call. Let $\beta = \mathbb{E}[\alpha_k]$ denote the mean acceptance rate. Assuming i.i.d. acceptance across positions (an approximation; see Stern et al. 2018 for tighter bounds):

Explanation: This geometric-series formula counts the expected number of tokens accepted before the first rejection plus the mandatory bonus token, showing that even moderate acceptance rates yield significant multi-token gains. When $\beta \to 1$ (draft nearly matches target), $\mathbb{E}[\text{tokens per step}] \to K+1$; when $\beta \to 0$ it degrades to 1, i.e., standard decoding. In practice $\beta \approx 0.7$–$0.85$ for well-matched draft/target pairs, giving speedups of 2–4×.

KV-Cache and Inference Optimization · Non-Autoregressive Generation · Knowledge Distillation · Beam Search and Decoding Strategies · Inference-Time Compute Scaling