Кенгуру: Без потерь самоспекулятивное декодирование с двойным ранним выходом.

Kangaroo: Lossless Self-Speculative Decoding via Double Early Exiting

April 29, 2024

Авторы: Fangcheng Liu, Yehui Tang, Zhenhua Liu, Yunsheng Ni, Kai Han, Yunhe Wang

cs.AI

Аннотация

Спекулятивное декодирование продемонстрировало свою эффективность в ускорении вывода больших языковых моделей, сохраняя при этом согласованное распределение выборки. Однако традиционный подход обучения отдельной черновой модели для достижения удовлетворительной скорости принятия токенов может быть затратным. Извлекая вдохновение из раннего завершения, мы предлагаем новую собственную спекулятивную декодирующую структуру Kangaroo, которая использует фиксированную неглубокую подсеть в качестве собственной черновой модели, а оставшиеся слои служат в качестве более крупной целевой модели. Мы обучаем легкий и эффективный модуль адаптера поверх подсети, чтобы преодолеть разрыв между способностью представления подсети и полной модели. Следует отметить, что задержка вывода собственной черновой модели может больше не быть незначительной по сравнению с крупной моделью, что требует стратегий для увеличения скорости принятия токенов, минимизируя при этом шаги черновой модели. Для решения этой проблемы мы вводим дополнительный механизм раннего завершения для генерации черновых токенов. Конкретно, мы останавливаем последующее предсказание маленькой модели во время фазы чернового процесса, как только уровень уверенности для текущего токена опускается ниже определенного порога. Обширные эксперименты на Spec-Bench демонстрируют эффективность Kangaroo. При верификации одиночной последовательности Kangaroo достигает ускорения до 1,68 раза на Spec-Bench, превосходя Medusa-1 с 88,7\% меньшим количеством дополнительных параметров (67M по сравнению с 591M). Код для Kangaroo доступен по ссылке https://github.com/Equationliu/Kangaroo.

English

Speculative decoding has demonstrated its effectiveness in accelerating the inference of large language models while maintaining a consistent sampling distribution. However, the conventional approach of training a separate draft model to achieve a satisfactory token acceptance rate can be costly. Drawing inspiration from early exiting, we propose a novel self-speculative decoding framework Kangaroo, which uses a fixed shallow sub-network as a self-draft model, with the remaining layers serving as the larger target model. We train a lightweight and efficient adapter module on top of the sub-network to bridge the gap between the sub-network and the full model's representation ability. It is noteworthy that the inference latency of the self-draft model may no longer be negligible compared to the large model, necessitating strategies to increase the token acceptance rate while minimizing the drafting steps of the small model. To address this challenge, we introduce an additional early exiting mechanism for generating draft tokens. Specifically, we halt the small model's subsequent prediction during the drafting phase once the confidence level for the current token falls below a certain threshold. Extensive experiments on the Spec-Bench demonstrate the effectiveness of Kangaroo. Under single-sequence verification, Kangaroo achieves speedups up to 1.68times on Spec-Bench, outperforming Medusa-1 with 88.7\% fewer additional parameters (67M compared to 591M). The code for Kangaroo is available at https://github.com/Equationliu/Kangaroo.