Flexible Feature Distillation for Large Language Models
作者: Khouloud Saadi, Di Wang
分类: cs.CL
发布日期: 2025-07-14 (更新: 2025-10-02)
💡 一句话要点
Flex-KD：一种面向大语言模型的灵活特征蒸馏方法，无需参数对齐，提升下游任务性能。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 知识蒸馏 大型语言模型 特征蒸馏 模型压缩 梯度选择 无参数方法 生成任务
📋 核心要点
- 现有LLM知识蒸馏方法侧重于logit，忽略了模型内部蕴含的丰富特征信息，特征蒸馏受限于教师-学生模型隐藏层大小需一致的假设。
- Flex-KD通过梯度选择教师模型隐藏层中最相关的维度子空间进行蒸馏，避免了投影引入的额外参数和信息扭曲。
- 实验表明，Flex-KD在分类和生成任务中均能有效提升学生模型性能，相比线性投影基线最高提升3.75%。
📝 摘要（中文）
知识蒸馏(KD)已成为压缩大型语言模型(LLM)的基石。然而，现有的LLM-KD方法主要集中在基于logits的方法上，这些方法虽然性能良好，但忽略了LLM丰富的内部表示。特征级别的KD可以利用这种结构来提供互补的优势，但由于当前特征KD方法通常假设教师和学生的隐藏层大小相同，这是一个限制性且不切实际的假设，因此仍未得到充分探索。一种常见的解决方法是训练一个线性投影器来对齐它们的特征空间；然而，这引入了额外的参数，扭曲了教师的嵌入，并且常常降低下游性能，尤其是在生成任务中。我们提出了Flex-KD，一个用于LLM的任务驱动特征蒸馏的无参数框架。Flex-KD不投影整个教师表示，而是使用基于梯度的分数来识别教师隐藏状态中最相关的维度，并仅将该子空间提炼到学生模型中。这确保了学生模型的有限容量被分配给信息丰富的组件，同时避免了投影器引起的失真和额外的参数。Flex-KD与现有的KD流程无缝集成，并支持不同的教师-学生隐藏层大小。在分类和生成任务（即指令跟随和摘要）中的大量实验表明，Flex-KD始终提高学生模型的性能，与线性投影基线相比，性能提升高达3.75%。
🔬 方法详解
问题定义：现有的大语言模型知识蒸馏方法，特别是特征蒸馏方法，通常假设教师模型和学生模型的隐藏层维度相同。当教师模型和学生模型的隐藏层维度不同时，常用的方法是训练一个线性投影层来对齐特征空间。然而，这种方法会引入额外的参数，并且可能会扭曲教师模型的特征表示，从而降低下游任务的性能，尤其是在生成任务中。
核心思路：Flex-KD的核心思路是只蒸馏教师模型中与特定任务最相关的特征维度，而不是蒸馏整个特征空间。通过这种方式，可以避免将学生模型的容量浪费在不重要的特征上，并且可以减少由于特征对齐而引入的失真。该方法无需训练额外的参数，因此更加高效。
技术框架：Flex-KD的整体框架包括以下几个步骤：1. 使用教师模型和学生模型处理输入数据，得到各自的隐藏层表示。2. 计算教师模型隐藏层表示中每个维度的梯度，用于衡量该维度与任务的相关性。3. 根据梯度选择最相关的维度子空间。4. 将教师模型中选定的维度子空间蒸馏到学生模型中。Flex-KD可以与现有的知识蒸馏流程无缝集成。
关键创新：Flex-KD的关键创新在于提出了一种无参数的特征选择方法，该方法基于梯度来识别教师模型中与任务最相关的特征维度。与传统的特征对齐方法相比，Flex-KD避免了引入额外的参数和失真，并且可以处理教师模型和学生模型隐藏层维度不同的情况。
关键设计：Flex-KD的关键设计包括：1. 使用梯度作为特征重要性的度量。梯度的绝对值越大，表示该维度对任务的影响越大。2. 选择top-k个梯度最大的维度作为最相关的维度子空间。k是一个超参数，可以根据具体任务进行调整。3. 使用标准的知识蒸馏损失函数，例如KL散度或MSE损失，来训练学生模型。
🖼️ 关键图片
📊 实验亮点
实验结果表明，Flex-KD在分类和生成任务中均能有效提升学生模型的性能。在指令跟随任务中，Flex-KD相比线性投影基线取得了高达3.75%的性能提升。在摘要任务中，Flex-KD也显著优于其他基线方法。这些结果表明，Flex-KD是一种有效的特征蒸馏方法，可以用于压缩和加速大型语言模型。
🎯 应用场景
Flex-KD可应用于各种需要压缩大型语言模型的场景，例如移动设备上的自然语言处理、边缘计算等。该方法尤其适用于生成任务，如文本摘要、机器翻译等，可以有效提升学生模型的性能，同时降低计算成本和模型大小。未来，Flex-KD可以进一步扩展到多模态学习等领域。
📄 摘要（原文）
Knowledge distillation (KD) has become a cornerstone for compressing large language models (LLMs). However, existing LLM-KD methods have primarily focused on logit-based approaches, which achieve good performance but overlook the rich internal representations of LLMs. Feature-level KD could leverage this structure to provide complementary benefits, yet it remains underexplored because current feature-KD approaches typically assume identical teacher-student hidden sizes, a restrictive and unrealistic assumption. A common workaround is to train a linear projector to align their feature spaces; however, this introduces additional parameters, distorts teacher embeddings, and often degrades downstream performance, especially in generative tasks. We propose Flex-KD, a parameter-free framework for task-driven feature distillation for LLMs. Instead of projecting the entire teacher representation, Flex-KD uses gradient-based scores to identify the most task-relevant dimensions of the teacher's hidden states and distills only this subspace into the student. This ensures that the student's limited capacity is allocated to informative components, while avoiding projector-induced distortion and extra parameters. Flex-KD integrates seamlessly with existing KD pipelines and supports differing teacher-student hidden sizes. Extensive experiments across both classification and generative tasks, i.e., instruction-following and summarization, show that Flex-KD consistently boosts student performance, achieving up to a 3.75 percent performance gain over the linear projection baseline.