Auf dem Weg zu optimaler Multi-Entwurfs-Spekulativdekodierung

Towards Optimal Multi-draft Speculative Decoding

February 26, 2025

Autoren: Zhengmian Hu, Tong Zheng, Vignesh Viswanathan, Ziyi Chen, Ryan A. Rossi, Yihan Wu, Dinesh Manocha, Heng Huang

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) sind zu einem unverzichtbaren Bestandteil von Aufgaben der natürlichen Sprachverarbeitung geworden. Allerdings ist die autoregressive Abtastung zu einem Effizienzengpass geworden. Das Multi-Draft Spekulative Dekodieren (MDSD) ist ein neuer Ansatz, bei dem bei der Generierung jedes Tokens ein kleines Entwurfsmodell mehrere Entwürfe erstellt und das Ziel-LLM sie parallel überprüft, um sicherzustellen, dass die endgültige Ausgabe der Zielmodellverteilung entspricht. Die beiden Hauptentwurfsentscheidungen bei MDSD sind die Entwurfsabtastmethode und der Überprüfungsalgorithmus. Für eine feste Entwurfsabtastmethode ist die optimale Akzeptanzrate eine Lösung für ein optimales Transportproblem, aber die Komplexität dieses Problems macht es schwierig, die optimale Akzeptanzrate zu bestimmen und die Kluft zwischen bestehenden Überprüfungsalgorithmen und der theoretischen Obergrenze zu messen. In diesem Papier wird das Duale des optimalen Transportproblems diskutiert, was einen Weg bietet, um die optimale Akzeptanzrate effizient zu berechnen. Zum ersten Mal messen wir die theoretische Obergrenze der Effizienz von MDSD für Vokabellisten in der Größenordnung von Tausenden und quantifizieren die Kluft zwischen bestehenden Überprüfungsalgorithmen und dieser Obergrenze. Wir vergleichen auch verschiedene Entwurfsabtastmethoden anhand ihrer optimalen Akzeptanzraten. Unsere Ergebnisse zeigen, dass die Entwurfsabtastmethode die optimale Akzeptanzrate maßgeblich beeinflusst, wobei die Abtastung ohne Zurücklegen die Abtastung mit Zurücklegen übertrifft. Darüber hinaus erreichen bestehende Überprüfungsalgorithmen sowohl bei der Abtastung ohne Zurücklegen als auch bei der Abtastung mit Zurücklegen nicht die theoretische Obergrenze. Unsere Ergebnisse legen nahe, dass sorgfältig gestaltete Entwurfsabtastmethoden das Potenzial haben, die optimale Akzeptanzrate zu verbessern und die Entwicklung von Überprüfungsalgorithmen zu ermöglichen, die der theoretischen Obergrenze nahekommen.

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.