Logit Standardization in Knowledge Distillation

TL;DR

This paper introduces a logit standardization method using Z-score pre-processing in knowledge distillation, allowing better transfer of essential logit relations and improving performance across various models and datasets.

Contribution

It proposes a plug-and-play Z-score standardization of logits, addressing the limitations of shared temperature assumptions in knowledge distillation.

Findings

Improves student model performance on CIFAR-100 and ImageNet.

Enhances existing distillation methods with the proposed pre-process.

Achieves state-of-the-art results with vanilla distillation using logit standardization.

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.