PaperlayerView exact revision on arXiv

Source-linked AI summary

SpecTr: Fast Speculative Decoding via Optimal Transport

Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, Felix Yu

arXiv:2310.15141v2cs.LGcs.CLcs.DScs.IT

TL;DR

Autoregressive decoding is slow because tokens are generated serially. This paper frames speculative decoding as optimal transport with membership cost, extends selection to multiple draft candidates, and proposes SpecTr, which preserves the large model’s output distribution while achieving measured speedups.

  • Problem

    Serial autoregressive generation makes inference slow or prohibitive, motivating principled methods that improve decoding speed without degrading the large model’s output distribution.

  • Method

    The paper formulates token-level draft selection as optimal transport with membership cost, permits k draft candidates, and develops the SpecTr autoregressive sampling algorithm.

  • Results

    SpecTr achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

  • Takeaways & Limitations

    SpecTr provides faster decoding while deriving valid sequences whose distribution follows the large model.

  • Takeaways & Limitations

    The optimal transport plan has best-known runtime exponential in k, and the computational model abstracts hardware overheads despite experimental verification of largely preserved gains.

Abstract

from arXiv · show

Autoregressive sampling from large language models has led to state-of-the-art results in several natural language tasks. However, autoregressive sampling generates tokens one at a time making it slow, and even prohibitive in certain tasks. One way to speed up sampling is $\textit{speculative decoding}$: use a small model to sample a $\textit{draft}$ (block or sequence of tokens), and then score all tokens in the draft by the large language model in parallel. A subset of the tokens in the draft are accepted (and the rest rejected) based on a statistical method to guarantee that the final output follows the distribution of the large model. In this work, we provide a principled understanding of speculative decoding through the lens of optimal transport (OT) with $\textit{membership cost}$. This framework can be viewed as an extension of the well-known $\textit{maximal-coupling}$ problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of $k$ candidates at the token-level, which leads to an improved optimal membership cost. We show that the optimal draft selection algorithm (transport plan) can be computed via linear programming, whose best-known runtime is exponential in $k$. We then propose a valid draft selection algorithm whose acceptance probability is $(1-1/e)$-optimal multiplicatively. Moreover, it can be computed in time almost linear with size of domain of a single token. Using this $new draft selection$ algorithm, we develop a new autoregressive sampling algorithm called $\textit{SpecTr}$, which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

1 Introduction

Autoregressive language models generate tokens serially, creating latency that can make inference slow or prohibitive. The paper motivates speculative decoding and develops a computational framework for analyzing parallelized inference.

  • Autoregressive decoding generates tokens one after another, which can make inference slow or prohibitive in some applications.
  • The computational model treats time-axis and batch-axis parallelization as taking O(1) time, while computation scales with the number of contexts or sequence length.The paper notes that real hardware introduces overhead, while experiments indicate theoretical gains are largely preserved for a large transformer.
  • Speculative decoding uses an efficient model to guess future tokens, evaluates corresponding large-model distributions in parallel, and statistically accepts or rejects the guesses.
  • Accepted speculative tokens are designed to remain distributed according to the large model, avoiding degradation in decoded-output quality.
  • The work aims to provide a principled understanding of speculative decoding, including its optimality conditions and algorithmic improvements.

2 Previous works and speculative decoding

Speculative decoding drafts multiple tokens with a small model, scores them with the large model in parallel, and selects a valid prefix while preserving the large model’s distribution. Its central selection step recursively applies token-level maximal coupling and corrects the first rejection using a residual sample.

  • Speculative decoding constructs a length-L draft, computes large-model conditional distributions for all draft positions in parallel, and selects an accepted prefix.
  • Draft selection applies token-level maximal coupling recursively, using the small-model and large-model conditional distributions to select a valid sequence.
  • At the first rejected draft token, the algorithm samples a correction from the residual distribution; otherwise it returns the accepted draft token.
  • The resulting output has the large model’s distribution when the draft sample is drawn from the draft distribution.
  • The probability of accepting a draft token is 1 − dTV(p, q), so closer draft and large-model distributions reduce serial calls to the large model.
  • When p = q, all L draft tokens are accepted and one iteration produces L + 1 tokens, yielding an idealized speedup of L + 1.

3 Our contributions

