SpecTr: Fast Speculative Decoding via Optimal Transport

📄 arXiv: 2310.15141v2 📥 PDF

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

分类: cs.LG, cs.CL, cs.DS, cs.IT

发布日期: 2023-10-23 (更新: 2024-01-18)

备注: NeurIPS 2023

💡 一句话要点

提出SpecTr以加速自回归解码过程

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 自回归采样 推测解码 最优传输 自然语言处理 算法优化

📋 核心要点

  1. 现有自回归采样方法逐个生成令牌，导致解码速度慢，影响实际应用。
  2. 论文提出通过最优传输理论优化推测解码，允许在令牌级别选择多个候选项，提高了采样效率。
  3. 实验结果显示，SpecTr在大型语言模型上实现了2.13倍的速度提升，相较于推测解码进一步提高了1.37倍。

📝 摘要（中文）

自回归采样在多个自然语言任务中取得了最先进的结果，但由于逐个生成令牌，速度较慢。本文提出了一种名为SpecTr的快速推测解码方法，通过小模型生成草稿，并利用大语言模型并行评分，优化了采样效率。我们通过最优传输理论提供了推测解码的原则性理解，并提出了一种新的草稿选择算法，显著提高了解码速度，同时确保输出质量不下降。实验表明，该方法在标准基准上实现了2.13倍的速度提升。

🔬 方法详解

问题定义：本文旨在解决自回归采样速度慢的问题，现有方法逐个生成令牌，导致在某些任务中不可行。

核心思路：通过引入最优传输理论，利用小模型生成草稿，并通过大模型并行评分，优化了推测解码的效率。

技术框架：整体流程包括草稿生成、并行评分和草稿选择三个主要模块。小模型生成草稿后，大模型对草稿中的令牌进行评分，最后根据评分选择最佳令牌。

关键创新：提出了一种新的草稿选择算法，基于最优传输理论，允许在令牌级别选择多个候选项，显著提高了选择的最优性。

关键设计：算法的接受概率达到(1-1/e)-最优，且计算时间几乎与单个令牌的域大小成线性关系，优化了计算效率。该方法的设计确保了输出质量与速度的平衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SpecTr在标准基准上实现了2.13倍的速度提升，相较于传统的推测解码方法，进一步提高了1.37倍，展现了其在大型语言模型中的优越性能。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过加速自回归解码，SpecTr可以在实时应用中提供更快的响应时间，提升用户体验，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.