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

1 Introduction

Knowledge distillation (KD) is one of the most potent model compression techniques in which knowledge is transferred from a cumbersome model (teacher) to a single small model (student) [Hinton et al., 2015]. In general, the objective of training a smaller student network in the KD framework is formed as a linear summation of two losses: cross-entropy (CE) loss with "hard" targets, which are onehot ground-truth vectors of the samples, and KullbackLeibler (KL) divergence loss with the teacher's predictions. Specifically, KL divergence loss has achieved considerable success by controlling the softness of "soft" targets via the temperature-scaling hyperparameter τ. Utilizing a larger value for this hyperparameter τ makes the softmax vectors smooth over classes. Such a re-scaled output probability vector by τ is called the softened probability distribution, or the softened softmax [Maddison et al., 2016; Jang et al., 2016]. Recent KD has evolved to give more importance to the KL divergence loss to improve performance when balancing the objective between CE loss and KL divergence loss [Hinton et al., 2015; Tian et al., 2019]. Hence, we focus on training a student network based solely on the KL divergence loss.

∗The authors contributed equally to this paper.

Recently, there has been an increasing demand for investigating the reasons for the superiority of KD. [Yuan et al., 2020; Tang et al., 2020] empirically showed that the facilitation of KD is attributed to not only the privileged information on similarities among classes but also the label smoothing regularization. In some cases, the theoretical reasoning for using "soft" targets is clear. For example, in deep linear neural networks, KD not only accelerates the training convergence but also helps in reliable training [Phuong and Lampert, 2019]. In the case of self-distillation (SD), where the teacher model and the student model are the same, such an approach progressively restricts the number of basis functions to represent the solution [Mobahi et al., 2020]. However, there is still a lack of understanding of how the degree of softness affects the performance.

In this paper, we first investigate the characteristics of a student trained with KL divergence loss with various τ, both theoretically and empirically. We find that the student's logit (i.e., an input of the softened softmax function) more closely resembles the teacher's logit as τ increases, but not completely. Therefore, we design a direct logit learning scheme by replacing the KL divergence loss between the softened probability distributions of a teacher network and a student network (Figure 1(a)) with the mean squared error (MSE) loss between the student's logit and the teacher's logit (Figure 1(b)). Our contributions are summarized as follows:

• We investigate the role of the softening hyperparameter τ theoretically and empirically. A large τ, that is, strong softening, leads to logit matching, whereas a small τ results in training label matching. In general, logit matching has a better generalization capacity than label matching.
• We propose a direct logit matching scheme with the MSE loss and show that the KL divergence loss with any value of τ cannot achieve complete logit matching as much as the MSE loss. Direct training results in the best performance in our experiments.
• We theoretically show that the KL divergence loss makes the model's penultimate layer representations elongated than those of the teacher, while the MSE loss does not. We visualize the representations using the method proposed by [M¨uller et al., 2019].
• We show that sequential distillation, the MSE loss after the KL divergence loss, can be a better strategy than direct distillation when the capacity gap between the teacher and the student is large, which contrasts [Cho and Hariharan, 2019].
• We observe that the KL divergence loss, with low τ in particular, is more efficient than the MSE loss when the data have incorrect labels (noisy label). In this situation, extreme logit matching provokes bad training, whereas the KL divergence loss mitigates this problem.

2 RelatedWork

2.1 Knowledge Distillation

KD has been extended to a wide range of methods. One attempted to distill not only the softened probabilities of the teacher network but also the hidden feature vector so that the student could be trained with rich information from the teacher [Romero et al., 2014; Zagoruyko and Komodakis, 2016a; Srinivas and Fleuret, 2018; Kim et al., 2018; Heo et al., 2019b; Heo et al., 2019a]. The KD approach can be leveraged to reduce the generalization errors in teacher models (i.e., self-distillation; SD) [Zhang et al., 2019; Park et al., 2019] as well as model compressions. In the generative models, a generator can be compressed by distilling the latent features from a cumbersome generator [Aguinaldo et al., 2019].

