VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs

2025-06-27

AI Summary

To address the high memory overhead and low inference efficiency in speculative decoding (SpD) caused by large target-model vocabularies, this paper proposes VocabTrim—a training-free, lightweight vocabulary pruning method. Its core innovation lies in the first systematic identification of redundancy in full-vocabulary sampling during SpD; accordingly, VocabTrim retains only the most frequent tokens from the target model’s vocabulary to construct a compact sub-vocabulary, and enables efficient adaptation via vocabulary remapping and lightweight drafter head tuning. Crucially, it requires no modifications to the target model or additional training. VocabTrim significantly alleviates memory bottlenecks while preserving accuracy. On Spec-Bench, it achieves a 16% improvement in memory-bound speedup for Llama-3.2-3B-Instruct. The method particularly enhances generation efficiency in resource-constrained settings, such as edge devices, without compromising decoding quality or introducing computational overhead from retraining.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Summarization
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

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.

Problem

• Reduces drafting overhead in speculative decoding for LLMs
• Improves generation speed in memory-bound environments
• Optimizes drafter vocabulary for faster token sampling

Innovation

• Vocabulary pruning for efficient speculative decoding
• Training-free technique to reduce drafting overhead
• Reconstructs drafter LM head with frequent tokens