[Paper Review] SpecTr: Fast Speculative Decoding via Optimal Transport

Ziteng Sun, Ananda Theertha Suresh+4 more|arXiv (Cornell University)|Oct 23, 2023

Natural Language Processing Techniques4 citations

TL;DR

SpecTr proposes a novel speculative decoding method using optimal transport theory to select high-quality draft tokens, enabling faster autoregressive generation with provably no quality degradation. It achieves a 2.13x wall clock speedup over baseline autoregressive decoding and a 1.37x improvement over standard speculative decoding on state-of-the-art language models.

ABSTRACT

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 $ extit{speculative decoding}$: use a small model to sample a $ extit{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 $ extit{membership cost}$. This framework can be viewed as an extension of the well-known $ extit{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 $ extit{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.

Motivation & Objective

• To develop a principled framework for speculative decoding using optimal transport with membership cost.
• To generalize speculative decoding to allow k-token drafts at the token level, improving draft quality and acceptance rates.
• To design a draft selection algorithm with (1−1/e)-optimal acceptance probability that runs in near-linear time per token.
• To ensure the final output distribution matches that of the large language model, guaranteeing no quality degradation.
• To empirically validate the method on state-of-the-art models and demonstrate significant speedups in real-world inference.

Proposed method

• Formulates speculative decoding as an optimal transport problem with membership cost, extending the maximal-coupling framework.
• Introduces a k-token draft selection strategy that optimizes the transport plan to maximize acceptance probability.
• Solves the optimal draft selection via linear programming, though proposes a practical approximation with (1−1/e)-optimality.
• Designs a fast algorithm to compute the draft selection in time nearly linear in the vocabulary size per token.
• Integrates the draft selection into a new autoregressive sampling algorithm called SpecTr, which enables parallel scoring of draft tokens by the large model.
• Uses a statistical acceptance mechanism to ensure the final output follows the true distribution of the large model.

Experimental results

Research questions

• RQ1 Can optimal transport theory be used to provide a principled framework for speculative decoding?
• RQ2 How can speculative decoding be generalized to support k-token drafts at the token level to improve acceptance rates?
• RQ3 What is the theoretical limit of draft selection quality under the optimal transport formulation?
• RQ4 Can a near-optimal draft selection algorithm be computed efficiently, in near-linear time per token?
• RQ5 Does the proposed method achieve significant speedups in practice without degrading output quality?

Key findings

• SpecTr achieves a 2.13x wall clock speedup over baseline autoregressive decoding on the LM1B benchmark using PALM-2-Bison as the large model.
• With L=8 and K=8, SpecTr achieves a 4.0 block efficiency, significantly outperforming standard speculative decoding with K=1, which achieves only 2.9 block efficiency.
• The method provides a 1.37x further speedup over standard speculative decoding when K=1 and L=8.
• The proposed draft selection algorithm achieves (1−1/e)-optimal acceptance probability, which is the best known multiplicative guarantee for such methods.
• Empirical results confirm that the output distribution of SpecTr matches that of the large model, ensuring no quality degradation.
• The theoretical gains are largely preserved in practice, even with implementation overheads and non-negligible small model inference times.