Rumo à Decodificação Especulativa Multi-rascunho Ótima

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

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) tornaram-se uma parte indispensável das tarefas de processamento de linguagem natural. No entanto, a amostragem autorregressiva tornou-se um gargalo de eficiência. A Decodificação Especulativa Multi-Rascunho (MDSD) é uma abordagem recente em que, ao gerar cada token, um pequeno modelo de rascunho gera vários rascunhos, e o LLM alvo os verifica em paralelo, garantindo que a saída final esteja de acordo com a distribuição do modelo alvo. As duas principais escolhas de design na MDSD são o método de amostragem de rascunho e o algoritmo de verificação. Para um método de amostragem de rascunho fixo, a taxa de aceitação ótima é uma solução para um problema de transporte ótimo, mas a complexidade desse problema torna difícil resolver a taxa de aceitação ótima e medir a diferença entre os algoritmos de verificação existentes e o limite superior teórico. Este artigo discute o dual do problema de transporte ótimo, fornecendo uma maneira de calcular eficientemente a taxa de aceitação ótima. Pela primeira vez, medimos o limite superior teórico da eficiência da MDSD para tamanhos de vocabulário na casa dos milhares e quantificamos a diferença entre os algoritmos de verificação existentes e esse limite. Também comparamos diferentes métodos de amostragem de rascunho com base em suas taxas de aceitação ótimas. Nossos resultados mostram que o método de amostragem de rascunho influencia fortemente a taxa de aceitação ótima, com a amostragem sem reposição superando a amostragem com reposição. Além disso, os algoritmos de verificação existentes não alcançam o limite superior teórico tanto para a amostragem sem reposição quanto para a amostragem com reposição. Nossas descobertas sugerem que métodos de amostragem de rascunho cuidadosamente projetados podem potencialmente melhorar a taxa de aceitação ótima e permitir o desenvolvimento de algoritmos de verificação que se aproximem do limite superior teórico.

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.