To explain the efficacy of KD, [Furlanello et al., 2018] asserted that the maximum value of a teacher's softmax probability was similar to weighted importance by showing that permuting all of the non-argmax elements could also improve performance. [Yuan et al., 2020] argued that "soft" targets served as a label smoothing regularizer rather than as a transfer of class similarity by showing that a poorly trained or smaller-size teacher model can boost performance. Recently, [Tang et al., 2020] modified the conjecture in [Furlanello et al., 2018] and showed that the sample was positively reweighted by the prediction of the teacher's logit vector.

2.2 Label Smoothing

Smoothing the label y is a common method for improving the performance of deep neural networks by preventing the overconfident predictions [Szegedy et al., 2016]. Label smoothing is a technique that facilitates the generalization by replacing a ground-truth one-hot vector y with a weighted mixture of hard targets yLS :

where k indicates the index, and β is a constant. This implicitly ensures that the model is well-calibrated [M¨uller et al., 2019]. Despite its improvements, [M¨uller et al., 2019] observed that the teacher model trained with LS improved its performance, whereas it could hurt the student's performance. [Yuan et al., 2020] demonstrated that KD might be a category of LS by using the adaptive noise, i.e., KD is a label regularization method.

3 Preliminaries: KD

We denote the softened probability vector with a temperaturescaling hyperparameter τ for a network f as pf (τ), given a sample x. The k-th value of the softened probability vector pf (τ) is denoted by pfk (τ) = exp(zfk /τ) ∑Kj=1 exp(zf j /τ) , where zf k is the k-th value of the logit vector zf , K is the number of classes, and exp(·) is the natural exponential function. Then, given a sample x, the typical loss L for a student network is a linear combination of the cross-entropy loss LCE and the Kullback-

Leibler divergence loss LKL:

where s indicates the student network, t indicates the teacher network, y is a one-hot label vector of a sample x, and α is a hyperparameter of the linear combination. For simplicity of notation, LCE(ps (1), y) and LKL(ps (τ ), pt (τ )) are denoted by LCE and LKL, respectively. The standard choices are α = 0.1 and τ ∈ {3, 4, 5} [Hinton et al., 2015; Zagoruyko and Komodakis, 2016a].

In [Hinton et al., 2015], given a single sample x, the gradient of LKL with respect to zsk is as follows:

When τ goes to ∞, this gradient is simplified with the approximation, i.e., exp(zf k /τ ) ≈ 1 + zf k /τ :

Here, the authors assumed the zero-mean teacher and student logit, i.e., ∑j zt j = 0 and ∑j zs j = 0, and hence ∂LKL ∂zsk ≈ 1K (zsk -ztk ). This indicates that minimizing LKL is equivalent to minimizing the mean squared error LMSE, that is, ||zs - zt ||22 , under a sufficiently large temperature τ and the zero-mean logit assumption for both the teacher and the student.

However, we observe that this assumption does not seem appropriate and hinders complete understanding by ignoring the hidden term in LKL when τ increases. Figure 2 describes the histograms for the magnitude of logit summations on the training dataset. The logit summation histogram from the teacher network trained with LCE is almost zero (Figure 2(a)), whereas that from the student network trained with LKL using the teacher's knowledge goes far from zero as τ

(a) Teacher.

increases (Figure 2(b)). This is discussed in detail in Section 4.2.

3.1 Experimental Setup

In this paper, we used an experimental setup similar to that in [Heo et al., 2019a; Cho and Hariharan, 2019]: image classification on CIFAR-100 with a family of Wide-ResNet (WRN) [Zagoruyko and Komodakis, 2016b] and ImageNet with a family of of ResNet (RN) [He et al., 2016]. We used a standard PyTorch SGD optimizer with a momentum of 0.9, weight decay, and apply standard data augmentation. Other than those mentioned, the training settings from the original papers [Heo et al., 2019a; Cho and Hariharan, 2019] were used.

4 Relationship between LKL and LMSE

In this section, we conduct extensive experiments and systematically break down the effects of τ in LKL based on theoretical and empirical results. Then, we highlight the relationship between LKL and LMSE. Then, we compare the models trained with LKL and LMSE in terms of performance and penultimate layer representations. Finally, we investigate the effects of a noisy teacher on the performance according to the objective.

4.1 Hyperparameter τ in LKL

