Source-linked AI summary
Logit Standardization in Knowledge Distillation
Shangquan Sun, Wenqi Ren, Jingzhi Li, Rui Wang, Xiaochun Cao
TL;DR
Conventional knowledge distillation shares a temperature that can force teacher–student logit-magnitude matching despite their capacity difference. This paper derives flexible temperature assignments and adds weighted Z-score logit standardization, which consistently improves existing logit-based methods across CIFAR-100 and ImageNet.
Problem
Shared temperatures in conventional logit-based distillation implicitly enforce teacher–student logit matching even though preserving innate logit relations can suffice and model capacities differ.
Method
The paper derives temperature-based softmax from entropy maximization and applies weighted Z-score logit standardization as a plug-and-play pre-process before distillation.
Results
The pre-process consistently improves four existing logit-based KD methods on CIFAR-100 and improves three logit-based methods on ImageNet.
Takeaways & Limitations
Z-score standardization lets students focus on essential teacher-logit relations rather than matching teacher logit magnitude and can benefit existing distillation variants.
Takeaways & Limitations
The theoretical analysis assumes a finally well-distilled student whose predicted probability distribution matches the teacher’s.
Abstract
from arXiv · showhide
Knowledge distillation involves transferring soft labels from a teacher to a student using a shared temperature-based softmax function. However, the assumption of a shared temperature between teacher and student implies a mandatory exact match between their logits in terms of logit range and variance. This side-effect limits the performance of student, considering the capacity discrepancy between them and the finding that the innate logit relations of teacher are sufficient for student to learn. To address this issue, we propose setting the temperature as the weighted standard deviation of logit and performing a plug-and-play Z-score pre-process of logit standardization before applying softmax and Kullback-Leibler divergence. Our pre-process enables student to focus on essential logit relations from teacher rather than requiring a magnitude match, and can improve the performance of existing logit-based distillation methods. We also show a typical case where the conventional setting of sharing temperature between teacher and student cannot reliably yield the authentic distillation evaluation; nonetheless, this challenge is successfully alleviated by our Z-score. We extensively evaluate our method for various student and teacher models on CIFAR-100 and ImageNet, showing its significant superiority. The vanilla knowledge distillation powered by our pre-process can achieve favorable performance against state-of-the-art methods, and other distillation variants can obtain considerable gain with the assistance of our pre-process.
1. Introduction
Knowledge distillation transfers knowledge from a heavy teacher to a lightweight student, but shared temperatures can force unnecessary logit-magnitude matching. The paper derives flexible temperature assignments and proposes Z-score standardization to preserve logit relations while improving existing methods.
- Knowledge distillation transfers a pre-trained heavy teacher’s knowledge to a small student through soft-label prediction matching.
- Shared temperatures implicitly require teacher and student logits to match in magnitude, range, and variance despite their capacity gap.
- The paper derives softmax in logit-based distillation from entropy maximization, showing temperatures arise from flexible Lagrangian multipliers.
- Z-score logit standardization adaptively allocates temperatures between models and across samples as a plug-and-play pre-process.
- Extensive CIFAR-100 and ImageNet experiments show consistent advantages across existing logit-based distillation methods.
2. Related Work
Related work frames knowledge distillation as transferring teacher soft labels through probability matching, while temperature controls how strongly the student attends to lower-valued logits. Existing approaches generally leave the theoretical basis for globally shared temperatures unexamined.
- Knowledge distillation transfers teacher “dark” knowledge by minimizing a divergence between student and teacher predictions.
- Temperature flattens logit-based probabilities and controls how much attention the student pays to logits below the average.
- Very low temperature makes the student focus mainly on the teacher’s largest logit while ignoring other logits.
- Prior work does not explain why teacher and student should share one globally predefined temperature.
3. Background and Notation
The background defines teacher and student logits, temperature-based softmax probabilities, and KL-divergence distillation over samples and classes. It also notes that KL divergence and cross-entropy are theoretically equivalent under a restricted optimization view but can have different empirical gradients.
- For each image-label sample, teacher and student networks output logit vectors over K classes.
- Temperature-based softmax converts teacher and student logits into class-probability vectors.
- Knowledge distillation trains the student to mimic the teacher’s probability for every class and sample using KL divergence.
- KL divergence is theoretically equivalent to cross-entropy when optimizing solely over student logits.
- The two losses are empirically nonequivalent because their gradients differ through the teacher-probability entropy term.
4. Methodology
The methodology derives temperature-flexible softmax formulations, analyzes the logit shift and variance constraints imposed by shared temperatures, and introduces weighted Z-score standardization. The pre-process is applied before existing logit-based distillation losses and evaluated as a plug-and-play component.
- 4.1. Derivation of softmax: Entropy-maximization derivations show teacher and student temperatures may be distinct and may vary across samples.
- 4.2. Drawbacks of Shared Temperatures: Under shared temperatures, a well-distilled student is forced to reproduce the teacher logits up to a sample-specific constant shift.
- 4.2. Drawbacks of Shared Temperatures: Shared temperatures also force student and teacher logit standard deviations to be equal.
- 4.3. Logit Standardization: The proposed weighted Z-score function standardizes logits with a base temperature before softmax and KL-divergence distillation.
- 4.3. Logit Standardization: The pre-process is inserted into existing logit-based KD pipelines and trains the student using the resulting distillation loss.
7 end
The proposed logit standardization uses logit statistics to remove magnitude and variance restrictions while preserving logit relations before distillation. In a toy case, it changes the loss comparison to agree with the students’ predictions.
- Logit Standardization: The method sets logit-affecting hyperparameters from each vector’s mean and weighted standard deviation, with a shared base temperature.This adaptively allocates temperatures between teacher and student while retaining a common base temperature.
- Logit Standardization: Z-score standardization gives transformed logits zero mean, standard deviation 1/τ, preserved ranking, and a bounded range.These properties remove restrictions on the original student logits’ variance and value range while preserving their innate relations.
- Toy Case: In the toy case, shared-temperature KL loss favors S1 at 0.1749 over S2 at 0.3457 despite S1 predicting “bird” incorrectly and S2 predicting “dog” correctly.S1 matches the teacher’s logit magnitudes more closely, whereas S2 preserves the teacher’s logit relations.
- Toy Case: After Z-score standardization, S2 receives loss 0 while S1 receives 0.0995, aligning the loss comparison with their observed predictions.The standardized evaluation emphasizes relations among logits rather than their magnitudes.
5. Experiments
Experiments across CIFAR-100 and ImageNet show that Z-score logit standardization consistently improves existing logit-based distillation methods. The method also relaxes student logit-range and variance matching while preserving standardized alignment with the teacher.
- Main Results: Z-score standardization consistently improves four existing logit-based KD methods on CIFAR-100, including KD, CTKD, DKD, and MLKD.Vanilla KD reaches performance comparable to state-of-the-art feature-based methods, while DKD, MLKD, and CTKD also benefit.
- Main Results: Z-score preprocessing yields consistent improvements for all three evaluated logit-based methods on ImageNet in both top-1 and top-5 accuracy.The comparison uses ImageNet validation-set top-1 and top-5 accuracy.
- Ablation Studies: As the KD-loss weight increases, vanilla KD gains little, whereas Z-score preprocessing produces significant enhancement.The ablation fixes the base temperature at 2 and examines different KD-loss weights.
- Extensions: Z-score lets students use divergent raw logit ranges while reducing standardized teacher–student mean distance from 0.27 to 0.18.Without preprocessing, the student’s target-label logit is 7.5 versus the teacher’s 12; standardized logits are closer to the teacher.
- Extensions: The method permits flexible raw logit variance, with student and teacher values of 0.48 and 3.10, while standardized variances both equal 0.99.Vanilla KD instead drives the raw student variance toward the teacher’s, measured as 3.78 versus 3.10.
- Extensions: Across teacher sizes and capacities, preprocessing consistently improves distillation and makes standardized student logit mean and standard deviation match the teacher.The paper also reports improved feature separability and discriminability across KD, CTKD, DKD, and MLKD.
6. Conclusion
The analysis finds no theoretical support for globally shared temperatures in conventional KD and proposes Z-score standardization so students learn teacher logit relations without matching magnitude.
- Entropy-maximization analysis derives temperature from a flexible Lagrangian multiplier, allowing distinct assignments rather than globally shared values.
- Shared temperatures can produce inauthentic student evaluations, while Z-score standardization addresses this concern by separating logit relations from magnitude.
- Extensive experiments show that the pre-process enhances existing logit-based knowledge-distillation methods.