Fast Inference from Transformers via Speculative Decoding

Problem

Framing

Autoregressive Transformer latency scales linearly with the number of decoded tokens because each token needs one serial target-model pass. The paper closes this with speculative decoding: a small draft model proposes $\gamma$ tokens, the large model verifies them in parallel, and exact output distribution is preserved while reaching roughly $2\times$–$3\times$ speedups on T5-XXL and Chinchilla-class models.

Currently Used Methods

Direct antecedents

• @hintonDistillation2015 — compresses a large model into a smaller student for cheaper inference.
  • Limitation in context: reduces quality or requires model replacement, not exact target-model decoding.
• @chenCodex2021 — uses lightweight drafting and verification for code generation.
  • Limitation in context: specialized to code-edit settings, not exact general autoregressive sampling.
• "Attention Is All You Need" — standard autoregressive Transformer decoding with KV caching.
  • Limitation in context: still needs one serial target-model pass per emitted token.
• "Fast Inference from Sequence Models via Speculative Execution" — speculative execution for sequence transduction.
  • Limitation in context: assumes aligned outputs, not arbitrary next-token Transformer decoding.

Proposed Method

Architecture

Two autoregressive models share the same tokenizer and sampling standardization: a large target model $M_p$ and a much cheaper approximation model $M_q$. Each step lets $M_q$ draft $\gamma$ continuations autoregressively, then runs $M_p$ once on the prefix plus drafted tokens to verify several positions in parallel.

Loss / Objective

The paper is algorithmic.It defines an adjusted fallback distribution after rejection:

Sampling Rule / Algorithm

Accept each drafted token with a modified rejection rule; on rejection, resample from the residual of $p$ over $q$.

Training Procedure

• Draft LM example: 6M-parameter GPT-like decoder on lm1b.
• Draft LM dimensions: hidden size 256.
• Feed-forward width: 1024.
• Layers: 2.
• Attention heads: 4.
• Context length: 8k tokens.
• Verification width: $\gamma$ drafted tokens per step.

Evaluation

Datasets

• lm1b unconditional language modeling.
• EnDe translation with T5 family.
• Summarization with T5 family.
• Chinchilla-class language modeling.

Metrics

• Empirical speedup factor.
• Acceptance rate $\alpha$.
• Cost coefficient $c$.
• Expected vs empirical runtime improvement.

Headline results

• T5-XXL EnDe, $T=0$, T5-small draft: $\alpha=0.75$, $c=0.02$, expected $3.2\times$, empirical $3.4\times$.
• T5-XXL EnDe, $T=0$, T5-base draft: $\alpha=0.80$, $c=0.04$, expected $3.3\times$, empirical $2.8\times$.
• T5-XXL EnDe, $T=0$, T5-large draft: $\alpha=0.82$, $c=0.11$, expected $2.5\times$, empirical $1.7\times$.
• T5-XXL with lenience on EnDe: improvement rises from $2.5\times$ at $l=1$ to $5\times$ at $l=0.1$.
• Example generation: 38 tokens require only 9 serial large-model runs.

Ablations

• Draft model size: larger drafts raise $\alpha$ but can hurt net speedup when $c$ grows.
• Sampling mode: argmax yields higher $\alpha$ and larger speedups than temperature-1 sampling.
• Lookahead $\gamma$: optimal $\gamma$ increases as $\alpha$ rises and decreases as $c$ rises.
• Lenience: relaxing exactness increases $\alpha$ and can push speedup to $5\times$ on T5-XXL.

Method Strengths and Weaknesses

Strengths

• Exact decoding distribution is preserved despite multi-token drafting.
• One target-model pass can emit several tokens in parallel.
• Theory predicts runtime gains well enough to choose $\gamma$.
• Works across language modeling, translation, and summarization setups.

Weaknesses

• Gains depend sharply on acceptance rate $\alpha$.
• A stronger draft model can erase speedup by increasing cost $c$.
• Exact method still needs target-model evaluation every iteration.
• Reported maximal gains are near $3\times$ without relaxing exactness.

Suggestions from the authors

• Design better approximation models that maximize $\alpha$ at fixed cost $c$.
• Exploit relaxed exactness with lenience for larger practical speedups.
• Optimize implementations to reduce system overhead beyond the theoretical model.
• Extend speculative decoding to broader Transformer inference settings.

Links

Prior Papers

• @hintonDistillation2015 — distillation is a core baseline for reducing inference cost with smaller models.
• @chenCodex2021 — provides a direct precursor that also uses draft-and-verify generation.

Further Papers

• @kwonvLLM2023 — extends the systems view of speculative decoding for LLM serving.
• @daoFlashAttention2_2023 — targets the same Transformer inference bottleneck from the kernel side.
• @jiangMixtral2024 — overlaps on fast large-model inference where speculative decoding is a natural complement.