最適なマルチドラフト推測デコーディングに向けて

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

要旨

大規模言語モデル（LLM）は、自然言語処理タスクにおいて不可欠な存在となっています。しかし、自己回帰的なサンプリングが効率のボトルネックとなっています。最近提案されたマルチドラフト推測的デコーディング（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.