Hacia la Decodificación Especulativa Óptima con Múltiples Borradores

Towards Optimal Multi-draft Speculative Decoding

February 26, 2025

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

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) se han convertido en una parte indispensable de las tareas de procesamiento del lenguaje natural. Sin embargo, el muestreo autorregresivo se ha convertido en un cuello de botella en términos de eficiencia. La Decodificación Especulativa de Múltiples Borradores (MDSD, por sus siglas en inglés) es un enfoque reciente en el que, al generar cada token, un modelo de borrador pequeño genera múltiples borradores, y el LLM objetivo los verifica en paralelo, asegurando que la salida final se ajuste a la distribución del modelo objetivo. Las dos principales decisiones de diseño en MDSD son el método de muestreo de borradores y el algoritmo de verificación. Para un método de muestreo de borradores fijo, la tasa de aceptación óptima es la solución a un problema de transporte óptimo, pero la complejidad de este problema dificulta la resolución de la tasa de aceptación óptima y la medición de la brecha entre los algoritmos de verificación existentes y el límite teórico superior. Este artículo discute el dual del problema de transporte óptimo, proporcionando una forma de calcular eficientemente la tasa de aceptación óptima. Por primera vez, medimos el límite teórico superior de la eficiencia de MDSD para tamaños de vocabulario en el orden de los miles y cuantificamos la brecha entre los algoritmos de verificación existentes y este límite. También comparamos diferentes métodos de muestreo de borradores basados en sus tasas de aceptación óptimas. Nuestros resultados muestran que el método de muestreo de borradores influye significativamente en la tasa de aceptación óptima, con el muestreo sin reemplazo superando al muestreo con reemplazo. Además, los algoritmos de verificación existentes no alcanzan el límite teórico superior tanto para el muestreo sin reemplazo como con reemplazo. Nuestros hallazgos sugieren que los métodos de muestreo de borradores cuidadosamente diseñados pueden potencialmente mejorar la tasa de aceptación óptima y permitir el desarrollo de algoritmos de verificación que se acerquen al límite teórico superior.

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.