VOCABTRIM: Обрезка словаря для эффективного спекулятивного декодирования в больших языковых моделях
VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs
June 28, 2025
Авторы: Raghavv Goel, Sudhanshu Agrawal, Mukul Gagrani, Junyoung Park, Yifan Zao, He Zhang, Tian Liu, Yiping Yang, Xin Yuan, Jiuyan Lu, Chris Lott, Mingu Lee
cs.AI
Аннотация
В данной статье мы представляем простую методику, не требующую обучения, для повышения производительности методов спекулятивного декодирования (SpD), основанных на использовании драфтеров, которая включает голову языковой модели (LM head) в процесс генерации черновиков. Спекулятивное декодирование на основе драфтеров использует одну или несколько меньших языковых моделей, называемых драфтерами или моделями-драфтерами, для выборки черновой последовательности или дерева, состоящего из нескольких токенов, с последующей проверкой базовой большой языковой моделью (LLM), целевой моделью, которая принимает подмножество токенов как валидную генерацию. Поскольку обычно считается, что спекулятивное декодирование требует однозначного соответствия между словарями целевой модели и модели-драфтера, естественным решением было использование общего словаря или даже общей LM head, как в методах EAGLE или Medusa. Мы впервые выявляем, что такая схема выборки токенов в черновиках изначально содержит избыточные вычислительные затраты, особенно для некоторых целевых LLM с очень большими словарями. Затем мы предлагаем простую методику, VocabTrim, для снижения этих затрат и повышения скорости генерации в условиях ограниченной памяти. VocabTrim перестраивает LM head модели-драфтера, ограничивая её набором токенов, которые чаще всего выбираются из словаря целевой модели. Хотя ограничение словаря в процессе генерации черновиков незначительно снижает процент принятия токенов, оно существенно уменьшает задержку генерации в условиях ограниченной памяти, что часто встречается на периферийных устройствах, что приводит к увеличению ускорения в условиях ограниченной памяти (MBSU). Мы показываем, что наш метод может повысить MBSU для моделей Llama-3 на Spec-Bench, в частности, на 16% для модели Llama-3.2-3B-Instruct.
English
In this paper, we introduce a simple training-free technique to improve the
performance of drafter-based speculative decoding (SpD) methods that
incorporates language modeling head (LM head) during drafting process. A
drafter-based speculative decoding leverages one or more smaller language
models, a.k.a. drafters or draft models, to sample a draft sequence or tree
consisting of multiple tokens, followed by verification by a base LLM, a target
model, accepting a subset as its valid generation. As it is usually considered
that the speculative decoding requires one-to-one mapping between vocabularies
of the target model and the draft model, it has been natural to share the
vocabulary between them, or even share the LM head as in EAGLE or Medusa. We
first identify that this draft token sampling scheme inherently contains an
unnecessary inference overhead in drafting, especially for some target LLMs
with very large vocabularies. Then, we propose a simple technique, VocabTrim,
to mitigate the drafting overhead to improve the generation speed in
memory-bound environment. VocabTrim reconstructs the drafter LM head to contain
only a limited set of tokens, selected by the most frequently sampled from the
vocabulary of the target model. While limiting the vocabulary in drafting
slightly degrades the acceptance rate, it significantly reduces the drafting
latency in memory-bound process which is often the case on edge devices,
resulting in higher memory-bound speed up (MBSU). We show that our method can
boost the memory-bound speed-up for Llama-3 models on Spec-Bench, specifically
by 16% for Llama-3.2-3B-Instruct.