Vocabulary-level Memory Efficiency for Language Model Fine-tuning (2025)

Abstract: The extensive memory footprint of language model (LM) fine-tuning poses a challenge for both researchers and practitioners. LMs use an embedding matrix to represent extensive vocabularies, forming a substantial proportion of the model parameters. While previous work towards memory-efficient fine-tuning has focused on minimizing the number of trainable parameters, reducing the memory footprint of the embedding matrix has yet to be explored. We first demonstrate that a significant proportion of the vocabulary remains unused during fine-tuning. We then propose a simple yet effective approach that leverages this finding to minimize memory usage. We show that our approach provides substantial reductions in memory usage across a wide range of models and tasks. Notably, our approach does not impact downstream task performance, while allowing more efficient use of computational resources.

## introduction language models (lms) @xcite @xcite form the foundation of contemporary natural language processing (nlp), however they require extensive computational resources to train @xcite . this is contrary to the democratization of nlp, exacerbating economic inequalities and hindering inclusivity @xcite . consequently, there is a growing focus towards developing efficient methods for lm training and fine-tuning @xcite . the memory footprint of lms is a major challenge for their application. storing model parameters requires extensive amounts of memory, constraining the size and architecture of the model @xcite . this problem is especially prominent during training as gradients and optimizer states must also be retained @xcite . this can be problematic when using consumer hardware or facing an academic budget @xcite . lms ordinarily use fixed vocabularies to derive vector representations from text, known as word embeddings. each element of the vocabulary has a corresponding word embedding, which collectively form an embedding matrix within the lm. the size of the embedding matrix scales with both the vocabulary size and embedding dimension, comprising a substantial proportion of the model parameters ( in this paper, we first demonstrate that our hypothesis holds for a variety of downstream tasks, with only a small subset of vocabulary used. we then propose a method to reduce memory usage during fine-tuning by excluding unused embeddings. finally, we empirically demonstrate the memory savings from our approach across a range of models and tasks. notably, our approach does not impact downstream task performance and is orthogonal to many existing lm memory efficiency techniques. ## vocabulary usage analysis to empirically assess the level of vocabulary usage during fine-tuning, we first examine the popular glue benchmark @xcite . this comprises a series of tasks that are varied in both size and domain (appendix c). for tokenization, we use the subword vocabulary from gpt-2, which was later adopted by models including roberta @xcite , bart @xcite , gpt-3 @xcite , and opt @xcite . figure 2 illustrates the relationship between unique tokens and total tokens in each of the glue datasets. notably, six out of nine datasets fail to use more than half of the vocabulary. moreover, the smallest dataset, wnli, uses less than 4%. interestingly, we observe that the glue datasets follow a trend resembling heaps' law @xcite . this states that as the size of a corpus grows, there are diminishing gains in new vocabulary. however, our use of a finite subword vocabulary means that the trend is asymptotic to the vocabulary size. separately, the statistical construction of subword vocabularies can reflect anomalies in their training data, creating tokens that may never be used. to examine the extent of the issue, we identify such tokens by evaluating a processed dump of english wikipedia, comprising over 20gb of text. peculiarly, we identify nearly 200 anomalous tokens without a single occurrence (see ## partial embedding matrix adaptation our empirical analysis (section 3) suggests that many fine-tuning datasets only use a fraction of lm vocabulary. we leverage this insight to propose partial embedding matrix adaptation (pema), a method that achieves substantial memory savings by selecting only the minimum subset of word embeddings needed for fine-tuning. notably, this does not impact task performance, as unused word embeddings are not updated during backpropagation. preliminaries. let each token in the vocabulary {w 1 , . . . , w k } be denoted by a unique integer i such that @xmath0. the embedding matrix @xmath1. before fine-tuning. suppose we have fine-tuning dataset @xmath2. we compute the partial vocabulary v ′ ⊂ v consisting of only the tokens in d. as the elements of v ′ are not necessarily consecutive integers, we define an arbitrary mapping f : we then construct the partial embedding matrix e ′ ∈ r |v ′ |×d with entries e ′ [:, f (i)] = e[:, i] for all @xmath3. that is, e ′ retains only embedding vectors corresponding to tokens in v ′ . to adapt d for the partial vocabulary v ′ , we create an intermediary dataset d ′ where each entry d ′ [i, j] = f (d[i, j]). finally, we use d ′ and e ′ in place of d and e. after fine-tuning. following fine-tuning, our partial embedding matrix e ′ holds the newly learned embeddings for the partial vocabulary. however, we do not wish to keep only the partial vocabulary, as this would limit future use of the model (i.e. tasks with different vocabulary). therefore, we merge the newly learned embeddings into the original embedding matrix (stored on-disk). more formally, we update e such that e[:, f -1 (i)] = e ′ [:, i] for all @xmath4. this ensures that the model remains structurally identical, with embeddings for the complete vocabulary. models. similarly, we select a variety of popular models used in existing work. however, we place an emphasis on having a variety of vocabularies ( ## results larger vocabularies see more memory savings. ## conclusion in this paper, we identified that many fine-tuning datasets do not use the majority of lm vocabulary. we then proposed partial embedding matrix adaptation (pema), a simple yet effective approach to minimize lm memory use during fine-tuning, that is orthogonal to many existing methods. finally, we empirically demonstrated that our approach offers substantial memory savings across a variety of popular tasks and models, without compromising performance. as future work, we are interested in adapting our approach for the output embedding matrix to offer further memory savings.