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

📄 arXiv: 2504.11426v1 📥 PDF

作者: Xue Zhang, Songming Zhang, Yunlong Liang, Fandong Meng, Yufeng Chen, Jinan Xu, Jie Zhou

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-04-15

备注: 19 pages, 9 figures, 11 tables, under review. Code is available at: https://github.com/songmzhang/DSKDv2. arXiv admin note: text overlap with arXiv:2406.17328

💡 一句话要点

提出双空间知识蒸馏框架DSKD，解决大语言模型通用知识蒸馏问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 大语言模型 模型压缩 双空间投影 Token对齐

📋 核心要点

  1. 现有白盒知识蒸馏方法在不同输出空间桥接概率分布，限制了模型相似性，且无法应用于不同词汇表的LLM。
  2. 提出双空间知识蒸馏框架DSKD，通过投影器统一教师和学生模型的预测头，实现输出空间的对齐。
  3. 实验表明，DSKD在指令跟随、数学推理和代码生成任务上显著优于现有方法，尤其是在跨词汇表蒸馏方面。

📝 摘要（中文）

知识蒸馏(KD)是一种通过将大型语言模型(LLM)的知识转移到较小模型来压缩LLM的有前景的解决方案。在这个过程中，白盒KD方法通常最小化教师模型和学生模型的输出分布之间的距离，以传递更多信息。然而，我们发现当前的白盒KD框架存在两个局限性：a)桥接来自不同输出空间的概率分布会限制教师模型和学生模型之间的相似性；b)该框架不能应用于具有不同词汇表的LLM。这些限制的一个根本原因是，用于KD的教师和学生模型的分布是由不同的预测头输出的，这导致分布在不同的输出空间和维度中。因此，在本文中，我们提出了一种双空间知识蒸馏(DSKD)框架，该框架统一了教师和学生模型的预测头以进行KD。具体来说，我们首先引入两个具有理想初始化的投影器，将教师/学生隐藏状态投影到学生/教师表示空间中。之后，来自不同模型的隐藏状态可以共享同一个头，并统一分布的输出空间。此外，我们开发了一种精确的token对齐(ETA)算法，以对齐两个不同token化序列中的相同token。基于以上，我们的DSKD框架是一个通用的KD框架，支持离策略和在策略KD，以及任意两个LLM之间的KD，而不管它们的词汇表如何。在指令跟随、数学推理和代码生成基准上的大量实验表明，DSKD显著优于基于当前白盒KD框架的现有方法，并且超过了其他用于具有不同词汇表的LLM的跨token器KD方法。

🔬 方法详解

问题定义：现有白盒知识蒸馏方法依赖于最小化教师模型和学生模型输出分布之间的距离。然而，当教师模型和学生模型具有不同的词汇表或输出空间时，这种方法会遇到困难。具体来说，不同模型的预测头产生不同维度和空间的概率分布，直接比较这些分布会限制知识传递的有效性。此外，现有方法难以处理具有不同tokenizer的LLM之间的知识蒸馏。

核心思路：DSKD的核心思路是通过将教师模型和学生模型的隐藏状态投影到统一的表示空间，从而消除输出空间差异带来的问题。具体来说，DSKD引入了两个投影器，分别将教师模型的隐藏状态投影到学生模型的表示空间，以及将学生模型的隐藏状态投影到教师模型的表示空间。这样，两个模型就可以共享相同的预测头，从而在相同的输出空间进行知识蒸馏。

技术框架：DSKD框架包含以下几个主要步骤：1) 初始化两个投影器，分别用于将教师模型和学生模型的隐藏状态投影到对方的表示空间。2) 使用投影器将教师模型和学生模型的隐藏状态进行投影。3) 使用统一的预测头，基于投影后的隐藏状态计算输出分布。4) 使用知识蒸馏损失函数，最小化教师模型和学生模型在统一输出空间中的分布差异。5) 使用精确token对齐(ETA)算法，对齐不同tokenizer产生的token序列。

关键创新：DSKD的关键创新在于双空间投影的思想，它通过将教师模型和学生模型的隐藏状态投影到统一的表示空间，解决了不同输出空间带来的问题。此外，ETA算法解决了不同tokenizer带来的token对齐问题，使得DSKD可以应用于具有不同词汇表的LLM之间的知识蒸馏。

关键设计：DSKD的关键设计包括：1) 投影器的初始化方式，论文中使用了理想初始化方法，以保证投影后的隐藏状态能够尽可能地保留原始信息。2) 知识蒸馏损失函数的选择，可以使用KL散度、交叉熵等常用的损失函数。3) ETA算法的具体实现，论文中提出了一种基于动态规划的ETA算法，可以有效地对齐不同tokenizer产生的token序列。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DSKD在指令跟随、数学推理和代码生成任务上显著优于现有方法。例如，在指令跟随任务上，DSKD相比于基线方法提升了X%。此外，DSKD在跨词汇表蒸馏方面也取得了显著的成果，超过了其他跨tokenizer KD方法。

🎯 应用场景

DSKD框架可广泛应用于大语言模型的压缩和加速，尤其适用于需要将大型模型知识迁移到资源受限设备上的场景。例如，可以将大型通用语言模型的知识蒸馏到移动设备上的小型模型，从而实现本地化的智能应用。此外，DSKD还可以用于构建特定领域的专业模型，通过将通用模型的知识迁移到领域模型，提高领域模型的性能。

📄 摘要（原文）

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.