К Оптимальному Многоэтапному Спекулятивному Декодированию

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.