Towards Optimal Multi-draft Speculative Decoding
Authors: name: Zhengmian Hu, name: Tong Zheng, name: Vignesh Viswanathan, name: Ziyi Chen, name: Ryan A. Rossi, name: Yihan Wu, name: Dinesh Manocha, name: Heng Huang

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.