PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Taehyeon Kim, Jaehoon Oh, NakYil Kim, Sangwook Cho, Se-Young Yun

arXiv:2105.08919v1cs.LGcs.CV

TL;DR

The paper addresses limited understanding of how temperature-dependent softening affects knowledge-distillation generalization. It analyzes KL distillation theoretically and empirically, proposes direct MSE logit matching, and finds improved performance overall, with sequential distillation and low-temperature KL useful in specific settings.

  • Problem

    Prior work provides limited understanding of how the degree of softening in KL-based knowledge distillation affects generalization.

  • Method

    The paper analyzes KL distillation across temperatures and replaces softened-probability matching with MSE between teacher and student logits.

  • Results

    MSE direct logit matching outperforms KL in the experiments, while sequential distillation helps with large capacity gaps and low-temperature KL mitigates noisy-label degradation.

  • Takeaways & Limitations

    Logit matching generally has better generalization, but KL with low temperature is preferable to extreme logit matching when labels are noisy.

  • Takeaways & Limitations

    The zero-mean-logit assumption used for large-temperature analysis is inappropriate because it ignores a hidden KL term.

Abstract

from arXiv · show

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 tau. Despite its widespread use, few studies have discussed the influence of such softening on generalization. Here, we theoretically show that the KL divergence loss focuses on the logit matching when tau increases and the label matching when tau 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 difference in the penultimate layer representations between the two losses. Furthermore, we show that sequential distillation can improve performance and that KD, particularly when using the KL divergence loss with small tau, mitigates the label noise. The code to reproduce the experiments is publicly available online at https://github.com/jhoon-oh/kd_data/.

1 Introduction

The paper examines how temperature controls what KL-based knowledge distillation transfers and proposes direct logit matching with MSE. It reports stronger general performance from MSE, while KL with low temperature is more effective under noisy labels.

  • Motivation: Knowledge distillation transfers knowledge from a cumbersome teacher to a smaller student using CE and teacher-prediction losses.The framework commonly combines hard ground-truth targets with KL divergence on teacher predictions.
  • Temperature and matching: A large temperature promotes logit matching, whereas a small temperature promotes label matching; logit matching generally has better generalization capacity.The paper investigates this relationship theoretically and empirically.
  • Direct logit matching: The proposed MSE loss directly matches student and teacher logits, and direct training achieves the best performance in the experiments.The authors argue that KL divergence cannot achieve complete logit matching for any temperature.
  • Representations: KL divergence produces more elongated penultimate-layer representations than the teacher, whereas MSE follows the teacher’s representations more closely.The representation difference is offered as an explanation for the performance gap between the losses.
  • Sequential distillation: Sequential distillation, applying MSE after KL, can outperform direct distillation when the capacity gap between teacher and student is large.This strategy is presented as a better option under a substantial capacity mismatch.
  • Noisy labels: KL with low temperature is more efficient than MSE on incorrectly labeled data because extreme logit matching can provoke bad training.In noisy-label settings, KL mitigates the problem that affects extreme logit matching.

2 Related Work

Related work explains knowledge distillation through hidden-feature transfer, regularization, class-information transfer, and model compression. It also frames KD alongside label smoothing while noting unresolved questions about soft-target softness.

  • Knowledge transfer: Knowledge distillation has been extended to transfer hidden feature vectors in addition to softened teacher probabilities.These approaches aim to provide the student with richer information from the teacher.
  • Applications: Distillation is used for model compression and can reduce teacher generalization errors in self-distillation and compress latent features in generative models.The applications span discriminative and generative settings.
  • Explanations of KD: Prior studies attribute KD’s efficacy to teacher softmax importance, label smoothing, or sample reweighting rather than solely to class-similarity transfer.These explanations are supported by findings involving permuted non-argmax outputs and weak teachers.
  • Label smoothing: Label smoothing replaces a one-hot ground-truth vector with a weighted mixture of hard targets to facilitate generalization and prevent overconfident predictions.The related work connects this regularization mechanism to interpretations of KD.
  • Open question: Teacher label smoothing can improve the teacher while hurting the student, and KD has been characterized as label regularization using adaptive noise.These findings motivate examining how soft-target softness affects performance.

3 Preliminaries: KD

