Vocabulary Expansion of Large Language Models via Kullback-Leibler-Based Self-Distillation
作者: Max Rehman Linder
分类: cs.CL, cs.AI
发布日期: 2025-08-14 (更新: 2026-01-12)
备注: Master's Thesis
💡 一句话要点
提出基于KL散度的自蒸馏方法，用于扩展冻结LLM的词汇表以适应特定领域术语。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 词汇表扩展 知识蒸馏 KL散度 大型语言模型 领域自适应
📋 核心要点
- 现有方法难以将领域特定术语融入大型语言模型，尤其是在资源有限的情况下，导致模型在特定任务上表现不佳。
- 论文提出基于KL散度的自蒸馏方法，使扩展后的模型能够从原始模型中学习知识，即使两者使用不同的分词方式。
- 实验表明，该方法在代码生成任务中优于传统交叉熵训练，并通过可解释性分析揭示了模型学习新token表示的机制。
📝 摘要（中文）
大型预训练语言模型在针对小型专业语料库进行微调时，通常难以整合新的领域特定术语。本文提出了一种基于KL散度的知识蒸馏方法，用于解决冻结LLM中的词汇表扩展问题，即使原始模型和扩展模型使用不同的分词方式也能有效工作。这种方法允许学生模型从教师模型继承分布知识，即使词汇表不同。我们将基于KL的蒸馏方法与传统的交叉熵训练进行比较，评估了两种方法在初始化新token嵌入的多种策略下的表现。在嵌入初始化之后，模型会进一步微调以整合新的词汇表。每个训练好的模型都在大约2000个代码生成任务上进行了基准测试，我们的方法在所有方面都取得了最佳性能。最后，通过机制可解释性，我们分析了模型如何学习新token的表示，为观察到的收益提供了解释，并深入了解了词汇表扩展期间嵌入空间的结构。
🔬 方法详解
问题定义：论文旨在解决大型预训练语言模型在特定领域微调时，难以有效整合新领域术语的问题。现有方法在处理不同分词方式以及知识迁移方面存在不足，导致模型无法充分利用预训练知识，影响在特定任务上的性能。
核心思路：核心思路是利用KL散度进行知识蒸馏，将原始模型（教师模型）的知识迁移到扩展了词汇表的模型（学生模型）。即使教师模型和学生模型使用不同的分词方式，KL散度也能有效地衡量两个概率分布之间的差异，从而实现知识迁移。
技术框架：整体框架包括以下几个阶段：1) 词汇表扩展：将新的领域特定术语添加到模型的词汇表中。2) 嵌入初始化：使用不同的策略初始化新token的嵌入向量。3) KL散度蒸馏：使用KL散度作为损失函数，训练学生模型，使其输出分布尽可能接近教师模型的输出分布。4) 微调：在特定领域的语料库上对学生模型进行微调，以进一步整合新的词汇表。5) 评估：在代码生成任务上评估模型的性能。
关键创新：最重要的创新点在于使用KL散度进行跨分词方式的知识蒸馏。传统方法通常依赖于交叉熵损失，但在不同分词方式下，交叉熵损失难以有效衡量两个模型输出分布的相似性。KL散度能够更准确地捕捉分布之间的差异，从而实现更有效的知识迁移。
关键设计：关键设计包括：1) KL散度损失函数：使用KL散度作为损失函数，衡量教师模型和学生模型输出分布之间的差异。2) 嵌入初始化策略：探索了不同的嵌入初始化策略，例如随机初始化、使用现有词汇的平均嵌入等。3) 微调策略：在特定领域的语料库上对学生模型进行微调，以进一步整合新的词汇表。
🖼️ 关键图片
📊 实验亮点
实验结果表明，基于KL散度的自蒸馏方法在代码生成任务中取得了最佳性能，显著优于传统的交叉熵训练方法。具体而言，该方法在所有测试场景下均取得了领先，证明了其在词汇表扩展方面的有效性。通过可解释性分析，论文还揭示了模型学习新token表示的机制。
🎯 应用场景
该研究成果可应用于各种需要扩展语言模型词汇表的场景，例如医学、法律、金融等专业领域。通过有效整合领域特定术语，可以提升语言模型在这些领域的应用性能，例如专业文档生成、智能问答、信息抽取等，具有重要的实际价值和广泛的应用前景。
📄 摘要（原文）
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.