Vers un décodage spéculatif multi-brouillon optimal

Towards Optimal Multi-draft Speculative Decoding

February 26, 2025

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

cs.AI

Résumé

Les grands modèles de langage (LLMs) sont devenus une partie indispensable des tâches de traitement du langage naturel. Cependant, l'échantillonnage autorégressif est devenu un goulot d'étranglement en termes d'efficacité. Le Décodage Spéculatif Multi-Brouillon (MDSD) est une approche récente où, lors de la génération de chaque jeton, un petit modèle de brouillon génère plusieurs brouillons, et le LLM cible les vérifie en parallèle, garantissant que la sortie finale est conforme à la distribution du modèle cible. Les deux principaux choix de conception dans le MDSD sont la méthode d'échantillonnage de brouillon et l'algorithme de vérification. Pour une méthode d'échantillonnage de brouillon fixe, le taux d'acceptation optimal est une solution à un problème de transport optimal, mais la complexité de ce problème rend difficile la résolution du taux d'acceptation optimal et la mesure de l'écart entre les algorithmes de vérification existants et la limite supérieure théorique. Cet article discute du dual du problème de transport optimal, fournissant un moyen de calculer efficacement le taux d'acceptation optimal. Pour la première fois, nous mesurons la limite supérieure théorique de l'efficacité du MDSD pour des tailles de vocabulaire dans les milliers et quantifions l'écart entre les algorithmes de vérification existants et cette limite. Nous comparons également différentes méthodes d'échantillonnage de brouillon en fonction de leurs taux d'acceptation optimaux. Nos résultats montrent que la méthode d'échantillonnage de brouillon influence fortement le taux d'acceptation optimal, l'échantillonnage sans remplacement surpassant l'échantillonnage avec remplacement. De plus, les algorithmes de vérification existants n'atteignent pas la limite supérieure théorique pour les échantillonnages sans remplacement et avec remplacement. Nos conclusions suggèrent que des méthodes d'échantillonnage de brouillon soigneusement conçues peuvent potentiellement améliorer le taux d'acceptation optimal et permettre le développement d'algorithmes de vérification qui se rapprochent de la limite supérieure théorique.

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.