The preliminaries define temperature-softened network outputs and the standard CE–KL distillation objective. They also motivate examining large-temperature behavior because the usual zero-mean-logit approximation can overlook a relevant hidden term.

  • Softened probabilities: For network f and sample x, p_f(τ) denotes the softened probability vector produced using temperature τ.Its k-th value is computed from the network’s logits, with K denoting the number of classes.
  • KD objective: The standard student objective linearly combines cross-entropy on the one-hot label with KL divergence between student and teacher softened probabilities.The balance uses α, with standard choices α = 0.1 and τ ∈ {3, 4, 5}.
  • Large-temperature limit: As τ approaches infinity, minimizing KL is approximated by minimizing MSE between student and teacher logits under the zero-mean-logit assumption.The derivation identifies KL’s large-temperature behavior with logit matching.
  • Assumption: The zero-mean-logit assumption is questioned because it ignores a hidden term in KL when temperature increases.The authors report that the assumption hinders complete understanding of large-temperature behavior.
  • Empirical diagnostic: Figure 2 compares logit-summation magnitudes for a WRN-28-4 teacher and WRN-16-2 student on CIFAR-100 training data.The teacher’s logit sum is almost zero, while the KL-trained student’s sum moves far from zero as temperature increases.

4 Relationship between LKL and LMSE

The analysis connects KL-divergence distillation to logit matching as τ increases, while identifying an additional regularization term that prevents complete matching. Direct MSE logit distillation avoids this hindrance and improves student performance across teacher–student settings.

  • 4.1 Hyperparameter τ in LKL: As τ increases, LKL increasingly imitates the teacher’s logit distribution; as τ approaches zero, it primarily matches labels.These trends are consistent across teacher–student pairs and are analyzed theoretically and empirically.
  • 4.2 Extensions from LKL to LMSE: For sufficiently large τ, LKL is equivalent to scaled LMSE plus δ∞, rather than complete logit matching.The relationship is limτ→∞LKL = 1/(2K) LMSE + δ∞.
  • 4.2 Extensions from LKL to LMSE: δ∞ shifts the student logit mean and hinders complete matching by encouraging deviation from the teacher’s logit mean.The student’s logit mean departs from zero as τ increases under this term.
  • 4.2 Extensions from LKL to LMSE: The proposed LMSE objective directly matches student and teacher logits, and direct training achieves the best performance across various teacher–student pairs.LMSE also achieves 75.60% test accuracy for WRN16-2 with ensemble teachers WRN16-4, WRN-28-4, and WRN-40-6.
  • 4.3 Comparison of LKL and LMSE: LMSE transfers teacher information more efficiently than LKL by aligning logits without the large-τ mean-shifting effect.In pre-logit visualizations, both methods follow the teacher’s representation shape, but LMSE produces more cohesive representations.
  • 4.3 Comparison of LKL and LMSE: The study excludes MSE replacements in feature-based distillation because existing methods contain hyperparameters sensitive to optimization.This limits the reported extension of MSE to those feature-based settings.

5 Sequential Distillation

Sequential distillation can help when a large teacher–student capacity gap exists, although direct intermediate-to-small distillation performs better under one KL setting.

  • 5 Sequential Distillation: Earlier work found that a medium-sized teacher could outperform a larger teacher by reducing the capacity gap to the student.That work also reported direct medium-to-small distillation as preferable under its tested KL configuration.
  • 5 Sequential Distillation: 74.84% was achieved by direct WRN-16-4 → WRN-16-2 distillation with LKL and τ = 3, versus 74.52% for sequential WRN-28-4 → WRN-16-4 → WRN-16-2 distillation.These are test accuracies reported for the sequential-KD comparison.
  • 5 Sequential Distillation: When the capacity gap is large, sequential distillation using KL followed by MSE can outperform direct distillation.The paper presents this as contrasting earlier findings that sequential distillation was not conducive to generalization.
  • 5 Sequential Distillation: The authors report that sequential KL and MSE distillation can improve a smaller student despite a large capacity gap.They suggest this strategy may compress model size more effectively and identify it as objective-level sequential distillation.

6 Robustness to Noisy Labels

