Fast Inference from Transformers via Speculative Decoding

Autoregressive Transformers decode one token at a time. Generating $K$ tokens normally means $K$ serial runs of the large target model. That loop is often memory-bandwidth bound: the hardware has spare arithmetic capacity, but each step still waits on the previous token.

Speculative decoding attacks that serial bottleneck. A cheaper draft model $M_q$ proposes $\gamma$ tokens. The large target $M_p$ then scores those guesses (and one extra prefix) in parallel. Speculative sampling accepts as long a prefix as possible while guaranteeing the emitted tokens still follow the target distribution alone — no architecture change, no retrain, identical outputs.

What this paper explains

Standard decoding from $M_p$ is:

1. Condition on the prefix.
2. Sample one token from $p(x_t \mid x_{<t})$.
3. Append and repeat.

Speculative decoding inserts a draft:

1. Sample $\gamma$ tokens autoregressively from a faster $M_q$.
2. Run $M_p$ once across the $\gamma+1$ prefixes needed to score those guesses in parallel.
3. Accept the longest prefix allowed by speculative sampling; on the first rejection, resample from an adjusted residual of $p$; if everything is accepted, still sample one bonus token from $M_p$.

So each parallel target pass produces at least one new token and at most $\gamma+1$.

Prior limits

Earlier speedups usually trade something away:

• Distillation, sparsity, early-exit, and adaptive compute often change architectures or training.
• Many adaptive methods do not preserve the exact target distribution.
• Naïve “run a small model instead” changes answers.

Speculative decoding’s claim is stricter: accelerate existing off-the-shelf models, keep identical sampling behavior, and exploit spare concurrency when memory bandwidth — not FLOPs — is the limiter.

The mechanism

Draft. $M_q$ proposes $x_1,\ldots,x_\gamma$ cheaply.

Verify. $M_p$ evaluates $p_1,\ldots,p_{\gamma+1}$ for prefixes $\text{prefix}$, $\text{prefix}+[x_1]$, …, $\text{prefix}+[x_1,\ldots,x_\gamma]$ in parallel.

Speculative sampling. For each draft token $x$ with draft probability $q(x)$ and target probability $p(x)$:

• If $q(x) \le p(x)$, accept.
• If $q(x) > p(x)$, reject with probability $1 - p(x)/q(x)$.
• On rejection, sample from the normalized residual $\mathrm{norm}(\max(0, p - q))$.

That rule is distribution-preserving: the accepted (or resampled) token is still distributed as $p$.

Acceptance rate. Let $\beta$ be the probability a draft token is accepted for a given prefix, and $\alpha = \mathbb{E}(\beta)$. Under an i.i.d. simplification, the number of tokens from one Algorithm-1 step is a capped geometric random variable:

between 1 and $\gamma+1$.

Algorithm walkthrough

1. Sample $\gamma$ draft tokens from $M_q$ autoregressively.
2. Run $M_p$ in parallel on the $\gamma+1$ prefixes that score those drafts.
3. Draw uniforms $r_i$ and accept until the first $i$ with $r_i > p_i(x_i)/q_i(x_i)$.
4. If a rejection happens at position $n$, resample token $n+1$ from the residual of $p_{n+1}$ vs $q_{n+1}$.
5. If all $\gamma$ drafts pass, still sample one more token from $p_{\gamma+1}$.

The opening figure is that story as a scene: small cube proposes a chain; large cube verifies; green checks keep a prefix, a red cross marks the first reject, and a star marks the corrected or bonus token.

What to notice when reading

• Speculative execution is old hardware lore; the novelty is stochastic speculative sampling that preserves $p$.
• Walltime gains need enough spare concurrency to run $\gamma+1$ target evaluations without stretching each one’s latency.
• Draft quality $\alpha$ and relative draft cost $c$ jointly set the best $\gamma$ (paper Figure 3 / Theorem 3.8).
• Even a trivial bigram draft can give $\alpha \approx 0.2$ and a small but real speedup on En→De.

Results and evidence

Implemented against T5X for T5-XXL (11B) on a single TPU-v4, batch size 1. Drafts are existing T5 checkpoints (small / base / large). Reported walltime speedups with identical outputs:

Larger drafts raise $\alpha$ but also raise $c$, so walltime can get worse even when guesses are better (T5-large rows). Argmax (temp=0) accepts more often than temp=1.

Limitations

• Needs spare parallel capacity; if $\gamma+1$ concurrent $M_p$ passes inflate latency, the win shrinks.
• Rejected drafts waste work; total FLOPs can rise even when walltime falls.
• Optimal fixed $\gamma$ depends on $\alpha$ and $c$; an oracle-varying $\gamma$ could do better (left for future work in the paper).
• Gains track how well $M_q$ mimics $M_p$ on the task — hard domains with low $\alpha$ help less.

How to read the paper

1. Abstract + §1 — serial decode vs speculative concurrency.
2. §2 — speculative sampling rule and Algorithm 1.
3. §3 — $\alpha$, expected tokens, walltime theorem, choosing $\gamma$.
4. §4 — T5-XXL walltimes (Table 2) and supporting tasks.
5. Appendix — correctness proof for speculative sampling.

Knowledge check

Keep reading

1. Original paper (arXiv:2211.17192) — Algorithm 1, Eq. 1, Table 2.
2. Related Fanout Daily: Orca continuous batching, GQA, FlashAttention, PagedAttention.
3. Follow-on systems literature on speculative decoding in production LLM servers (same draft-verify idea; implementations differ).

Sources

• Leviathan, Kalman, Matias, Fast Inference from Transformers via Speculative Decoding, ICML 2023 — arXiv:2211.17192