We investigate the training and test accuracies according to the change in α in L and τ in LKL (Figure 3). First, we empirically observe that the generalization error of a student model decreases as α in L increases. This means that "soft" targets are more efficient than "hard" targets in training a student if "soft" targets are extracted from a well-trained teacher. This result is consistent with prior studies that addressed the efficacy of "soft" targets [Furlanello et al., 2018; Tang et al., 2020]. Therefore, we focus on the situation where "soft" targets are used to train a student model solely, that is, α = 1.0, in the remainder of this paper.

When α = 1.0, the generalization error of the student model decreases as τ in LKL increases. These consistent tendencies according to the two hyperparameters, α and τ, are the same across various teacher-student pairs. To explain this phenomenon, we extend the gradient analysis in Section 3 without the assumption that the mean of the logit vector is zero.

Proposition 1. Let K be the number of classes in the dataset, and 1[·] be the indicator function, which is 1 when the statement inside the brackets is true and 0 otherwise. Then,

Proposition1 explains the consistent trends as follows. In the course of regularizing LKL with sufficiently large τ, the student model attempts to imitate the logit distribution of the teacher model. Specifically, a larger τ is linked to a larger LKL, making the logit vector of the student similar to that of the teacher (i.e., logit matching). Hence, "soft" targets are being fully used as τ increases. This is implemented using a handcrafted gradient (top row of Figure 3). On the other hand, when τ is close to 0, the gradient of LKL does not consider the logit distributions and only identifies whether the student and the teacher share the same output (i.e., label matching), which transfers limited information. In addition, there is a scaling issue when τ approaches 0. As τ decreases, LKL increasingly loses its quality and eventually becomes less involved in learning. The scaling problem can be easily fixed by multiplying 1/τ by LKL when τ is close to zero.

From this proposition, it is recommended to modify the original LKL in Eq. (2), considering τ ∈ (0,∞) as follows:

The key difference between our analysis and the preliminary analysis on a sufficiently large τ, i.e., in Eq. (5), is that the latter term is generated by removing the existing assumption on the logit mean, which is discussed in Section 4.2 at the loss-function level.

4.2 Extensions fromLKL to LMSE

In this subsection, we focus on Eq. (5) to investigate the reason as to why the efficacy of KD is observed when τ is greater than 1 in the KD environment, as shown in Figure 3. Eq. (5) can be understood as a biased regression of the vector expression as follows:

where 1 is a vector whose elements are equal to one. Furthermore, we can derive the relationship between limτ→∞LKL and LMSE as follows:

In Figure 2(a), the sum of the logit values of the teacher model is almost zero. With the teacher's logit value, δ∞ is approximated as - 1 2K2 (∑Kj=1 zs j )2 . Therefore, δ∞ can make the logit mean of the student trained with LKL depart from zero. From this analysis, it is unreasonable to assume that the student's logit mean is zero. We empirically find that the student's logit mean breaks the existing assumption as τ increases (Figure 2(b)). In summary, δ∞ hinders complete logit matching by shifting the mean of the elements in the logit. In other words, as derived from Eq. (8), optimizing LKL with sufficiently large τ is equivalent to optimizing LMSE with the additional regularization term δ∞, and it seems to rather hinder logit matching.

Therefore, we propose the direct logit learning objective for enhanced logit matching as follows:

Although this direct logit learning was used in [Ba and Caruana, 2013; Urban et al., 2016], they did not investigate the wide range of temperature scaling and the effects of MSE in the latent space. In this respect, our work differs.

4.3 Comparison of LKL and LMSE

We empirically compared the objectives LKL and LMSE in terms of performance gains and measured the distance between the logit distributions. Following the previous analysis, we also focused on "soft" targets in L′ . Table 2 presents the top-1 test accuracies on CIFAR-100 according to the student learning scheme for various teacher-student pairs. The students trained with LCE are vanilla models without a teacher. The students trained with LKL or LMSE are trained following the KD framework without using the "hard" targets, i.e., α = 1 in L and L′ , respectively. It is shown that distillation with LMSE, that is, direct logit distillation without hindering term δ∞, is the best training scheme for various teacher-student pairs. We also found the consistent improvements in ensemble distillation [Hinton et al., 2015]. For the ensemble distillation using MSE loss, an ensemble of logit predictions (i.e., an average of logit predictions) are used by multiple teachers. We obtained the test accuracy of WRN162 (75.60%) when the WRN16-4, WRN-28-4, and WRN-40-6 models were used as ensemble teachers in this manner. Moreover, the model trained with LMSE has similar or better performance when compared to other existing KD methods, as described in Table 3. 1

