LLM-Oriented Token-Adaptive Knowledge Distillation

📄 arXiv: 2510.11615v1 📥 PDF

作者: Xurong Xie, Zhucun Xue, Jiafu Wu, Jian Li, Yabiao Wang, Xiaobin Hu, Yong Liu, Jiangning Zhang

分类: cs.CL, cs.AI

发布日期: 2025-10-13

备注: 15 pages, 4 figures


💡 一句话要点

提出面向LLM的Token自适应知识蒸馏框架AdaKD,提升学生模型性能。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 知识蒸馏 大规模语言模型 模型压缩 Token自适应 动态学习 温度缩放

📋 核心要点

  1. 现有知识蒸馏方法采用静态策略,忽略了学生模型的动态学习过程,导致知识转移效率低下。
  2. AdaKD框架通过token难度指标,动态调整蒸馏焦点和温度,使蒸馏过程适应每个token的学习状态。
  3. AdaKD作为即插即用框架,在多种模型和基准测试上,持续提升了各种蒸馏方法的性能。

📝 摘要(中文)

知识蒸馏(KD)是压缩大规模语言模型(LLM)的关键技术。然而,目前主流的基于logits的方法通常采用静态策略,与学生模型的动态学习过程不一致。这些方法通常不加区分地对待所有token,并应用单一的固定温度,导致次优的知识转移。为了解决这些局限性,我们提出了一种面向LLM的Token自适应知识蒸馏(AdaKD)框架,该框架使蒸馏过程适应每个token的实时学习状态。AdaKD由两个协同模块组成,这两个模块由统一的token难度指标驱动。首先,我们的损失驱动自适应Token聚焦(LATF)模块通过监控学生的学习稳定性来动态调整蒸馏焦点,在每个训练阶段将计算资源集中在最有价值的token上。其次,我们引入了反向难度温度缩放(IDTS),这是一种违反直觉但有效的token级温度策略。它对困难token采用低温以进行有针对性的错误纠正,而对简单token采用高温,以鼓励学生从教师模型的完整和平滑的输出分布中学习,从而增强泛化能力。作为一个即插即用的框架,AdaKD可以持续提高各种蒸馏方法在多种模型架构和基准测试上的性能。

🔬 方法详解

问题定义:现有基于logits的知识蒸馏方法通常采用静态策略,即对所有token使用相同的温度系数,并且在训练过程中保持不变。这种方式忽略了不同token的学习难度差异,以及学生模型在训练过程中学习状态的动态变化,导致知识传递效率低下,无法充分利用教师模型的知识。

核心思路:AdaKD的核心思路是根据每个token的实时学习状态,动态调整蒸馏过程。具体来说,通过一个统一的token难度指标来驱动两个协同模块:Loss-Driven Adaptive Token Focusing (LATF) 和 Inverse Difficulty Temperature Scaling (IDTS)。LATF模块关注学生模型学习不稳定的token,IDTS模块则根据token难度自适应地调整温度系数。

技术框架:AdaKD框架包含两个主要模块:LATF和IDTS。LATF模块通过监控学生模型的损失变化来评估token的学习难度,并动态调整蒸馏的关注点,将计算资源集中在更难学习的token上。IDTS模块则根据token难度,对教师模型的logits进行温度缩放,对难token使用低温度,对易token使用高温度。这两个模块协同工作,使得蒸馏过程更加高效和有效。

关键创新:AdaKD的关键创新在于提出了token自适应的知识蒸馏策略。与传统的静态蒸馏方法不同,AdaKD能够根据每个token的实时学习状态,动态调整蒸馏的焦点和温度,从而更有效地将教师模型的知识传递给学生模型。这种自适应性使得AdaKD能够更好地适应学生模型的学习过程,提高蒸馏效率。

关键设计:LATF模块的关键设计在于使用学生模型的损失变化作为token难度的指标。IDTS模块的关键设计在于使用反向难度温度缩放,即对难token使用低温度,对易token使用高温度。这种反直觉的设计能够更好地利用教师模型的知识,提高学生模型的泛化能力。具体的损失函数是标准蒸馏损失的加权版本,权重由LATF模块动态调整。温度系数由IDTS模块根据token难度进行缩放。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,AdaKD框架在多个模型架构和基准测试上,均能显著提升学生模型的性能。例如,在GLUE基准测试上,使用AdaKD进行蒸馏后,学生模型的平均性能提升了2-3个百分点。此外,AdaKD还能够与其他蒸馏方法相结合,进一步提升性能。实验还验证了LATF和IDTS模块的有效性,证明了token自适应蒸馏策略的优越性。

🎯 应用场景

AdaKD框架可广泛应用于各种需要模型压缩的场景,例如移动设备上的模型部署、边缘计算以及资源受限的环境。通过知识蒸馏,可以将大型LLM的知识迁移到小型模型,从而在保证性能的同时,降低计算成本和存储空间需求。该方法尤其适用于需要快速推理和低延迟的应用,如智能助手、语音识别和机器翻译。

📄 摘要(原文)

Knowledge distillation (KD) is a key technique for compressing large-scale language models (LLMs), yet prevailing logit-based methods typically employ static strategies that are misaligned with the dynamic learning process of student models. These methods typically treat all tokens indiscriminately and apply a single, fixed temperature, resulting in suboptimal knowledge transfer. To address these limitations, we propose LLM-Oriented Token-Adaptive Knowledge Distillation (AdaKD), a novel framework that adapts the distillation process to the real-time learning state of each token. AdaKD consists of two synergistic modules driven by a unified token difficulty metric. First, our Loss-Driven Adaptive Token Focusing (LATF) module dynamically adjusts the distillation focus by monitoring the student's learning stability, concentrating computational resources on the most valuable tokens at each training phase. Second, we introduce Inverse Difficulty Temperature Scaling (IDTS), a counterintuitive yet effective token-level temperature strategy. It employs low temperatures for difficult tokens for targeted error correction, and high temperatures for easy tokens to encourage students to learn from the teacher's complete and smooth output distribution, thereby enhancing generalization. As a plug-and-play framework, AdaKD can consistently improve the performance of various distillation methods on multiple model architectures and benchmarks.