Towards Optimal Multi-draft Speculative Decoding

2025-02-26

AI Summary

This work addresses the fundamental inefficiency of autoregressive decoding in large language models (LLMs). We propose the Multi-Draft Speculative Decoding (MDSD) optimal design framework. Methodologically, we formulate the MDSD acceptance rate as an optimal transport problem and derive its theoretical upper bound via the dual formulation—enabling exact quantification even over vocabularies spanning thousands of tokens. We further demonstrate that draft token sampling without replacement substantially outperforms sampling with replacement, and systematically benchmark mainstream verification algorithms against the derived theoretical limit. Our primary contribution is establishing a rigorous theoretical analysis paradigm for MDSD efficiency: we prove that all existing algorithms fall short of the optimum, thereby providing interpretable, computationally tractable guidance for designing high-acceptance-rate draft sampling strategies and verification algorithms that approach the theoretical bound.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Sampling/Simulation-based Search
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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 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.

Problem

• Optimize multi-draft speculative decoding efficiency
• Compute optimal acceptance rate via duality
• Compare draft sampling methods' impact on acceptance

Innovation

• Dual optimal transport problem solution
• Multi-Draft Speculative Decoding efficiency
• Draft sampling without replacement outperforms