arXiv:2504.11426cs.CLcs.AI2025-04被引 2

解决大模型压缩中跨词汇量知识蒸馏难题,实现不同模型间高效知识迁移。

A Dual-Space Framework for General Knowledge Distillation of Large Language Models

  • 设计双空间框架,统一师生模型输出表示空间,消除分布差异
  • 提出精确词元对齐算法,支持不同分词器的模型间蒸馏
  • 通用性强,适用于任意两模型,涵盖指令遵循、数学推理等任务

知识蒸馏(KD)是压缩大型语言模型(LLMs)的有效方法,通过将教师模型的知识传递给学生模型。现有白盒蒸馏框架存在两大局限:一是不同输出空间的概率分布难以对齐,限制了师生模型的相似性;二是无法应用于不同词汇表的模型。根本原因在于师生模型使用不同的预测头,导致输出空间和维度不一致。为此,本文提出双空间知识蒸馏(DSKD)框架,通过引入两个理想初始化的投影器,将教师/学生隐藏状态分别映射到学生/教师表示空间,使两者共享同一预测头,统一输出空间。此外,提出精确词元对齐(ETA)算法,对齐不同分词序列中的相同词元。基于此,DSKD可支持离策略与在线策略蒸馏,并适用于任意两模型间的蒸馏,无论其词汇表是否相同。在指令遵循、数学推理和代码生成等多个基准上的实验表明,DSKD显著优于现有白盒蒸馏方法,并超越其他跨分词器蒸馏方法。

原文摘要 · Abstract (English)

Knowledge distillation (KD) is a promising solution to compress large language models (LLMs) by transferring their knowledge to smaller models. During this process, white-box KD methods usually minimize the distance between the output distributions of the teacher model and the student model to transfer more information. However, we reveal that the current white-box KD framework exhibits two limitations: a) bridging probability distributions from different output spaces will limit the similarity between the teacher model and the student model; b) this framework cannot be applied to LLMs with different vocabularies. One of the root causes for these limitations is that the distributions from the teacher and the student for KD are output by different prediction heads, which yield distributions in different output spaces and dimensions. Therefore, in this paper, we propose a dual-space knowledge distillation (DSKD) framework that unifies the prediction heads of the teacher and the student models for KD. Specifically, we first introduce two projectors with ideal initialization to project the teacher/student hidden states into the student/teacher representation spaces. After this, the hidden states from different models can share the same head and unify the output spaces of the distributions. Furthermore, we develop an exact token alignment (ETA) algorithm to align the same tokens in two differently-tokenized sequences. Based on the above, our DSKD framework is a general KD framework that supports both off-policy and on-policy KD, and KD between any two LLMs regardless of their vocabularies. Extensive experiments on instruction-following, mathematical reasoning, and code generation benchmarks show that DSKD significantly outperforms existing methods based on the current white-box KD framework and surpasses other cross-tokenizer KD methods for LLMs with different vocabularies.

知识蒸馏大模型压缩跨词汇蒸馏双空间框架

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。