KL-based self-distillation for large language models
This addresses the challenge of incorporating domain-specific terminology into pre-trained LLMs, which is important for practitioners needing to adapt models to specialized domains, though it appears incremental as it builds on existing distillation techniques.
The paper tackles the problem of vocabulary expansion in frozen large language models when fine-tuned on small specialized corpora, introducing a KL-based self-distillation method that allows models with different tokenizations to share distributional knowledge, achieving the best performance on approximately 2000 code-generation tasks.
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.