Comparing Kullback-Leibler Divergence and Mean Squared Error Loss in Knowledge Distillation

By Taehyeon Kim, Jaehoon Oh, Nakyil Kim, Sangwook Cho, Se-Young Yun. 2021

337 citations recorded by the source corpus.

Knowledge distillation (KD), transferring knowledge from a cumbersome teacher model to a lightweight student model, has been investigated to design efficient neural architectures. Generally, the objective function of KD is the Kullback-Leibler (KL) divergence loss between the softened probability distributions of the teacher model and the student model with the temperature scaling hyperparameter τ. Despite its widespread use, few studies have discussed how such softening influences generalization. Here, we theoretically show that the KL divergence loss focuses on the logit matching when τ increases and the label matching when τ goes to 0 and empirically show that the logit matching is positively correlated to performance improvement in general. From this observation, we consider an intuitive KD loss function, the mean squared error (MSE) between the logit vectors, so that the student model can directly learn the logit of the teacher model. The MSE loss outperforms the KL divergence loss, explained by the penultimate layer representations difference between the two losses. Furthermore, we show that sequential distillation can improve performance and that KD, using the KL divergence loss with small τ particularly, mitigates the label noise. The code to reproduce the experiments is publicly available online at https://github.com/jhoon-oh/kd_data/.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Jingshan Tang
2. Yuanhao Yu
3. Li Gu
4. Huan Liu
5. Zhixiang Chi
6. Yang Wang

In Computer Vision and Pattern Recognition

1. Jingshan Tang
2. Yuanhao Yu
3. Li Gu
4. Huan Liu
5. Jun Chen
6. Zhixiang Chi
7. Yang Wang
8. Xiaogang Wang

Citations in the corpus

• Few-Shot Class-Incremental Learning via Entropy-Regularized Data-Free Replay (2022, methodology): “Note that we adopt mean square error (MSE) as objective in both phases other than KL divergence adopted in [20], because logit matching has better generalization capacity [16].”