The paper answers theoretical questions about speculative decoding by connecting it to optimal transport and extending draft selection to multiple candidates. It introduces SpecTr, which preserves output quality while achieving substantial decoding speedups.

  • A multi-draft formulation enables batch-axis parallelization, addressing whether multiple draft sequences can further accelerate decoding.The approach uses multiple drafts from the small model and a new selection algorithm to obtain valid sequences.
  • The work connects speculative decoding to discrete optimal transport and maximal coupling, identifying token-level draft selection as an optimal transport solution.The formulation uses an indicator cost and provides a principled basis for analyzing speculative decoding.
  • The proposed transport plan is efficiently computable and achieves a (1 −1/e)-approximation of the optimal acceptance probability.The exact linear-program formulation has exponentially many variables in batch size, motivating the efficient approximation.
  • SpecTr recursively selects valid tokens from multiple draft sequences and generates four tokens in the illustrated L = 4 example.The example selects ‘be’, ‘liked’, and ‘by’, then produces ‘three’ as a correction when the final draft is rejected.
  • SpecTr provably avoids degradation in decoded-output quality relative to the large model while improving decoding speed.The quality guarantee parallels speculative decoding, and the method is designed to derive valid sequences from the large model.
  • 2.13X wall-clock speedup was achieved for state-of-the-art large language models, including a further 1.37X speedup over speculative decoding.These results were demonstrated on standard benchmarks.

4 Token-level draft selection and optimal transport

The paper formulates token-level draft selection as a coupling problem between draft and target distributions, requiring valid large-model samples while maximizing draft acceptance. For one draft, speculative decoding is optimal under an indicator resampling cost; multiple drafts expose why naive repeated acceptance fails.

  • Draft-selection objectives: The selection algorithm maps draft tokens to an output token whose distribution matches the large model while maximizing acceptance.Validity preserves the large model’s output distribution, while acceptance saves serial computation through parallelization.
  • Multiple-draft challenge: Applying one-draft acceptance independently to multiple drafts can produce an invalid output distribution.For p = Ber(1) and q = Ber(1/2), repeated acceptance makes the probability of outputting 1 at least 1 − 1/2^k for k > 1.
  • Coupling formulation: A coupling between P and Q is a joint distribution whose marginals are P and Q, and its conditional mapping defines a valid draft-selection algorithm.If X follows P, the mapped output follows Q.
  • Optimal transport formulation: Optimal transport selects the coupling minimizing a specified transportation cost, which here represents draft-token resampling.The indicator cost penalizes outputs Y different from the draft token X.
  • One-draft case: For one draft token, maximal coupling achieves the optimal indicator-cost transport plan, so speculative sampling has optimal cost.The token domain is Ω, with the draft and target distributions given by p and q.

5 Optimal transport with multiple draft tokens

The paper extends draft selection to k candidate tokens using optimal transport with membership cost, which measures whether the output belongs to the candidate set. The exact optimum is computable by linear programming but has exponential dependence on k, motivating efficient approximation.

  • Membership-cost formulation: With k draft tokens, membership cost assigns zero cost when the output belongs to the draft set and one otherwise.For a draft tuple x, S(x) is the set of distinct tokens appearing in x.
  • Multiple-draft objective: The framework uses k i.i.d. samples from p to produce one valid token from q, with optimal acceptance probability α_k(p, q).For SpecTr, p is the small-model conditional distribution and q is the large-model conditional distribution.
  • Acceptance properties: Acceptance probability is monotone in the number of drafts, and it converges to one as k grows when q(x)/p(x) is bounded.The stated consistency property gives the convergence condition.
  • Exact optimization: The exact OTM-k plan can be computed in time O(|Ω|^O(k)), reflecting the exponential dependence on the number of drafts.The associated linear program has |Ω|^(k+1) variables and |Ω|^k + |Ω| equality constraints.
  • Computational limitation: The optimal acceptance bound is tight up to the multiplicative factor 1 − (1 − 1/k)^k ≥ 1 − 1/e, while exact optimization can be prohibitive for large vocabulary or k.The paper notes that even entropy-regularized alternatives retain linear dependence on |Ω|^k.

6 Draft selection via k-sequential selection

