Exploring and Enhancing the Transfer of Distribution in Knowledge Distillation for Autoregressive Language Models

📄 arXiv: 2409.12512v2 📥 PDF

作者: Jun Rao, Xuebo Liu, Zepeng Lin, Liang Ding, Jing Li, Dacheng Tao, Min Zhang

分类: cs.CL

发布日期: 2024-09-19 (更新: 2024-09-20)

💡 一句话要点

提出在线知识蒸馏(OKD)方法，提升自回归语言模型蒸馏效率与性能。

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

关键词: 知识蒸馏 自回归语言模型 在线学习 模型压缩 模型加速

📋 核心要点

  1. 现有知识蒸馏方法，如反向KL散度和学生生成输出(SGO)，在捕捉教师模型行为和计算效率上存在不足。
  2. 论文提出在线知识蒸馏(OKD)方法，通过集成小型在线模块到教师网络，实现教师模型动态适应学生分布。
  3. 实验结果表明，OKD在多个生成数据集上优于现有方法，并在各种模型架构和大小上实现了显著的训练时间缩短。

📝 摘要（中文）

知识蒸馏(KD)是一种通过训练小型学生模型模仿大型教师模型来压缩模型的技术。KD在自回归语言模型中的成功主要依赖于用于模式搜索的反向KL散度和用于对抗暴露偏差的学生生成输出(SGO)。我们的理论分析和实验验证表明，虽然反向KL有效地模仿了教师分布的某些特征，但它未能捕捉到其大部分行为。相反，SGO会产生更高的计算成本，并在优化方面带来挑战，尤其是在学生模型远小于教师模型时。这些约束主要归因于教师模型的不可变分布，该分布无法自适应地调整到不同大小的模型。我们引入了在线知识蒸馏(OKD)，其中教师网络集成了小型在线模块，与学生模型同时训练。这种策略消除了对on-policy采样的需求，并且仅需要在训练期间对教师在线模块的参数进行最小更新，从而允许动态适应学生分布以更好地进行蒸馏。跨多个生成数据集的广泛结果表明，OKD在各种模型架构和大小中实现了或超过了领先方法的性能，并将训练时间减少了高达四倍。

🔬 方法详解

问题定义：现有知识蒸馏方法在自回归语言模型中存在局限性。反向KL散度虽然能模仿教师分布的部分特征，但无法完全捕捉其行为。学生生成输出(SGO)虽然能缓解暴露偏差，但计算成本高昂，且在学生模型远小于教师模型时优化困难。根本原因是教师模型的分布是固定的，无法根据学生模型的大小进行调整。

核心思路：论文的核心思路是让教师模型能够动态适应学生模型的分布。通过在教师模型中集成小型在线模块，并与学生模型同步训练，使得教师模型能够根据学生模型的学习情况进行调整，从而提供更有效的蒸馏信号。

技术框架：OKD的整体框架包括一个学生模型和一个教师模型，教师模型包含一个预训练的主干网络和一个小型在线模块。在训练过程中，学生模型和教师模型的在线模块同时进行更新。学生模型的目标是模仿教师模型的输出，而教师模型的在线模块的目标是更好地适应学生模型的分布。这种同步训练的方式使得教师模型能够动态地为学生模型提供指导。

关键创新：OKD的关键创新在于引入了在线学习机制，使得教师模型能够动态适应学生模型的分布。这与传统的知识蒸馏方法中教师模型分布固定不变形成了鲜明对比。通过这种动态适应，OKD能够更有效地将知识从教师模型传递到学生模型。

关键设计：OKD的关键设计包括：1) 小型在线模块的结构，可以是简单的线性层或更复杂的神经网络；2) 损失函数的设计，通常包括模仿损失（例如KL散度）和正则化项，以防止在线模块过度拟合；3) 教师模型主干网络的参数更新策略，通常采用较小的学习率或固定参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，OKD在多个生成数据集上取得了显著的性能提升，并且在各种模型架构和大小上都优于现有的知识蒸馏方法。更重要的是，OKD能够将训练时间减少高达四倍，这使得训练大型语言模型的学生模型变得更加高效。

🎯 应用场景

该研究成果可应用于各种需要模型压缩的场景，例如移动设备上的自然语言处理、边缘计算和低资源环境下的模型部署。通过OKD，可以在保证模型性能的同时，显著降低模型的大小和计算复杂度，从而提高模型的实用性和适用性。未来，该方法还可以扩展到其他类型的模型和任务中。

📄 摘要（原文）

Knowledge distillation (KD) is a technique that compresses large teacher models by training smaller student models to mimic them. The success of KD in auto-regressive language models mainly relies on Reverse KL for mode-seeking and student-generated output (SGO) to combat exposure bias. Our theoretical analyses and experimental validation reveal that while Reverse KL effectively mimics certain features of the teacher distribution, it fails to capture most of its behaviors. Conversely, SGO incurs higher computational costs and presents challenges in optimization, particularly when the student model is significantly smaller than the teacher model. These constraints are primarily due to the immutable distribution of the teacher model, which fails to adjust adaptively to models of varying sizes. We introduce Online Knowledge Distillation (OKD), where the teacher network integrates small online modules to concurrently train with the student model. This strategy abolishes the necessity for on-policy sampling and merely requires minimal updates to the parameters of the teacher's online module during training, thereby allowing dynamic adaptation to the student's distribution to make distillation better. Extensive results across multiple generation datasets show that OKD achieves or exceeds the performance of leading methods in various model architectures and sizes, reducing training time by up to fourfold.