Furthermore, to measure the distance between the student's logit zs and the teacher's logit zt sample by sample, we describe the probabilistic density function (pdf) from the histogram for ‖zs - zt ‖2 on the CIFAR-100 training dataset (Figure 4(a)). The logit distribution of the student with a large τ is closer to that of the teacher than with a small τ when LKL is used. Moreover, LMSE is more efficient in transferring the teacher's information to a student than LKL. Optimizing LMSE aligns the student's logit with the teacher's logit. On the other hand, when τ becomes significantly large, LKL has the δ∞, and optimizing δ∞ makes the student's logit mean deviate from that of the teacher's logit mean.

1 We excluded the additional experiments for the replacement with MSE loss in feature-based distillation methods. It is difficult to add the MSE loss or replace the KL loss with MSE loss in the existing works because of the sensitivity to hyperparameter optimization. Their methods included various types of hyperparameters that need to be optimized for their settings.

We further investigate the effect of δ∞ on the penultimate layer representations (i.e., pre-logits). Based on δ∞ ≈ - 1 2K2 (∑Kj=1 zs j )2 , we can reformulate Eq. (8). Let rs ∈ Rd be the penultimate representation of student s from an instance x, and W s ∈ RK×d be the weight matrix of the student's fully connected layer. Then,

As derived in Eq. (10), training a network with LKL encourages the pre-logits to be dilated via δ∞ (Figure 4b). For visualization, following [M¨uller et al., 2019], we first find an orthonormal basis constructed from the templates (i.e., the mean of the representations of the samples within the same class) of the three selected classes (apple, aquarium fish, and baby in our experiments). Then, the penultimate layer representations are projected onto the hyperplane based on the identified orthonormal basis. WRN-28-4 (t) is used as a teacher, and WRN-16-2 (s) is used as a student on the CIFAR-100 training dataset. As shown in the first row of Figure 5, when WRN-like architectures are trained with LCE based on ground-truth hard targets, clusters are tightened as the model's complexity increases. As shown in the second row of Figure 5, when the student s is trained with LKL with infinite τ or with LMSE, both representations attempt to follow the shape of the teacher's representations but differ in the degree of cohesion. This is because δ∞ makes the pre-logits become much more widely clustered. Therefore, LMSE can shrink the representations more than LKL along with the teacher.

4.4 Effects of a Noisy Teacher

We investigate the effects of a noisy teacher (i.e., a model poorly fitted to the training dataset) according to the objective. It is believed that the label matching (LKL with a small τ) is more appropriate than the logit matching (LKL with a large τ or the LMSE) under a noisy teacher. This is because label matching neglects the negative information of the outputs of an untrained teacher. Table 4 describes top-1 test accuracies on CIFAR-100, where the used teacher network (WRN-28-4) has a training accuracy of 53.77%, which is achieved in 10 epochs. When poor knowledge is distilled, the students following the label matching scheme performed better than the students following the logit matching scheme, and the extreme logit matching through LMSE has the worst performance. Similarly, it seems that logit matching is not suitable for large-scale tasks. Table 5 presents top-1 test accuracies on ImageNet, where the used teacher network (ResNet152) has a training accuracy of 81.16%, which is provided in PyTorch. Even in this case, the extreme logit matching exhibits the worst performance. The utility of negative logits (i.e., negligible aspect when τ is small) was discussed in [Hinton et al., 2015].

5 Sequential Distillation

In [Cho and Hariharan, 2019], the authors showed that more extensive teachers do not mean better teachers, insisting that the capacity gap between the teacher and the student is a more important factor than the teacher itself. In their results, using a medium-sized network instead of a large-scale network as a teacher can improve the performance of a small network by reducing the capacity gap between the teacher and the student. They also showed that sequential KD (large network →

medium network → small network) is not conducive to generalization when (α, τ) = (0.1, 4) in Eq. (2). In other words, the best approach is a direct distillation from the medium model to the small model.

