Swapped Logit Distillation via Bi-level Teacher Alignment

TL;DR

The paper addresses a core challenge in knowledge distillation: transferring knowledge from a large teacher to a smaller student without propagating incorrect predictions from the teacher's logit distribution. It introduces Swapped Logit Distillation (SLD), which swaps the ground-truth target with the highest non-target in logits and employs a pseudo-teacher via a second swap, combined with loss scheduling to avoid early training conflicts. SLD uses prediction augmentation across multiple temperatures and two KL-based losses, $L_{TS}$ and $L_{SS}$, to realize a final objective $L_{SLD} = L_{TS} + L_{SS}$ that aligns both teacher and pseudo-teacher with the student. Extensive experiments on CIFAR-100 and ImageNet show that SLD consistently outperforms state-of-the-art KD methods, can complement feature-based distillation, and operates efficiently without adding extra parameters.

Abstract

Knowledge distillation (KD) compresses the network capacity by transferring knowledge from a large (teacher) network to a smaller one (student). It has been mainstream that the teacher directly transfers knowledge to the student with its original distribution, which can possibly lead to incorrect predictions. In this article, we propose a logit-based distillation via swapped logit processing, namely Swapped Logit Distillation (SLD). SLD is proposed under two assumptions: (1) the wrong prediction occurs when the prediction label confidence is not the maximum; (2) the "natural" limit of probability remains uncertain as the best value addition to the target cannot be determined. To address these issues, we propose a swapped logit processing scheme. Through this approach, we find that the swap method can be effectively extended to teacher and student outputs, transforming into two teachers. We further introduce loss scheduling to boost the performance of two teachers' alignment. Extensive experiments on image classification tasks demonstrate that SLD consistently performs best among previous state-of-the-art methods.

Figures (6)

• Figure 1: Illustration of intuitive logit processing schemes with a false prediction case.
• Figure 2: Overview of SLD framework. After obtaining the logit outputs, swapping methods are applied to teacher logits, producing main teacher and pseudo-teacher logits. Loss scheduling is proposed to avoid teacher conflicts when both teachers are integrated. Prediction augmentation is used to generate outputs at various softness levels. Real misclassified samples of the teacher's prediction illustrate how a model prediction can mislead the student model's learning since the target and the prediction share similarities in color, shape, and textures (where the red text is the target, and the black one is the false prediction).
• Figure 3: Effects of scheduling on ResNet50 as a teacher and MobileNet-V2 as a student on CIFAR-100.
• Figure 4: Training time (per batch). We set ResNet32$\times$4 as the teacher model and ResNet8$\times$4 as the student model on CIFAR-100. The table shows the accuracy and the number of additional parameters.
• Figure 5: Prediction distribution of KD (left) and KD with swap method (right). Class 4 (beaver) is the ground-truth, class 55 (otter), and class 63 (porcupine).