The paper introduces K-SEQ, an efficient sequential approximation to the multiple-draft transport problem. It processes drafts in order using likelihood ratios, preserves validity through a tunable parameter, and attains a multiplicative acceptance guarantee.

  • Algorithm: K-SEQ processes k draft samples sequentially and accepts the first one according to a likelihood-ratio rule, otherwise sampling from a residual distribution.The algorithm compares q(X_i)/p(X_i) and returns a corrected sample if no draft is accepted.
  • Validity mechanism: K-SEQ uses an appropriate ρ and accepts X_i with probability min(1, q(X_i)/(ρ·p(X_i))) to remain a valid transport plan.Optimizing within this one-parameter family substantially reduces computation while preserving the target distribution.
  • Guarantee: K-SEQ achieves acceptance at least (1 − e^-1)α_k(p, q) when ρ = ρ*, giving a (1 − 1/e)-multiplicative guarantee relative to optimum.The parameter ρ* is defined through the paper’s stated identity involving β_p,q(ρ).
  • Computation: The parameter ρ* can be computed to accuracy δ in time O(|Ω| log((k − 1)/δ)).This computation is nearly linear in vocabulary size for fixed accuracy and draft count.
  • Scope of approximation: K-SEQ can have strictly lower acceptance than optimal OTM in some distributions, although it is optimal in other non-trivial cases.The paper reports this comparison for canonical examples in its appendix.

7 SpecTr: Application of OTM in autoregressive sampling

SpecTr applies optimal-transport-based draft selection to multiple sampled sequences, recursively selecting valid tokens while preserving the large model’s distribution. Its implementation uses parallel conditional-probability computation and supports an extra decoded token when the final draft token is accepted.

  • Three phases: SpecTr constructs a draft set, computes large-model conditional probabilities in parallel, and selects the decoded prefix.The draft set contains K sequences of length L, with small- and large-model probabilities retained for candidate prefixes.
  • Draft construction: The draft set can be formed by independently sampling K sequences from the small model, while a prefix-tree generalization did not improve experimental performance.The prefix-tree construction is included for completeness rather than used as a demonstrated improvement.
  • Draft selection: SpecTr selects valid tokens from multiple draft sequences using a token-level transport plan from the small model to the large model.The plan is recursively applied to candidate prefixes, filtering sequences that do not match the selected token.
  • Validity and output length: The decoded token is always a valid sample from the large model, and each iteration decodes a random number of tokens between 1 and L+1.When the final token is accepted, SpecTr can sample one additional token from the already computed large-model conditional distribution.
  • Quality guarantee: SpecTr’s quality guarantee is formalized in Theorem 2 under the assumption that all drafts are generated from the small model with the current context.The theorem states that the newly decoded sequence is distributed according to the large model’s conditional distribution.

8 Experiments

The experiments evaluate SpecTr against autoregressive decoding and speculative decoding while measuring speed because all methods sample from the same large-model distribution. On PALM-2 models and LM1B prompts, SpecTr reaches a 2.13x wall-clock speedup and improves over speculative decoding.

  • Experimental setup: The evaluation compares SpecTr with baseline autoregressive decoding and speculative decoding with K = 1, focusing on speedup because output quality is distributionally neutral.All three methods effectively generate samples from the same baseline large model.
  • Deployment considerations: Wall-clock speedup is affected by draft-model time, batch and time parallelization, and implementation overhead.These factors can make actual end-to-end speedup lower than theoretical block efficiency.
  • Experimental setup: The LM1B experiment uses PALM-2-Gecko as the small model and PALM-2-Bison as the large model, averaging results over 1000 prompts and 3 random seeds.The wall-clock speedup is normalized by baseline autoregressive-decoding latency.
  • Results: 2.13x wall-clock speedup is achieved with L = 8 and K = 8, a further 1.37x increase over speculative decoding with K = 1.The measured benefit remains after accounting for the listed deployment overheads.
  • Acceptance analysis: The theoretical acceptance bound is tight up to a multiplicative factor of 1 − (1 − 1/k)^k ≥ 1 − 1/e.For k = 1, the upper bound matches the optimal acceptance probability; illustrative examples also attain the bound.

A.2 Proof of Lemma 1

The proof establishes properties of the optimal acceptance probability by comparing transport plans across draft-set sizes and handling both bounded and unbounded likelihood-ratio cases. It then uses the efficient K-SEQ guarantee to obtain the stated lower bound.

  • Monotonicity: The proof first establishes monotonicity by extending a transport plan with an independent sample from p.This constructs a valid plan for one additional draft candidate.
  • Consistency: For bounded q(x)/p(x), the proof uses ρmax = maxx∈Ω q(x)/p(x) to establish consistency of the acceptance probability.The resulting relation implies limk→∞ αk(p, q) = 1.
  • Unbounded likelihood ratios: For unbounded q(x)/p(x), the proof identifies an outcome with q(x) > 0 and p(x) = 0 and modifies q to obtain a bounded-support comparison.The construction transfers the off-support mass to an outcome x0 with p(x0) > 0.
  • Upper bound: The upper-bound argument considers every transport plan and every subset Ω0 of the token domain.The proof uses this general subset condition to control the transport cost.
  • Lower bound: K-SEQ achieves an acceptance probability at least (1 − (1 − 1/k)^k) ᾱk(p, q), yielding the lower-bound guarantee.The proof compares this efficient construction with the optimal acceptance quantity.

