Towards Optimal Multi-draft Speculative Decoding

📄 arXiv: 2502.18779v1 📥 PDF

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

分类: cs.CL, cs.DS

发布日期: 2025-02-26

💡 一句话要点

提出基于最优传输理论的多Draft推测解码效率分析与优化方法

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

关键词: 大型语言模型 推测解码 最优传输 效率优化 draft采样 验证算法

📋 核心要点

  1. 自回归采样是LLM的效率瓶颈，多Draft推测解码(MDSD)通过并行验证多个draft来加速推理。
  2. 论文利用最优传输理论的对偶性，高效计算MDSD的最优接受率，从而分析现有算法的效率。
  3. 实验表明，draft采样方法显著影响最优接受率，且现有验证算法与理论上限存在差距。

📝 摘要（中文）

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

🔬 方法详解

问题定义：论文旨在解决多Draft推测解码（MDSD）中，如何确定最优的draft接受率，以及如何评估现有验证算法与理论最优性能之间的差距的问题。现有方法难以有效计算最优接受率，也无法准确衡量现有算法的效率损失。

核心思路：论文的核心思路是将最优接受率问题转化为最优传输问题的对偶问题。通过求解对偶问题，可以高效地计算出最优接受率。同时，利用计算出的最优接受率，可以评估不同draft采样方法（如放回采样和不放回采样）的性能，并量化现有验证算法与理论上限之间的差距。

技术框架：论文的技术框架主要包括以下几个部分：1) 将MDSD的最优接受率问题建模为最优传输问题；2) 利用最优传输理论的对偶性，将原问题转化为对偶问题；3) 设计算法高效求解对偶问题，得到最优接受率；4) 基于最优接受率，评估不同draft采样方法的性能；5) 量化现有验证算法与理论上限之间的差距。

关键创新：论文最重要的技术创新点在于利用最优传输理论的对偶性，为MDSD的最优接受率计算提供了一种高效的解决方案。此前，由于最优传输问题的复杂性，难以直接求解最优接受率。通过对偶转换，问题变得更容易求解，从而可以对MDSD的性能进行更深入的分析和优化。

关键设计：论文的关键设计包括：1) 精确地将MDSD问题映射到最优传输框架；2) 选择合适的对偶变量和约束条件，使得对偶问题易于求解；3) 设计高效的数值算法来求解对偶问题，例如使用线性规划或凸优化方法；4) 设计实验来比较不同draft采样方法和验证算法的性能，并分析其与理论上限的差距。

🖼️ 关键图片

img_0

📊 实验亮点

论文首次测量了数千词汇量下MDSD效率的理论上限，并量化了现有验证算法与该上限之间的差距。实验结果表明，不放回采样优于放回采样，且现有验证算法与理论上限之间存在显著差距。这些发现为未来MDSD算法的设计提供了重要的指导。

🎯 应用场景

该研究成果可应用于各种需要加速LLM推理的场景，例如在线对话系统、机器翻译、文本生成等。通过优化draft采样方法和验证算法，可以显著提高LLM的推理效率，降低计算成本，并提升用户体验。未来的研究可以进一步探索更高效的draft采样策略和验证算法，以充分发挥MDSD的潜力。

📄 摘要（原文）

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.