Towards Optimal Multi-draft Speculative Decoding
Zhengmian Hu, Tong Zheng, Vignesh Viswanathan, Ziyi Chen, Ryan A. Rossi, Yihan Wu, Dinesh Manocha, Heng Huang
2025-02-27
Summary
This paper talks about improving a method called Multi-Draft Speculative Decoding (MDSD) that helps large language models (LLMs) generate text faster and more efficiently.
What's the problem?
Large language models are great at processing and generating text, but they can be slow because they generate one word at a time. MDSD tries to speed this up by using a smaller, faster model to guess multiple words at once, but it's tricky to figure out the best way to do this guessing and checking process.
What's the solution?
The researchers found a way to calculate the best possible rate at which the guesses from the smaller model should be accepted. They compared different methods of making these guesses and showed that guessing without repeating words works better than allowing repeats. They also measured how close current methods are to being as efficient as theoretically possible.
Why it matters?
This research matters because it could make AI language models work much faster without losing quality. Faster language models could improve many applications we use every day, like chatbots, translation services, and writing assistants. By showing where there's room for improvement, this study also guides future research to make these systems even more efficient.
Abstract
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 <PRE_TAG>acceptance rate</POST_TAG> is a solution to an optimal transport problem, but the complexity of this problem makes it difficult to solve for the optimal <PRE_TAG>acceptance rate</POST_TAG> and measure the gap between existing <PRE_TAG>verification algorithms</POST_TAG> and the theoretical upper bound. This paper discusses the dual of the optimal transport problem, providing a way to efficiently compute the optimal <PRE_TAG>acceptance rate</POST_TAG>. 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 <PRE_TAG>verification algorithms</POST_TAG> and this bound. We also compare different draft sampling methods based on their optimal <PRE_TAG>acceptance rate</POST_TAG>s. Our results show that the draft sampling method strongly influences the optimal <PRE_TAG>acceptance rate</POST_TAG>, with sampling without replacement outperforming sampling with replacement. Additionally, existing <PRE_TAG>verification algorithms</POST_TAG> 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 <PRE_TAG>verification algorithms</POST_TAG> that closely match the theoretical upper bound.