VocabTailor: Dynamic Vocabulary Selection for Downstream Tasks in Small Language Models

TL;DR

VocabTailor is a dynamic vocabulary selection framework for small language models that significantly reduces memory usage while maintaining task performance, enabling more efficient deployment on resource-constrained devices.

Contribution

It introduces a decoupled, hybrid static-dynamic vocabulary selection method that adapts vocabulary components on demand, outperforming static pruning techniques.

Findings

Achieves up to 99% reduction in vocabulary-related memory usage.

Maintains minimal or no performance degradation across diverse tasks.

Outperforms existing static vocabulary pruning methods.

Abstract

Small Language Models (SLMs) provide computational advantages in resource-constrained environments, yet memory limitations remain a critical bottleneck for edge device deployment. A substantial portion of SLMs' memory footprint stems from vocabulary-related components, particularly embeddings and language modeling (LM) heads, due to large vocabulary sizes. Existing static vocabulary pruning, while reducing memory usage, suffers from rigid, one-size-fits-all designs that cause information loss during the prefill stage and lack flexibility. In this work, we identify two key principles underlying the vocabulary reduction challenge: the lexical locality principle, the observation that only a small subset of tokens is required during any single inference, and the asymmetry in computational characteristics between vocabulary-related components of SLM. Based on these insights, we introduce VocabTailor, a novel decoupled dynamic vocabulary selection framework that addresses memory constraints through offloading embedding and implements a hybrid static-dynamic vocabulary selection strategy for LM Head, enabling on-demand loading of vocabulary components. Comprehensive experiments across diverse downstream tasks demonstrate that VocabTailor achieves a reduction of up to 99% in the memory usage of vocabulary-related components with minimal or no degradation in task performance, substantially outperforming existing static vocabulary pruning. Our code is available at https://github.com/AwakenedInsects/VocabTailor.