迈向最优的多草案推测解码

Towards Optimal Multi-draft Speculative Decoding

February 26, 2025

作者: Zhengmian Hu, Tong Zheng, Vignesh Viswanathan, Ziyi Chen, Ryan A. Rossi, Yihan Wu, Dinesh Manocha, Heng Huang

cs.AI

摘要

大型语言模型（LLMs）已成为自然语言处理任务中不可或缺的一部分。然而，自回归采样已成为效率瓶颈。多草案推测解码（MDSD）是近期提出的一种方法，在生成每个词元时，小型草案模型会生成多个草案，并由目标LLM并行验证，确保最终输出符合目标模型的分布。MDSD中的两个主要设计选择是草案采样方法和验证算法。对于固定的草案采样方法，最优接受率是一个最优传输问题的解，但该问题的复杂性使得求解最优接受率并衡量现有验证算法与理论上限之间的差距变得困难。本文讨论了最优传输问题的对偶问题，提供了一种高效计算最优接受率的方法。我们首次测量了词汇量在数千级别时MDSD效率的理论上限，并量化了现有验证算法与此上限之间的差距。我们还基于最优接受率比较了不同的草案采样方法。结果表明，草案采样方法对最优接受率有显著影响，其中无放回采样优于有放回采样。此外，现有验证算法在无放回和有放回采样下均未达到理论上限。我们的研究结果表明，精心设计的草案采样方法有望提高最优接受率，并推动开发出更接近理论上限的验证算法。

English

Large Language Models (LLMs) have become an indispensable part of natural language processing tasks. However, autoregressive sampling has become an efficiency bottleneck. Multi-Draft Speculative Decoding (MDSD) is a recent approach where, when generating each token, a small draft model generates multiple drafts, and the target LLM verifies them in parallel, ensuring that the final output conforms to the target model distribution. The two main design choices in MDSD are the draft sampling method and the verification algorithm. For a fixed draft sampling method, the optimal acceptance rate is a solution to an optimal transport problem, but the complexity of this problem makes it difficult to solve for the optimal acceptance rate and measure the gap between existing verification algorithms and the theoretical upper bound. This paper discusses the dual of the optimal transport problem, providing a way to efficiently compute the optimal acceptance rate. For the first time, we measure the theoretical upper bound of MDSD efficiency for vocabulary sizes in the thousands and quantify the gap between existing verification algorithms and this bound. We also compare different draft sampling methods based on their optimal acceptance rates. Our results show that the draft sampling method strongly influences the optimal acceptance rate, with sampling without replacement outperforming sampling with replacement. Additionally, existing verification algorithms do not reach the theoretical upper bound for both without replacement and with replacement sampling. Our findings suggest that carefully designed draft sampling methods can potentially improve the optimal acceptance rate and enable the development of verification algorithms that closely match the theoretical upper bound.