B Comparison between α(π

The comparisons show how acceptance probability changes with the number of candidates k for Bernoulli and uniform distributions. OTM-k improves on K-SEQ in some Bernoulli settings, while K-SEQ is optimal for the uniform example and both methods approach acceptance probability one as k grows.

  • Bernoulli distributions: For Bernoulli distributions, both OTM-k and K-SEQ reach acceptance probability one when p = q, and their acceptance probabilities increase with k when p ≠ q.When b = 0.25, p = q and both methods always achieve one.
  • Bernoulli distributions: When b = 0.1 or 0.75, K-SEQ has lower acceptance probability than OTM-k, whereas the methods perform equally when b = 1.The comparison is made with p = Ber(0.25) and q = Ber(b).
  • Uniform distributions: For uniform distributions, K-SEQ achieves the optimal acceptance probability α_k, which increases monotonically and tends to one as k approaches infinity.The plotted setting uses p = U(d) with d = 120 and varying uniform q distributions.
  • Acceptance calculations: The transport-plan calculations construct probabilities using minima such as min(p^k, q) and min((1 − p)^k, 1 − q) to match the optimal acceptance expression.The resulting expression is evaluated across cases determined by p^k relative to q and (1 − p)^k relative to 1 − q.
  • Uniform distributions: For p = U(d) and q = U(d/r), the optimal acceptance probability is 1 − (1 − 1/r)^k.The bound is achieved by K-SEQ in this example.
  • Comparison with MULTI-ROUND: Compared with MULTI-ROUND, both OTM and K-SEQ perform better on the considered Bernoulli examples, while OTM is optimal under the validity guarantee.The authors leave systematic comparison of the algorithms for future work.

C.1 Proof of Theorem 1

The proof establishes that K-SEQ is computationally efficient, valid, and near-optimal for token-level draft selection, then uses induction to extend distributional validity to the decoding procedure and motivates structured draft sets.

  • Existence of ρ*: The proof establishes existence and uniqueness of ρ* by showing that f(ρ) is decreasing and changes sign over the relevant interval.The argument uses β_p,q(ρ)'s monotonicity and endpoint inequalities.
  • Computation efficiency: The root ρ* is computed to arbitrary accuracy δ in time O(|Ω| log((k − 1)/δ)) using binary search over [1, k].This follows from the characterization of ρ* and supports efficient computation of the K-SEQ parameters.
  • Validity: For ρ ≥ ρ*, the constructed transport plan is valid because the acceptance and rejection probabilities form a valid residual distribution.The proof then verifies that the output marginal is q.
  • Acceptance probability: K-SEQ achieves an acceptance probability at least 1 − (1 − β_p,q(ρ*))^k and at least the benchmark c_k(p, q).The guarantee uses monotonicity of β_p,q(ρ) and the upper bound from Lemma 3.
  • Distributional validity: Induction shows that the decoding procedure preserves the large model's conditional distribution at each output position and sequence length.The base case uses a valid transport plan, and the induction step applies the same argument after conditioning on the first output token.
  • Draft-set construction: Because hardware limits the draft-set size, the paper also constructs prefix-tree draft sets whose children extend parent prefixes using samples from the small model.The tree is built level by level with specified expansion factors and returns the leaves as the draft set.

E Additional experiments

Additional experiments evaluate SpecTr with transformer models of different sizes, verify the assumed parallelization behavior, and examine how draft-model size affects block efficiency.

  • Experimental setup: The experiments train decoder-only transformers on LM1B, using 2M, 6M, and 20M-parameter draft models with a 97M-parameter large model.Results are reported over 1,000 test prompts, three random seeds, and sampling temperature 1.0 for both models.
  • Computational model: The large-model latency stays roughly constant when computation is parallelized over the time and batch axes.This experiment verifies the simplified computational model used in the paper.
  • Computational model: The 6M draft-model measurements report latency relative to the 97M large model to quantify the relative cost of sampling multiple drafts.The comparison uses parallelization across the time and batch axes.
  • Draft-model size: Larger draft models improve block efficiency for all speculative methods, with SpecTr maintaining the best performance across the tested draft-model sizes.The comparison uses the same 97M large model and varying draft-model architectures.
Loading 2310.15141v2…