[论文解读] SpecTr: Fast Speculative Decoding via Optimal Transport

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

Natural Language Processing Techniques被引用 4

一句话总结

SpecTr 提出了一种新颖的推测解码方法，利用最优传输理论选择高质量的草稿标记，实现更快的自回归生成，且可证明无质量下降。该方法在最先进语言模型上相较基线自回归解码实现 2.13 倍的时钟周期加速，相较标准推测解码实现 1.37 倍的性能提升。

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.

研究动机与目标

• 提出一种基于最优传输与成员成本的推测解码原理性框架。
• 将推测解码推广至支持逐标记的 k 个标记草稿，以提升草稿质量和接受率。
• 设计一种草稿选择算法，其接受概率达到 (1−1/e)-最优，且每标记运行时间接近线性。
• 确保最终输出分布与大语言模型一致，从而保证无质量下降。
• 在最先进模型上实证验证该方法，并在真实推理中展示显著加速。

提出的方法

• 将推测解码建模为带有成员成本的最优传输问题，扩展了最大耦合框架。
• 引入一种 k 个标记草稿选择策略，通过优化传输计划以最大化接受概率。
• 通过线性规划求解最优草稿选择，但提出一种具有 (1−1/e)-最优性的实用近似方法。
• 设计一种快速算法，使每标记的草稿选择计算时间接近词汇表大小的线性时间。
• 将草稿选择集成到一种新型自回归采样算法 SpecTr 中，使大模型能并行评分草稿标记。
• 采用统计接受机制，确保最终输出遵循大模型的真实分布。

实验结果

研究问题

• RQ1 最优传输理论能否为推测解码提供一个原理性框架？
• RQ2 如何将推测解码推广至支持逐标记的 k 个标记草稿，以提升接受率？
• RQ3 在最优传输框架下，草稿选择质量的理论极限是什么？
• RQ4 能否高效地计算出接近最优的草稿选择算法，使每标记的运行时间接近线性？
• RQ5 所提出的方法在实践中是否实现了显著加速，且不降低输出质量？

主要发现

• 在 LM1B 基准上，使用 PALM-2-Bison 作为大模型时，SpecTr 相较基线自回归解码实现 2.13 倍的时钟周期加速。
• 当 L=8 且 K=8 时，SpecTr 实现 4.0 的块效率，显著优于标准推测解码（K=1 时仅 2.9 的块效率）。
• 当 K=1 且 L=8 时，该方法相较标准推测解码实现 1.37 倍的进一步加速。
• 所提出的草稿选择算法实现了 (1−1/e)-最优接受概率，这是此类方法中已知的最佳乘法保证。
• 实证结果证实，SpecTr 的输出分布与大模型一致，确保无质量下降。
• 理论上的性能增益在实践中基本得以保留，即使存在实现开销和不可忽略的小模型推理时间。