Table 6 describes the test accuracies of sequential KD, where the largest model is WRN-28-4, the intermediate model is WRN-16-4, and the smallest model is WRN-16-2. Similar to the previous study, when LKL with τ = 3 is used to train the small network iteratively, the direct distillation from the intermediate network to the small network is better (i.e., WRN-16-4 → WRN-16-2, 74.84%) than the sequential distillation (i.e., WRN-28-4 → WRN-16-4 → WRN-162, 74.52%) and direct distillation from a large network to a small network (i.e., WRN-28-4 →WRN-16-2, 74.24%). The same trend occurs in LMSE iterations.

On the other hand, we find that the medium-sized teacher can improve the performance of a smaller-scale student when LKL and LMSE are used sequentially (the last fourth row) despite the large capacity gap between the teacher and the student. KD iterations with such a strategy might compress the model size more effectively, and hence should also be considered in future work. Furthermore, our work is the first study on the sequential distillation at the objective level, not at the architecture level such as [Cho and Hariharan, 2019; Mirzadeh et al., 2020].

6 Robustness to Noisy Labels

In this section, we investigate how noisy labels, samples annotated with incorrect labels in the training dataset, affect the distillation ability when training a teacher network. This setting is related to the capacity for memorization and generalization. Modern deep neural networks even attempt to memorize samples perfectly [Zhang et al., 2016]; hence, the teacher might transfer corrupted knowledge to the student in this situation. Therefore, it is thought that logit matching might not be the best strategy when the teacher is trained using a noisy label dataset.

From this insight, we simulate the noisy label setting to evaluate the robustness on CIFAR-100 by randomly flipping a certain fraction of the labels in the training dataset following a symmetric uniform distribution. Figure 6 shows the test accuracy graphs as the loss function changes. First, we observe that a small network (WRN-16-2 (s), orange dotted line) has a better generalization performance than an extensive network (WRN-28-4 (t), purple dotted line) when models are trained with LCE. This implies that a complex model can memorize the training dataset better than a simple model, but cannot generalize to the test dataset. Next, WRN-28-4 (purple dotted line) is used as the teacher model. When the noise is less than 50%, extreme logit matching (LMSE, green dotted line) and logit matching with δ∞ (LKL(τ = ∞), blue dotted line) can mitigate the label noise problem compared with the model trained with LCE. However, when the noise is more than 50%, these training cannot mitigate this problem because it follows corrupted knowledge more often than correct knowledge.

Interestingly, the best generalization performance is achieved when we use LKL with τ ≤ 1.0. In Figure 6, the blue solid line represents the test accuracy using the rescaled loss function from the black dotted line when τ ≤ 1.0. As expected, logit matching might transfer the teacher's overconfidence, even for incorrect predictions. However, the proper objective derived from both logit matching and label matching enables similar effects of label smoothing, as studied in [Lukasik et al., 2020; Yuan et al., 2020]. Therefore, LKL with τ = 0.5 appears to significantly mitigate the problem of noisy labels.

7 Conclusion

In this paper, we first showed the characteristics of a student trained with LKL according to the temperature-scaling hyperparameter τ. As τ goes to 0, the trained student has the label matching property. In contrast, as τ goes to ∞, the trained student has the logit matching property. Nevertheless, LKL with a sufficiently large τ cannot achieve complete logit matching owing to δ∞. To achieve this goal, we proposed a direct logit learning framework using LMSE and improved the performance based on this loss function. In addition, we showed that the model trained with LMSE followed the teacher's penultimate layer representations more than that with LKL. We observed that sequential distillation can be a better strategy when the capacity gap between the teacher and the student is large. Furthermore, we empirically observed that, in the noisy label setting, using LKL with τ near 1 mitigates the performance degradation rather than extreme logit matching, such as LKL with τ = ∞ or LMSE.

Acknowledgments

This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) [No.20190-00075, Artificial Intelligence Graduate School Program (KAIST)] and [No. 2021-0-00907, Development of Adaptive and Lightweight Edge-Collaborative Analysis Technology for Enabling Proactively Immediate Response and Rapid Learning].