Fast Inference from Transformers via Speculative Decoding

Yaniv Leviathan; Matan Kalman; Yossi Matias

Source-linked AI summary

Fast Inference from Transformers via Speculative Decoding

Yaniv Leviathan, Matan Kalman, Yossi Matias

arXiv:2211.17192v2Nov 30, 2022cs.LGcs.CL

TL;DR

Large autoregressive models decode slowly because token generation is serial, motivating faster inference without changing model outputs. The paper introduces speculative decoding, which proposes tokens with an efficient model and verifies them in parallel with the target model. It reports a 2X-3X latency improvement for T5-XXL versus T5X with unchanged outputs, while identifying additional-compute availability as a key boundary.

Problem
Decoding K tokens from large autoregressive models requires K serial runs, making inference slower than decoding with smaller models.
Method
Speculative decoding uses a more efficient approximation model to propose tokens and speculative sampling to verify or correct them while preserving the target distribution.
Results
2X-3X latency improvement was demonstrated for T5-XXL versus the T5X implementation, without changing outputs.
Takeaways & Limitations
The method can accelerate off-the-shelf autoregressive models without architecture changes, retraining, or altered output distributions.
Takeaways & Limitations
The method requires additional computation resources and is not helpful when those resources are unavailable.

Abstract

from arXiv · show

Inference from large autoregressive models like Transformers is slow - decoding K tokens takes K serial runs of the model. In this work we introduce speculative decoding - an algorithm to sample from autoregressive models faster without any changes to the outputs, by computing several tokens in parallel. At the heart of our approach lie the observations that (1) hard language-modeling tasks often include easier subtasks that can be approximated well by more efficient models, and (2) using speculative execution and a novel sampling method, we can make exact decoding from the large models faster, by running them in parallel on the outputs of the approximation models, potentially generating several tokens concurrently, and without changing the distribution. Our method can accelerate existing off-the-shelf models without retraining or architecture changes. We demonstrate it on T5-XXL and show a 2X-3X acceleration compared to the standard T5X implementation, with identical outputs.

1. Introduction

Large autoregressive models are capable but slow because decoding remains serial. Speculative decoding uses efficient approximations and parallel verification to accelerate exact sampling without changing architectures, training, or output distributions.

Decoding K tokens from large autoregressive models requires K serial model runs, making inference substantially slower than smaller-model inference.
Speculative decoding uses an approximation model to propose tokens and a target model to verify them in parallel.The method generalizes speculative execution to stochastic sampling through speculative sampling.
2X-3X walltime improvement was achieved for T5-XXL versus the T5X implementation, with unchanged outputs.Experiments also covered GPT-like unconditional generation, translation, summarization, and dialog models.
The approach is positioned for production use when memory bandwidth limits inference and additional compute resources are available.
The method accelerates decoding without changing model architectures, training procedures, retraining requirements, or output distributions.

2. Speculative Decoding

Speculative decoding samples several candidate continuations with a smaller model, then evaluates them concurrently with the target model. Accepted candidates are retained, while rejected candidates are corrected using an adjusted distribution that preserves exact target-model sampling.

2.1. Overview: The approximation model Mq generates γ candidate completions, and the target model Mp evaluates their probabilities in parallel.
2.1. Overview: Each target-model run produces at least one new token and can produce up to γ + 1 tokens when the approximation is sufficiently accurate.
2.2. Standardized Sampling: Sampling from q(x) accepts a candidate when q(x) ≤ p(x); otherwise, rejection is followed by resampling from p′(x) = norm(max(0, p(x) − q(x))).
2.1. Overview: When a proposed token is rejected, its downstream speculative computation is discarded and the token is resampled from the adjusted target distribution.
2.1. Overview: The algorithm samples γ guesses autoregressively, runs Mp on all resulting prefixes in parallel, determines the accepted prefix length, and returns corrected output.

3. Analysis

Speculative decoding reduces serial target-model calls by using an approximation model to propose tokens and the target model to verify them in parallel. Its gains depend on acceptance rate, approximation cost, concurrency, and available compute resources.

3.1. Number of Generated Tokens: The acceptance rate β measures how often speculative samples are accepted, and under an i.i.d. assumption α = E(β) determines the expected tokens generated per algorithm run.The generated-token count is modeled as a capped geometric variable with cap γ + 1.
3.1. Number of Generated Tokens: Speculative decoding reduces target-model calls by a factor of 1−α^(γ+1).The reduction follows from the expected number of tokens produced in one run of Algorithm 1.
3.3. Walltime Improvement: The approximation model’s cost coefficient c depends on hardware and software, while in experiments it was below 0.05 and often nearly zero.The acceptance rate α is described as intrinsic to the models and task, unlike c.
3.3. Walltime Improvement: The expected walltime improvement factor is (1−α^(γ+1))/((1−α)(γc+1)), where c is the approximation-to-target runtime ratio.The analysis assumes γ + 1 concurrent target-model evaluations can run without increasing walltime and that generations are long enough.
3.4. Arithmetic Operations: Algorithm 1 increases concurrent arithmetic operations by γ + 1, while rejected guesses can increase total arithmetic work.Accepted samples can make the additional computation effectively free, but rejected guesses waste computation.
3.5. Choosing γ: For a given c and α, the optimal integer γ maximizes the walltime improvement factor and can be found numerically.Figure 3 reports the optimal γ as a function of α for different c values.
3.6. Approximation Models: With negligible-cost approximations, c ≈ 0, and a trivial bigram model achieved α ≈ 0.2 and a 1.25X speed improvement on English-German translation.The passage also identifies n-gram models as negligible-cost approximations because evaluation is a table lookup.

4. Experiments

