Source-linked AI summary
Accelerating Large Language Model Decoding with Speculative Sampling
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, John Jumper
TL;DR
Large transformer decoding is costly because autoregressive, memory-bandwidth-bound generation requires repeated model calls and can incur distributed communication overhead. The paper introduces speculative sampling, which drafts short continuations with a faster model, scores them with the target model, and applies modified rejection sampling to preserve its distribution. Sampling Chinchilla achieves a 2–2.5× decoding speedup without modifying the target model or compromising sample quality.
Problem
Autoregressive decoding of large transformers is costly because repeated memory-bandwidth-bound calls and model parallelism add latency.
Method
Speculative sampling drafts K tokens with a faster model, scores them with the target model, and accepts them using modified rejection sampling.
Results
2–2.5× decoding speedup is achieved when sampling from Chinchilla, without modifying the target model or biasing the sample distribution.
Takeaways & Limitations
Speculative sampling accelerates decoding while preserving the target distribution within hardware numerics and integrating with existing target-model infrastructure.
Takeaways & Limitations
Larger draft lengths can increase latency variance, which is problematic when P90 or P99 latency matters.
Abstract
from arXiv · showhide
We present speculative sampling, an algorithm for accelerating transformer decoding by enabling the generation of multiple tokens from each transformer call. Our algorithm relies on the observation that the latency of parallel scoring of short continuations, generated by a faster but less powerful draft model, is comparable to that of sampling a single token from the larger target model. This is combined with a novel modified rejection sampling scheme which preserves the distribution of the target model within hardware numerics. We benchmark speculative sampling with Chinchilla, a 70 billion parameter language model, achieving a 2-2.5x decoding speedup in a distributed setup, without compromising the sample quality or making modifications to the model itself.
Introduction
Large transformer models improve performance but make autoregressive decoding costly. Speculative sampling reduces this latency by drafting tokens with a faster model, scoring them with the target model, and using modified rejection sampling.
- Motivation: Transformer decoding becomes costly because each new token requires another memory-bandwidth-bound, model-parallel transformer call.Parameter movement, transformer memory, and inter-device communication contribute to latency.
- Method: Speculative sampling accelerates decoding by generating a short draft with a faster model and scoring it with the larger target model.The draft model may generate K tokens autoregressively or in parallel.
- Mechanism: Strong agreement between draft and target distributions permits multiple tokens to be generated per target-model call.The approach exploits short continuations whose parallel scoring has comparable latency to single-token sampling.
- Results: 2–2.5× speedup is achieved when sampling from Chinchilla, depending on the evaluation domain.The method reduces sampling latency without modifying the target model or biasing the sample distribution.
Related Work
Prior approaches reduce transformer sampling cost through compression, architectural changes, or parallel decoding. Speculative sampling addresses limitations of earlier parallel methods, including restricted sampling settings and limited distributed-scale evaluation.
- Model efficiency: Quantization and distillation reduce sampling latency with little to no performance penalty, while smaller models are also motivated by scaling observations.These approaches target the model-size contribution to memory-bound sampling.
- Architectural optimization: Multi-query attention shrinks the key-value cache and is most effective for throughput at larger batch sizes rather than latency for large models.Large-model parameters consume most of the memory-bandwidth budget in the cited setting.
- Prior systems: Low-level TPU optimizations and related techniques substantially improved PaLM 540B serving latency and efficiency.This work is cited as part of the broader effort to reduce large-model sampling latency.
- Parallel decoding: Earlier parallel decoding methods were limited by greedy sampling, biased results, other modalities, or lack of distributed evaluation at very large scale.The passage identifies these as limitations of block parallel sampling, aggressive decoding, and related methods.
- Positioning: Speculative sampling was developed independently of contemporaneous speculative decoding work and shares its core underlying idea.The paper emphasizes distributed serving for large models and adds incremental optimizations.
Auto-regressive Sampling
Autoregressive sampling generates tokens sequentially, making large transformer models memory-bandwidth bound. Model size and multi-accelerator serving impose hard latency ceilings and communication overheads.
- Auto-regressive Sampling: Autoregressive sampling draws each next token from the target model conditioned on all preceding tokens until the target length is reached.The cited algorithm initializes at the prompt length and repeatedly samples one token.
- Auto-regressive Sampling: Memory-bound model calls generate only a single token for each sequence in the batch.Consequently, generating multiple tokens introduces substantial latency.
- Auto-regressive Sampling: Model size divided by aggregate accelerator memory bandwidth sets a hard ceiling on maximum autoregressive sampling speed.Larger models also incur inter-device communication latency when served across multiple accelerators.
- Speculative Sampling: Speculative sampling uses a lookahead K with separate autoregressive target and draft models.The corresponding algorithm is specified for an initial prompt and target sequence length.
- Speculative Sampling: The speculative procedure samples draft tokens autoregressively, then computes target-model logits for the draft continuation in parallel.The algorithm description includes the draft sequence and K+1 target-logit sets.
Speculative Sampling
Speculative sampling combines short draft continuations with parallel target scoring and sequential modified rejection sampling. The procedure preserves the target distribution within hardware numerics while enabling multiple accepted tokens per loop.
- Speculative Sampling: Parallel scoring of a short K-token continuation has latency similar to sampling one token from a large target transformer.This comparison motivates using a faster draft model to propose multiple tokens before each target-model call.
- Speculative Sampling: For small K, memory-bound linear layers and attention remain central components of the sampling-time cost.Implementation details such as positional encoding, decoding method, and hardware can create additional scoring differences.
- Modified Rejection Sampling: Draft tokens are accepted left to right using a modified rejection-sampling scheme based on target and draft probabilities.Rejected tokens are resampled from the positive part of the target-minus-draft distribution.
- Modified Rejection Sampling: The sequential procedure recovers the target-model distribution for accepted tokens within hardware numerics.At least one token is generated per draft-accept loop, and a fully accepted draft permits one additional token.
- Modified Rejection Sampling: Standard sampling controls such as nucleus sampling, top-k sampling, and temperature adjustment can be applied before rejection sampling.The observed overall acceptance rate is robust to the exact parameters used.
- Compatibility: Because it does not modify the transformer body, speculative sampling can be combined with quantization and multi-query attention.These techniques target memory use or sampling efficiency independently of the speculative procedure.
Choice of Draft Models
Speculative sampling can use a separate draft model as long as it exposes logits and achieves sufficient acceptance or low latency. The paper favors a smaller target-like model as a practical draft, while noting that jointly trained or distilled alternatives require substantial resources or target-model changes.
- Choice of Draft Models: Any drafting method is acceptable if it exposes logits and provides sufficiently high acceptance or sufficiently low latency to break even.The modified rejection criterion preserves the target distribution under this choice of draft.
- Choice of Draft Models: Drafting methods may incorporate generation into the target model, use sequence-level distillation, or train a draft model from target activations.These approaches are presented as alternatives for generating continuations whose tokens can be evaluated by the target model.
- Choice of Draft Models: Jointly trained or distilled draft methods require target-generated data, target-model changes, or substantial compute, limiting their practicality at large scale.The paper specifically identifies sequence-level distillation as requiring a large compute budget.
- Choice of Draft Models: A smaller version of the target language model can serve as the draft because some tokens are easier to predict and may be handled adequately by smaller models.The paper also notes engineering convenience because tooling for smaller models may already exist.
Results
The evaluation compares speculative and naive sampling with Chinchilla across XSum and HumanEval, including distributed hardware and varying draft lengths. Speculative sampling substantially accelerates both tasks while preserving benchmark parity, but longer drafts can increase variance and eventually reduce speedup.
- Results: Chinchilla is served on 16 TPU v4s, whereas a Chinchilla-optimal 7B model is fastest on 4 TPU v4s at 5 ms/token.Serving the 7B model on 16 TPUs increases latency because communication overhead outweighs added memory bandwidth and flops.
- Results: Speculative sampling preserves parity in benchmark metrics because its underlying sample distribution is provably the same up to hardware numerics.The authors empirically verify that the draft model does not bias results.
- Results: In HumanEval and greedy XSum, speculative sampling exceeds the theoretical memory-bandwidth limit for autoregressive sampling.The limit is defined by model size divided by total hardware memory bandwidth.
- Trade off between longer drafts and more frequent scoring: Increasing K reduces scoring calls but raises drafting and scoring time approximately linearly, while accepted-token efficiency decreases with longer lookahead.Consequently, speedup can plateau or degrade; XSum nucleus latency is minimized at K=3.
- Trade off between longer drafts and more frequent scoring: Larger K can increase variance in full-sequence generation time even when it yields marginally greater mean speedups.This matters when P90 or P99 latency is the operational concern.
Conclusion
Speculative sampling accelerates language-model decoding without modifying the target model, while preserving the target distribution within numerics. Scaled experiments with Chinchilla show large speedups across benchmark tasks and decoding methods.
- Speculative sampling requires no modifications to the target language model’s parameters or architecture.
- The method is provably lossless within numerics and empirically verifies unchanged downstream sample behavior.
- The technique scales to Chinchilla 70B using a draft model trained with existing infrastructure.
- It produces large speedups across benchmark tasks and common decoding methods.
Hyperparams
This section presents the draft-model hyperparameter table and proves that modified rejection sampling recovers the desired target distribution.
- Table 2 lists the hyperparameters used for the draft model.
- Theorem 1 considers discrete distributions q and p, with a draft sample drawn from p and a final sample X.
- The proof accounts for accepted draft samples and samples resampled after rejection using the acceptance and resampling rules.
- The rejection probability is calculated using the positive part of q(x) − p(x) and its denominator.
- The resulting procedure recovers the desired target distribution.