[论文解读] Multi-head Knowledge Distillation for Model Compression

Huan Wang, Suhas Lohit+2 more|arXiv (Cornell University)|Dec 5, 2020

Advanced Neural Network Applications参考文献 24被引用 5

一句话总结

该论文提出了一种名为多头知识蒸馏（MHKD）的简单而有效的方法，用于模型压缩。该方法在教师网络和学生网络的中间层添加辅助分类头。通过最小化这些分类头输出之间的KL散度，MHKD在多种架构和数据集上均提升了学生模型的性能，始终优于以往的知识蒸馏方法，并在结合对比学习时实现了最先进（SOTA）的结果。

ABSTRACT

Several methods of knowledge distillation have been developed for neural network compression. While they all use the KL divergence loss to align the soft outputs of the student model more closely with that of the teacher, the various methods differ in how the intermediate features of the student are encouraged to match those of the teacher. In this paper, we propose a simple-to-implement method using auxiliary classifiers at intermediate layers for matching features, which we refer to as multi-head knowledge distillation (MHKD). We add loss terms for training the student that measure the dissimilarity between student and teacher outputs of the auxiliary classifiers. At the same time, the proposed method also provides a natural way to measure differences at the intermediate layers even though the dimensions of the internal teacher and student features may be different. Through several experiments in image classification on multiple datasets we show that the proposed method outperforms prior relevant approaches presented in the literature.

研究动机与目标

• 通过利用中间层监督来改进模型压缩中的知识蒸馏。
• 解决蒸馏过程中学生网络与教师网络之间特征维度不匹配的挑战。
• 开发一种简单、即插即用的方法，以提升学生模型性能，而无需引入复杂的知识定义或网络架构。
• 确保在多种网络架构和数据集上具有鲁棒性，包括跨架构组合。
• 提供一种可与其他蒸馏技术（如对比学习）兼容的方法，以实现进一步的性能提升。

提出的方法

• 该方法在教师网络和学生网络的多个中间层添加辅助分类头。
• 对于每个中间层，使用KL散度损失比较学生和教师辅助分类头的输出。
• 在学生训练过程中，将KL散度损失与标准交叉熵损失结合，以指导特征学习。
• 由于辅助分类头被训练为生成可比较的输出，无论输入维度如何，该方法自然地处理了特征维度不匹配的问题。
• 辅助分类头在蒸馏过程中进行训练，但在推理阶段被移除，仅保留学生网络的主分类头。
• 该方法可与其他蒸馏技术兼容，例如带有记忆库的对比学习，从而实现进一步的性能提升。

实验结果

研究问题

• RQ1 通过辅助分类头实现的中间层监督能否提升模型压缩中知识蒸馏的性能？
• RQ2 MHKD在多种教师-学生架构组合（包括跨架构对）中的表现如何？
• RQ3 与FitNet、RKD和CC等现有蒸馏方法相比，MHKD在准确率和鲁棒性方面是否表现更优？
• RQ4 MHKD能否有效与其它蒸馏技术（如带有负样本对的对比学习）结合？
• RQ5 MHKD对辅助分类头所选中间层的敏感性如何？使用多个分类头是否能进一步提升性能？

主要发现

• MHKD在所有评估的数据集和网络组合中始终优于标准知识蒸馏（KD）。
• 在CIFAR-100上，当教师网络为ResNet32x4、学生网络为ShuffleNetV2时，MHKD实现了96.68%的top-1准确率，较KD最高提升1.2个百分点。
• 对于VGG13/MobileNetV2组合，与FitNet、RKD和CC相比，MHKD的准确率提升超过1%。
• 消融实验表明，使用多个辅助分类头（如头1–3）的性能优于任意单个头，从而消除了手动选择分类头的需要。
• 当与CRD结合时，MHKD实现了最先进性能，证明了其与对比学习具有良好的兼容性和协同效应。
• 该方法在不同数据集上表现出强鲁棒性：在CIFAR-100和SVHN上均优于所有基线方法，即使其他方法在不同数据集上的性能趋势存在差异。