[论文解读] An Embarrassingly Simple Approach for Knowledge Distillation

Mengya Gao, Yujun Shen+6 more|arXiv (Cornell University)|Dec 5, 2018

Advanced Neural Network Applications参考文献 40被引用 12

一句话总结

该论文提出了一种分阶段知识蒸馏（SSKD）方法，这是一种简单但高效的方法，将知识蒸馏过程解耦为两个阶段：第一阶段，无需使用标签，将教师网络的中间特征表示蒸馏到学生主干网络；第二阶段，仅使用真实标签对任务特定的分类头进行微调。SSKD 消除了对损失权重超参数调优的需求，并在 ImageNet、CIFAR-100、COCO 目标检测和 IJB-A 人脸识别基准上实现了最先进性能，其中在 COCO 上最高提升 2.8% AP，在 ImageNet 上最高提升 1.8% 准确率。

ABSTRACT

Knowledge Distillation (KD) aims at improving the performance of a low-capacity student model by inheriting knowledge from a high-capacity teacher model. Previous KD methods typically train a student by minimizing a task-related loss and the KD loss simultaneously, using a pre-defined loss weight to balance these two terms. In this work, we propose to first transfer the backbone knowledge from a teacher to the student, and then only learn the task-head of the student network. Such a decomposition of the training process circumvents the need of choosing an appropriate loss weight, which is often difficult in practice, and thus makes it easier to apply to different datasets and tasks. Importantly, the decomposition permits the core of our method, Stage-by-Stage Knowledge Distillation (SSKD), which facilitates progressive feature mimicking from teacher to student. Extensive experiments on CIFAR-100 and ImageNet suggest that SSKD significantly narrows down the performance gap between student and teacher, outperforming state-of-the-art approaches. We also demonstrate the generalization ability of SSKD on other challenging benchmarks, including face recognition on IJB-A dataset as well as object detection on COCO dataset.

研究动机与目标

• 为了解决知识蒸馏中超参数敏感的问题，特别是难以调节任务损失与 KD 损失之间权重平衡的问题。
• 通过将训练过程解耦为两个独立阶段，提升知识蒸馏在多样化数据集和任务上的泛化能力与稳定性。
• 证明渐进式、分阶段的中间特征蒸馏相比单步蒸馏能带来更优的知识迁移效果。
• 将知识蒸馏扩展到非分类任务，如目标检测和人脸识别，这些任务中先前方法并不直接适用。
• 通过在 COCO 上对 FPN 和 RetinaNet 应用 SSKD，建立目标检测领域知识蒸馏的新基准。

提出的方法

• 将学生网络的训练过程分解为两个阶段：(1) 在不使用真实标签的情况下，将教师网络的中间特征表示蒸馏到学生主干网络；(2) 固定主干网络，仅使用带标签数据对任务特定分类头进行训练。
• 实现分阶段知识蒸馏（SSKD），即按网络阶段顺序依次进行蒸馏，逐个子网络地模仿特征。
• 在第一阶段中，使用特征级蒸馏损失（如 MSE 或 KL 散度）比较教师与学生网络对应层之间的特征。
• 在蒸馏完成后固定学生主干网络，并在第二阶段使用交叉熵或其他任务特定损失，仅对任务头进行微调。
• 将 SSKD 应用于多种架构与任务，包括在 ImageNet 和 CIFAR-100 上的 ResNets，以及在 COCO 目标检测任务中基于 FPN 和 RetinaNet 的 ResNets。
• 通过在每个阶段保持网络大部分参数冻结，确保训练效率，从而避免训练时间显著增加。

实验结果

研究问题

• RQ1 将知识蒸馏解耦为骨干特征模仿与任务头微调是否能消除对损失权重超参数调优的需求？
• RQ2 分阶段、渐进式地蒸馏中间特征是否能带来优于端到端联合训练（固定损失权重）的性能表现？
• RQ3 所提出的方法是否能在图像分类、目标检测和人脸识别等多样化任务上实现良好泛化？
• RQ4 在 ImageNet 和 COCO 等标准基准上，SSKD 与最先进知识蒸馏方法相比表现如何？
• RQ5 SSKD 是否能有效应用于 FPN 和 RetinaNet 等复杂检测架构，这些架构中 KD 尚未被充分探索？

主要发现

• 在 ImageNet 上，以 ResNet-18 为学生、ResNet-34 为教师时，SSKD 将 top-1 准确率提升 1.8%，达到 80.5%，超过第二好的方法。
• 在 COCO 目标检测任务中，使用 FPN 时，SSKD 将 ResNet-50 的 AP 从 37.7% 提升至 40.5%，提升 2.8%，蒸馏来源为 ResNet-152。
• 在 RetinaNet 上，SSKD 将 ResNet-18 的 AP 从 32.4% 提升至 35.6%，提升 3.2%，蒸馏来源为 ResNet-152。
• 在 IJB-A 人脸识别任务中，SSKD 展现出强大的泛化能力，持续优于基线方法。
• 在 COCO 上，SSKD 在所有目标尺寸类别（AP_S、AP_M、AP_L）中均实现一致性能提升，表明对尺度变化具有鲁棒性。
• 消融实验证实，分阶段蒸馏优于单步蒸馏，且两阶段训练方案在稳定性和有效性上均优于联合优化。