Logit Standardization in Knowledge Distillation

TL;DR

The proposed Z-score 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.

Abstract

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.

Figures (9)

• Figure 1: Vanilla knowledge distillation implicitly enforces an exact match between the magnitudes of teacher and student logits. It is an unnecessary side-effect because it is found sufficient to preserve the innate relations between their logits. Given the capacity gap between them, it is also challenging for a lightweight student to produce logits with the same magnitude as a cumbersome teacher. In contrast, the proposed $\mathcal{Z}$-score logit standardization pre-process mitigates the side-effect. The standardized student logits have arbitrary magnitude suitable for the student's capacity while preserving the essential relations learned from the teacher.
• Figure 2: A toy case where two students, $\mathcal{S}_1$ and $\mathcal{S}_2$, learning from the same teacher with an identical temperature (assumed 1 for simplicity). Student $\mathcal{S}_1$ generates the logits much closer to the teacher's in terms of magnitude and thus has lower loss of 0.1749, but it returns a wrong prediction of "bird". In contrast, Student $\mathcal{S}_2$ outputs the logits far from the teacher's and yields greater loss value of 0.3457, but it returns the correct prediction of "dog". After the proposed logit standardization, the issue is addressed.
• Figure 3: 1st Row: An example bar plot of logit output. 2nd Row: The heatmap of the average logit difference between the teacher and student. Our pre-process indeed enables the student to generate the logits of divergent range from the teacher as shown in \ref{['fig::diff_wo_stand']}, while its standardized logits (\ref{['fig::diff_w_stand']}) are more closer to the teacher's than vanilla KD (\ref{['fig::diff_kd']}).
• Figure 4: Teacher
• Figure 5: KD hinton2015distilling