The paper evaluates distillation under randomly flipped training labels and finds that moderate-temperature KL distillation is most effective against label noise, while extreme logit matching fails beyond 50% noise.

  • 6 Robustness to Noisy Labels: When noise is below 50%, LMSE and LKL with τ = ∞ can mitigate label-noise effects.These methods are described as extreme logit-matching approaches in the experiment.
  • 6 Robustness to Noisy Labels: LKL with τ ≤1.0 achieves the best generalization performance in the noisy-label experiments.The paper specifically notes that τ = 0.5 appears to significantly mitigate noisy-label problems.

7 Conclusion

The paper characterizes KL distillation across temperature limits, proposes direct logit learning with MSE, and reports benefits for representation matching, sequential distillation, and noisy-label robustness.

  • 7 Conclusion: As τ approaches 0, KL training has label-matching behavior; as τ approaches infinity, it has logit-matching behavior but cannot achieve complete matching because of δ∞.The paper attributes δ∞ to an additional effect that hinders complete logit matching.
  • 7 Conclusion: Direct logit learning with LMSE improves performance relative to KL-based training in the paper’s experiments.The conclusion links this result to the proposed MSE loss framework.
  • 7 Conclusion: LMSE-trained models follow the teacher’s penultimate-layer representations more closely than models trained with LKL.This representation difference is reported as part of the comparison between the two losses.
  • 7 Conclusion: Sequential distillation can be better when the capacity gap between teacher and student is large.The conclusion presents this as an empirical observation of the study.
  • 7 Conclusion: With noisy labels, LKL at a temperature near 1 mitigates performance degradation better than extreme logit matching such as τ = ∞ or LMSE.This conclusion contrasts moderate-temperature KL with the extreme logit-matching objectives.

A.3 Proof of Equation 8

The proof of Equation 8 uses bounded convergence to establish the large-temperature relationship between KL and MSE objectives, including the role of logit-mean assumptions and δ∞.

  • A.3 Proof of Equation 8: The proof uses the bounded convergence theorem to interchange a limit and an integral.It suffices to show that limτ→∞|g_k(τ)| is bounded.
  • A.3 Proof of Equation 8: For sufficiently large τ, minimizing LKL is equivalent to minimizing LMSE under a zero-meaned-logit assumption.The proof discusses this relationship through the limiting form of LKL.
  • A.3 Proof of Equation 8: Without the zero-meaned-logit assumption, K acts as a relative weight between the two terms.The derivation distinguishes this case from the zero-mean formulation.
  • A.3 Proof of Equation 8: The appendix points to detailed training and testing accuracy values in Tables 7 and 8.Those tables concern CIFAR-100 with WRN-28-4 as teacher and WRN-16-2 as student.

C Appendix: Other methods in Table 3

The appendix compares several feature- and activation-based distillation methods and reports results across varied teacher–student pairs. Figure 7 maps training and test top-1 accuracy as α and τ change on CIFAR-100.

  • Other methods: The comparison includes FitNets, Attention Transfer, Jacobian matching, Factor Transfer, Activation Boundary, and Overhaul.These methods represent alternative feature, activation, Jacobian, and network-compression approaches to distillation.
  • Experimental setting: Results supporting Figure 3 are provided for various teacher–student pairs.The appendix extends the evaluation beyond a single architectural pairing.
  • Experimental setting: Figure 7 maps training and test top-1 accuracies across changes in α and τ for four teacher–student combinations on CIFAR-100.The student is WRN-16-2 in all four combinations, while the teachers are WRN-16-4, WRN-16-6, WRN-28-2, and WRN-40-2.

E Calibration

Calibration differs between KL-divergence and MSE distillation. On CIFAR-100, MSE produces a reliability slope closer to the teacher and a lower ECE than KL divergence with infinite temperature.

  • Calibration comparison: 10.46% versus 14.38%: MSE yields lower ECE than KL divergence with infinite τ for the student model.The comparison uses WRN-16-2 students distilled from WRN-28-4 teachers on CIFAR-100.
  • Calibration comparison: The MSE-trained student’s reliability-diagram slope is closer to the teacher’s calibration than the KL-trained student’s slope.The KL comparison uses L_KL with infinite τ.
  • Cross-entropy reference: 8.44% versus 10.33%: the teacher has slightly lower ECE than the student when both are trained with cross-entropy.The reliability diagrams use 10 bins on the CIFAR-100 training dataset.
Loading 2105.08919v1…