SpecTr: Fast Speculative Decoding via Optimal Transport Article Swipe

MAKE IMPACT

Ziteng Sun , Ananda Theertha Suresh , Jae Hun Ro , Ahmad Beirami , Himanshu Jain , Felix Yu · YOU? · · 2023 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2310.15141

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.

Type

preprint

Language

en

Landing Page

http://arxiv.org/abs/2310.15141

PDF

https://arxiv.org/pdf/2310.15141

OA Status

green

Cited By

4

Related Works

10

OpenAlex ID

https://openalex.org/W4387929784

Raw OpenAlex JSON

OpenAlex ID

https://openalex.org/W4387929784

Canonical identifier for this work in OpenAlex

DOI

https://doi.org/10.48550/arxiv.2310.15141

Digital Object Identifier

Title

SpecTr: Fast Speculative Decoding via Optimal Transport

Work title

Type

preprint

OpenAlex work type

Language

en

Primary language

Publication year

2023

Year of publication

Publication date

2023-10-23

Full publication date if available

Authors

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

List of authors in order

Landing page

https://arxiv.org/abs/2310.15141

Publisher landing page

PDF URL

https://arxiv.org/pdf/2310.15141

Direct link to full text PDF

Open access

Yes

Whether a free full text is available

OA status

green

Open access status per OpenAlex

OA URL

https://arxiv.org/pdf/2310.15141

Direct OA link when available

Concepts

Decoding methods, Computer science, Security token, Speedup, Selection (genetic algorithm), Autoregressive model, Algorithm, Sampling (signal processing), Sequence (biology), Set (abstract data type), Theoretical computer science, Mathematics, Parallel computing, Filter (signal processing), Artificial intelligence, Statistics, Biology, Genetics, Computer security, Computer vision, Programming language

Top concepts (fields/topics) attached by OpenAlex

Cited by

4

Total citation count in OpenAlex

Citations by year (recent)

2025: 3, 2024: 1

Per-year citation counts (last 5 years)

Related works (count)

10

Other works algorithmically related by OpenAlex

Full payload