LoCa: Logit Calibration for Knowledge Distillation

📄 arXiv: 2409.04778v1 📥 PDF

作者: Runming Yang, Taiqiang Wu, Yujiu Yang

分类: cs.CL, cs.LG

发布日期: 2024-09-07

备注: Accepted by ECAI 2024

💡 一句话要点

提出LoCa：一种用于知识蒸馏的Logit校准方法，解决教师模型误导问题。

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

关键词: 知识蒸馏 模型压缩 Logit校准 暗知识 图像分类

📋 核心要点

  1. 现有知识蒸馏方法在教师模型预测错误时会误导学生模型，影响蒸馏效果。
  2. LoCa通过ground-truth标签校准教师模型的logits，纠正错误预测并保留有用的暗知识。
  3. 实验表明，LoCa在图像分类和文本生成任务上均能有效提升知识蒸馏的性能。

📝 摘要（中文）

知识蒸馏(KD)在模型压缩中扮演着重要角色，其目标是通过模仿教师模型来训练更好的学生模型。一种典型的方法是对齐输出logits。然而，我们发现一个常见的问题，即当基于教师logits的预测不符合标签时，学生会被误导，称之为误指导。同时，logits中还存在其他有用的暗知识，例如类区分性，这对蒸馏至关重要。在本文中，我们提出了一种简单而有效的Logit校准(LoCa)方法，该方法基于ground-truth标签校准来自教师模型的logits。核心思想是纠正预测（以解决误指导问题）并同时保持有用的暗知识。我们提出的LoCa不需要任何额外的参数。在图像分类和文本生成任务上的实验结果表明，LoCa可以有效地提高基线的性能。

🔬 方法详解

问题定义：知识蒸馏旨在将大型教师模型的知识迁移到小型学生模型。现有方法，特别是基于logits对齐的方法，容易受到教师模型错误预测的影响。当教师模型给出错误的logits分布时，学生模型会学习到错误的知识，导致性能下降。这种“误指导”问题是现有方法的痛点。

核心思路：LoCa的核心思路是利用ground-truth标签来校准教师模型的logits。通过校准，可以纠正教师模型的错误预测，避免学生模型学习到错误的知识。同时，LoCa的设计目标是在校准logits的同时，尽可能保留教师模型logits中包含的有用信息，例如类之间的区分性信息。

技术框架：LoCa方法不需要额外的网络结构或参数。其主要流程是在知识蒸馏过程中，首先利用ground-truth标签对教师模型的logits进行校准，然后使用校准后的logits作为目标，指导学生模型的训练。整个过程可以看作是在标准知识蒸馏框架中加入了一个logits校准的预处理步骤。

关键创新：LoCa的关键创新在于提出了一个简单有效的logits校准方法，该方法能够同时解决“误指导”问题并保留有用的暗知识。与现有方法相比，LoCa不需要额外的参数，易于实现，并且可以与现有的知识蒸馏方法相结合。

关键设计：LoCa的具体实现细节取决于所使用的知识蒸馏方法。一种常见的实现方式是，对于每个训练样本，首先计算教师模型的logits。然后，根据ground-truth标签，对logits进行校准。例如，可以增强正确类别的logits值，同时降低错误类别的logits值。校准后的logits将作为学生模型训练的目标。具体的校准函数可以根据实际情况进行选择，论文中可能给出了具体的函数形式。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文在图像分类和文本生成任务上验证了LoCa的有效性。实验结果表明，LoCa能够显著提升知识蒸馏的性能，优于现有的基线方法。具体的性能提升幅度取决于所使用的任务和数据集，但总体趋势是LoCa能够带来一致的性能提升。

🎯 应用场景

LoCa方法可广泛应用于各种需要模型压缩的场景，例如移动设备上的图像识别、自然语言处理等。通过知识蒸馏，可以将大型模型部署到资源受限的设备上，同时保持较高的性能。LoCa的提出进一步提升了知识蒸馏的性能，使得小型模型能够更好地学习大型模型的知识，具有重要的实际应用价值。

📄 摘要（原文）

Knowledge Distillation (KD), aiming to train a better student model by mimicking the teacher model, plays an important role in model compression. One typical way is to align the output logits. However, we find a common issue named mis-instruction, that the student would be misled when the predictions based on teacher logits do not follow the labels. Meanwhile, there is other useful dark knowledge in the logits such as the class discriminability, which is vital for distillation. In this paper, we propose a simple yet effective Logit Calibration (LoCa) method, which calibrates the logits from the teacher model based on the ground-truth labels. The key insight is to correct the prediction (to address the mis-instruction issue) and maintain useful dark knowledge simultaneously. Our proposed LoCa does not require any additional parameters. Empirical results on image classification and text generation tasks demonstrate that LoCa can effectively improve the performance of baselines.