知识蒸馏中 Kullback-Leibler 散度损失与均方误差损失的比较
机器学习
2021-05-20 v1 计算机视觉与模式识别
摘要
知识蒸馏(KD)将知识从笨重的教师模型迁移到轻量级学生模型,已被研究用于设计高效的神经架构。通常,KD的目标函数是带温度缩放超参数tau的教师模型与学生模型软化概率分布之间的Kullback-Leibler(KL)散度损失。尽管广泛使用,极少研究讨论此种软化对泛化的影响。在此,我们从理论上表明KL散度损失在tau增大时聚焦于logit匹配,在tau趋于0时聚焦于标签匹配,并从经验上表明logit匹配通常与性能提升正相关。由该观察,我们考虑一种直观的KD损失函数,即logit向量间的均方误差(MSE),使学生模型可直接学习教师模型的logit。MSE损失优于KL散度损失,这由两种损失在倒数第二层表征上的差异解释。此外,我们展示顺序蒸馏可提升性能,且KD特别是使用小tau的KL散度损失时可缓解标签噪声。复现实验的代码公开于 https://github.com/jhoon-oh/kd_data/。
引用
@article{arxiv.2105.08919,
title = {Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation},
author = {Taehyeon Kim and Jaehoon Oh and NakYil Kim and Sangwook Cho and Se-Young Yun},
journal= {arXiv preprint arXiv:2105.08919},
year = {2021}
}
备注
Proceedings of International Joint Conference on Artificial Intelligence (IJCAI), 2021