Experiments evaluate speculative decoding across model sizes, tasks, and sampling settings, finding substantial T5-XXL walltime improvements and non-negligible acceptance from small approximations.

T5-XXL walltime results: T5-small (77M) provided the highest speedup among the tested approximation models by balancing computational cost and acceptance probability.The experiments used T5-large, T5-base, and T5-small checkpoints with T5-XXL (11B) as the target.
T5-XXL walltime results: 2.6X and 3.4X speedups were observed for translation with standard sampling and argmax sampling, respectively.Summarization achieved 2.3X and 3.1X speedups under the same settings.
Sampling and approximation effects: Approximation-model acceptance probability increased with model size, while argmax sampling produced higher acceptance probabilities and walltime improvements than standard sampling.The empirical results broadly matched theoretical predictions, with variance attributed to implementation details.

5. Related work

Related work includes methods that reduce inference cost, adapt computation, or use speculative execution; speculative decoding extends these ideas to stochastic sampling without retraining or output-distribution changes.

Inference acceleration: Prior efficiency methods include distillation, sparsification, quantization, and architecture modification applied broadly across inference tokens.These approaches generally target inference efficiency for all inputs or tokens.
Speculative execution: Blockwise Parallel Decoding supports greedy decoding but requires additional training and does not address the general stochastic setting.Its focus is preserving downstream task quality rather than guaranteeing the original output distribution.
Speculative execution: SAD copies the input rather than using general approximation models, limiting it to tasks where inputs and outputs are similar.SAD also does not support general stochastic sampling.
Subsequent work: An independent implementation later reported 2X-2.5X improvements on Chinchilla 70B.This provides an external implementation result related to speculative decoding.

6. Discussion

The discussion presents speculative decoding as a practical, distribution-preserving acceleration method, while identifying compute availability and several untested extensions as important boundaries.

Conclusions: Speculative sampling enables stochastic speculative execution, yielding practical 2X-3X speedups versus T5X when sufficient compute resources are available.The method preserves the target model’s output distribution while accelerating autoregressive decoding.
Practical scope: The method is most suitable when memory bandwidth is the bottleneck and additional arithmetic computation is available.In that setting, it can use out-of-the-box models without architecture changes or retraining.
Limitation: Speculative execution increases concurrency at the cost of additional arithmetic operations, limiting usefulness when extra computation resources are unavailable.This is identified as a general limitation of speculative execution and speculative decoding.
Future directions: Future work includes testing beam-search compatibility, custom approximation models, hierarchical acceleration, and varying approximation models or guess counts during inference.The paper also suggests alternative transformations and applications beyond text.
Broader applications: The authors suggest stochastic speculative execution may apply beyond autoregressive decoding, including physics simulations and reinforcement learning.The proposed pattern applies when one slow function generates a distribution from which another function samples its input.

A.1. Correctness of Speculative Sampling

Speculative sampling is analyzed for arbitrary target and approximation distributions by separating accepted guesses from a residual adjusted distribution.

Correctness claim: For any distributions p(x) and q(x), speculative sampling produces tokens distributed identically to samples from p(x) alone.The proof introduces β as the acceptance probability.
Residual distribution: The adjusted distribution p′(x) is formed by normalizing the nonnegative residual max(0, p(x) − q(x)).This residual accounts for probability mass not covered by the approximation distribution.
Residual distribution: The residual distribution has normalizing constant 1 − β.The normalization follows from the stated lemma and theorem.
Proof structure: The output probability is decomposed into cases where the approximation guess is accepted or rejected.This case split supports the distribution-preservation argument.

A.2. Speculative Sampling vs. Rejection Sampling

Speculative sampling preserves the target distribution while improving on a non-iterative rejection-sampling alternative. Its expected acceptance probability is higher than sampling from q(x) and accepting directly against p(x).

The non-iterative alternative samples x from q(x), accepts it using p(x)/Mq(x), and otherwise samples directly from p(x).
Speculative sampling produces tokens distributed identically to samples from p(x) alone.
The alternative’s expected accept probability, E_x∼q(x) min(p(x), q(x)), can be much lower than speculative sampling’s α.

A.3. Theoretical Predictions vs. Empirical Runtimes

Theoretical runtime predictions largely agree with measured runtimes. Remaining discrepancies are attributed to implementation optimization differences and an approximate i.i.d. assumption for the βs.

Theoretical expected runtime improvements mostly match empirically measured runtimes.The comparison uses Theorem 3.8 predictions and runtimes measured in Table 2.
Larger prediction differences reflect optimization differences between the implementation and baseline.
The simplifying assumption that the βs are i.i.d. is only an approximation and also contributes to discrepancies.

A.4. Application to Beam Search

Speculative decoding extends to beam search by generating candidates with an approximation model and checking them in parallel against the target model. Lenience can trade exactness for additional speed while bounding sampled-token probabilities.

Application to Beam Search: Beam search uses approximation-model candidates of width u ≥ w for γ steps, then checks them in parallel with the target model.The target-model compute budget is (w + uγ) runs, and guesses are accepted when topw(Mp) ⊆ topu(Mq).
Application to Beam Search: The beam-search procedure can produce identical results to regular beam search with Mp alone.
Lenience: Lenience preserves an upper bound on sampled-token probabilities but can reduce sample diversity because it provides no minimum-probability guarantee.
Lenience: 2.5X, 3.1X, 3.6X, and 5X improvement factors result from lenience values of 1, 0.5, 0.3, and 0.1, respectively, for T5-XXL with T5-small.
Lenience: 3.3X, 3.3X, 3.9X, and 4.9X speed improvement factors result for lenience values of 1, 0.5, 0.3, and 0.1 under argmax sampling.The setting uses T5-XXL with T5-small for English-German translation, c = 0.015, and γ = 8.