[论文解读] Efficient Knowledge Distillation from Model Checkpoints

Chaofei Wang, Qisen Yang+3 more|arXiv (Cornell University)|Oct 12, 2022

Machine Learning and Data Classification被引用 7

一句话总结

本文提出使用中间模型检查点（而非完全收敛的模型）作为知识蒸馏中的更优教师模型。基于信息瓶颈理论，研究发现中间模型通过保留更高的输入-特征互信息，能保存更多的“暗知识”，从而带来更好的学生模型性能；实验表明，最优的中间教师模型在性能上优于最终模型和强集成模型，所提出的筛选算法使CIFAR-100上的蒸馏准确率最高提升0.8%。

ABSTRACT

Knowledge distillation is an effective approach to learn compact models (students) with the supervision of large and strong models (teachers). As empirically there exists a strong correlation between the performance of teacher and student models, it is commonly believed that a high performing teacher is preferred. Consequently, practitioners tend to use a well trained network or an ensemble of them as the teacher. In this paper, we make an intriguing observation that an intermediate model, i.e., a checkpoint in the middle of the training procedure, often serves as a better teacher compared to the fully converged model, although the former has much lower accuracy. More surprisingly, a weak snapshot ensemble of several intermediate models from a same training trajectory can outperform a strong ensemble of independently trained and fully converged models, when they are used as teachers. We show that this phenomenon can be partially explained by the information bottleneck principle: the feature representations of intermediate models can have higher mutual information regarding the input, and thus contain more "dark knowledge" for effective distillation. We further propose an optimal intermediate teacher selection algorithm based on maximizing the total task-related mutual information. Experiments verify its effectiveness and applicability.

研究动机与目标

• 挑战传统假设，即高性能、完全训练的教师模型能为学生模型带来最佳效果。
• 探究为何尽管准确率较低，中间模型仍可作为比完全收敛模型更优的教师模型。
• 基于信息论，提出一种理论基础扎实且高效的最优中间教师模型选择方法。
• 证明基于单次训练轨迹的弱快照集成中间模型，可超越强独立训练模型的集成效果。
• 通过战略性地选择中间模型，降低训练成本，同时提升蒸馏性能。

提出的方法

• 该方法使用非参数熵估计器，在每个训练检查点处估计特征与输入之间的互信息（I(X;F)）以及特征与目标之间的互信息（I(Y;F)）。
• 基于最大化 I(X;F) 与 I(Y;F) 之和的准则，构建最优中间教师模型选择标准，以平衡知识丰富度与任务相关性。
• 所提筛选算法可在单一训练轨迹中识别最佳检查点，无需进行多次独立教师模型的训练。
• 该方法与标准KD损失函数兼容，可应用于任意学生-教师架构。
• 通过归一化的互信息得分比较不同检查点，选择总任务相关信息量最大的模型。
• 该方法通过用单个、最优时间点的中间检查点替代昂贵的全训练模型或集成教师模型，实现高效蒸馏。

实验结果

研究问题

• RQ1 为何中间模型检查点尽管准确率较低，却常在知识蒸馏中优于完全收敛的模型？
• RQ2 能否通过单次训练轨迹中获得的弱快照集成中间模型，超越由独立训练模型构成的强集成模型？
• RQ3 信息瓶颈原理如何解释中间模型更优的知识迁移能力？
• RQ4 选择能最大化蒸馏性能的中间教师模型的最优策略是什么？
• RQ5 能否通过单个、最优时间点的检查点替代昂贵的集成模型，同时保持或提升学生模型的准确率？

主要发现

• 通过互信息最大化选择的最优中间教师模型，在CIFAR-100上使用WRN-40-1学生模型时，KD准确率达到73.26%，比完全收敛教师模型（72.68%）高出0.58%。
• 对于ResNet-110，最优中间教师模型达到72.63%的KD准确率，超过全模型的72.48%和0.7x检查点的72.46%。
• 在所有评估架构中，半程检查点（0.5x）始终优于完全收敛模型，CIFAR-100上的平均准确率提升达0.31%。
• 基于单次训练轨迹的弱快照集成中间模型，其性能超越了强独立训练模型的全集成，证明了所提筛选策略的有效性。
• 该方法通过消除训练多个全模型或集成模型的需求，显著降低训练成本，同时仍实现更优的蒸馏性能。
• 研究证实，过度训练会抑制类别相关性信息，从而降低可用于蒸馏的“暗知识”质量。