VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs
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.
68637193588cea0da970c914
Raghavv Goel
68637193588cea0da970c915
Sudhanshu Agrawal
68637193588cea0da970c916
Mukul Gagrani
68637193588cea0da970c917
Junyoung Park
68637193588cea0da970c918
Yifan Zao
68637193588cea0da970c919
He Zhang
68637193588cea0da970c91a
Tian Liu
68637193588cea0da970c91b
Yiping Yang
68637193588cea0da970c91c
Xin Yuan
68637193588cea0da970c91d
Jiuyan Lu
68637193588cea0da970c91e
Chris Lott
68637193588cea0da970c91f
Mingu Lee
VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs
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.
A technique called VocabTrim improves the generation speed of drafter-based speculative decoding by reducing the vocabulary size of the drafter LM head, leading to higher memory-bound speed-up.
drafter-based speculative decoding language modeling head drafters draft models base LLM target model speculative decoding vocabulary EAGLE Medusa VocabTrim drafting overhead drafting latency memory-bound environment memory-bound speed-up MBSU Llama-3 Spec-Bench