邁向最優的多草稿推測解碼

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.