SWITCH: Studying with Teacher for Knowledge Distillation of Large Language Models

📄 arXiv: 2410.19503v2 📥 PDF

作者: Jahyun Koo, Yerin Hwang, Yongil Kim, Taegwan Kang, Hyunkyung Bae, Kyomin Jung

分类: cs.CL

发布日期: 2024-10-25 (更新: 2025-04-22)

备注: NAACL 2025 Findings

💡 一句话要点

提出SWITCH，通过教师模型干预解决大语言模型知识蒸馏中的长序列偏差问题。

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

关键词: 知识蒸馏 大型语言模型 模型压缩 长序列生成 教师模型 学生模型 指令跟随 概率分布差异

📋 核心要点

  1. 现有知识蒸馏方法依赖学生模型生成的数据，但长序列生成易引入噪声和偏差，导致教师模型误导。
  2. SWITCH的核心思想是在学生模型生成序列时，根据教师模型和学生模型的差异，选择性地引入教师模型的指导。
  3. 实验表明，SWITCH在长序列生成任务上优于传统知识蒸馏方法，提升了模型性能。

📝 摘要（中文）

大型语言模型（LLMs）虽然取得了显著成功，但仍面临高推理成本和内存需求等挑战。为了解决这些问题，知识蒸馏（KD）已成为一种流行的模型压缩方法，其中学生生成的输出（SGOs）作为训练数据，因能减少训练和推理之间的不匹配而备受关注。然而，SGOs经常产生噪声和有偏序列，这可能导致教师模型的误导，尤其是在长序列中。为了缓解这些挑战，我们提出了一种新颖的方法SWITCH（Studying WIth TeaCHer for Knowledge Distillation），该方法在学生的序列生成过程中策略性地结合了教师模型。SWITCH识别教师模型和学生模型的token概率之间的差异，允许教师有选择地进行干预，特别是在更容易出现教师误导的长序列中。在三个模型系列和五个指令跟随数据集上的大量实验结果表明，SWITCH超越了传统的KD方法，尤其擅长生成长序列数据。

🔬 方法详解

问题定义：论文旨在解决知识蒸馏过程中，学生模型在生成长序列时，由于自身偏差和噪声，导致教师模型产生误导的问题。现有方法直接使用学生模型生成的输出作为训练数据，忽略了长序列生成过程中累积的误差，使得学生模型难以有效地学习教师模型的知识。

核心思路：SWITCH的核心思路是在学生模型生成序列的过程中，动态地评估学生模型和教师模型之间的差异，并根据差异的大小决定是否引入教师模型的指导。通过这种方式，可以有效地纠正学生模型在生成过程中的偏差，从而提高知识蒸馏的效率和效果。

技术框架：SWITCH方法主要包含以下几个步骤：1) 学生模型生成token序列；2) 计算学生模型和教师模型在每个token上的概率分布差异；3) 根据差异大小，决定是否使用教师模型生成的token替换学生模型生成的token；4) 使用替换后的token序列作为训练数据，对学生模型进行知识蒸馏。

关键创新：SWITCH的关键创新在于其动态干预机制。传统知识蒸馏方法通常只在训练开始时使用教师模型，而在训练过程中完全依赖学生模型。SWITCH则根据学生模型和教师模型的差异，动态地调整教师模型的参与程度，从而更好地利用教师模型的知识。

关键设计：SWITCH的关键设计包括：1) 使用KL散度或交叉熵等方法来衡量学生模型和教师模型之间的概率分布差异；2) 设置一个阈值，当概率分布差异超过该阈值时，才使用教师模型生成的token替换学生模型生成的token；3) 可以使用不同的策略来选择教师模型生成的token，例如选择概率最高的token或随机选择token。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，SWITCH在多个指令跟随数据集上优于传统的知识蒸馏方法。例如，在长序列生成任务上，SWITCH相比于基线方法取得了显著的性能提升。具体的数据指标（如BLEU值、ROUGE值等）在论文中进行了详细的展示，证明了SWITCH在知识蒸馏方面的有效性。

🎯 应用场景

SWITCH方法可应用于各种需要模型压缩的场景，例如移动设备上的自然语言处理、边缘计算等。通过知识蒸馏，可以将大型语言模型的知识迁移到小型模型上，从而降低模型的计算成本和内存需求，使其能够在资源受限的环境中运行。该方法尤其适用于需要生成长文本序列的任务，例如机器翻译、文本摘要等。

📄 摘要（原文）

Despite the success of Large Language Models (LLMs), they still face challenges related to high inference costs and memory requirements. To address these issues, Knowledge Distillation (KD) has emerged as a popular method for model compression, with student-generated outputs (SGOs) as training data being particularly notable for reducing the mismatch between training and inference. However, SGOs often produce noisy and biased sequences, which can lead to misguidance from the teacher model, especially in long sequences. To mitigate these challenges, we propose SWITCH (Studying WIth TeaCHer for Knowledge Distillation), a novel approach that strategically incorporates the teacher model during the student's sequence generation. SWITCH identifies discrepancies between the token probabilities of the teacher and student models, allowing the teacher to intervene selectively, particularly in long sequences that are more prone to teacher misguidance. Extensive experimental results across three model families and five instruction-following datasets show that SWITCH surpasses traditional KD methods, particularly excelling in the generation of long sequential data.