[论文解读] Knowledge Distillation as Semiparametric Inference

Tri Dao, Govinda M. Kamath+2 more|arXiv (Cornell University)|Apr 20, 2021

Adversarial Robustness in Machine Learning参考文献 58被引用 4

一句话总结

本文将知识蒸馏建模为半参数推断问题，其中学生模型从教师模型提供的类别概率中学习，这些概率被视为未知贝叶斯类别概率的插补估计。通过应用半参数理论，作者推导出新的误差保证，并提出两种改进方法——交叉拟合（cross-fitting）和损失校正（loss correction），以减轻教师模型过拟合与欠拟合的负面影响，从而在表格数据和图像数据上实现一致的性能提升。

ABSTRACT

A popular approach to model compression is to train an inexpensive student model to mimic the class probabilities of a highly accurate but cumbersome teacher model. Surprisingly, this two-step knowledge distillation process often leads to higher accuracy than training the student directly on labeled data. To explain and enhance this phenomenon, we cast knowledge distillation as a semiparametric inference problem with the optimal student model as the target, the unknown Bayes class probabilities as nuisance, and the teacher probabilities as a plug-in nuisance estimate. By adapting modern semiparametric tools, we derive new guarantees for the prediction error of standard distillation and develop two enhancements -- cross-fitting and loss correction -- to mitigate the impact of teacher overfitting and underfitting on student performance. We validate our findings empirically on both tabular and image data and observe consistent improvements from our knowledge distillation enhancements.

研究动机与目标

• 解释为何知识蒸馏在使用教师模型软标签时，通常优于直接的学生模型训练。
• 将知识蒸馏形式化为一个半参数推断问题，以未知的贝叶斯类别概率为目标，教师模型的概率作为干扰估计量。
• 开发基于理论的改进方法，以缓解因教师模型过拟合与欠拟合导致的性能下降。
• 在表格数据与图像分类基准上，通过实证方法验证所提方法的有效性。

提出的方法

• 作者将知识蒸馏建模为半参数推断问题，将最优学生模型视为目标参数，真实贝叶斯类别概率视为非参数干扰成分。
• 使用教师模型输出作为干扰成分的插补估计，从而推导出学生预测误差的渐近偏差与方差界。
• 提出交叉拟合方法，通过数据分割来估计干扰参数，以减少教师模型过拟合带来的偏差。
• 提出损失校正方法，调整学生模型的训练目标，以补偿教师模型可能存在的欠拟合或校准不足问题。
• 利用现代半参数工具（包括影响函数与高效估计理论）推导理论保证。
• 该框架支持误差分解，并提供一种系统化方法来评估与提升学生模型的泛化能力，超越标准蒸馏方法。

实验结果

研究问题

• RQ1 为何使用教师模型的知识蒸馏通常能获得比直接监督训练更好的学生模型性能？
• RQ2 如何在知识蒸馏中形式化量化并缓解教师模型过拟合与欠拟合的影响？
• RQ3 能否将半参数推断理论应用于推导出更鲁棒的新知识蒸馏方法？
• RQ4 针对通过知识蒸馏训练的学生模型，可建立哪些预测误差的理论保证？

主要发现

• 所提出的交叉拟合与损失校正方法在表格数据与图像数据上均一致地提升了学生模型的准确率，有效减少了教师模型误差带来的偏差。
• 理论分析表明，标准知识蒸馏可能继承来自过拟合教师模型的偏差，而新方法通过正确估计干扰参数有效缓解了该问题。
• 损失校正通过调整学生模型的训练目标，使其更贴近真实的贝叶斯决策边界，从而有效降低教师模型欠拟合的影响。
• 实证结果表明，改进后的蒸馏方法在教师模型校准不足或过拟合时，显著优于标准知识蒸馏方法。