Vocabulary Expansion of Large Language Models via Kullback-Leibler-Based Self-Distillation

TL;DR

This paper introduces a mathematically grounded KL divergence-based self-distillation method for expanding the vocabulary of large language models, enabling better incorporation of new domain-specific terms without retraining from scratch.

Contribution

It presents a novel knowledge distillation technique that handles different tokenizations, improving vocabulary expansion in frozen large language models.

Findings

Outperforms conventional cross-entropy training in vocabulary expansion tasks

Achieves top performance on 2000 code-generation benchmarks

Provides interpretability insights into token representation learning

Abstract

Large pre-trained language models often struggle to incorporate new domain-specific terminology when fine-tuned on small, specialized corpora. In this work, we address the challenge of vocabulary expansion in frozen LLMs by introducing a mathematically grounded method for knowledge distillation via KL divergence, even when the original and extended models use different tokenizations. This allows the student model to inherit distributional knowledge from the teacher despite differing vocabularies. We compare our KL-based distillation approach to conventional cross-entropy training, evaluating both methods across multiple strategies for initializing new token embeddings. After embedding initialization, models are further fine-tuned to integrate the new vocabulary. Each trained model is benchmarked on approximately 2000 code-generation tasks, where our approach achieves the best performance across the board. Finally, through mechanistic interpretability, we analyze how models learn representations for the new tokens, providing an explanation for the observed gains and offering insight into the structure of embedding space during vocabulary expansion.