Swapped Logit Distillation via Bi-level Teacher Alignment (2504.20108v1)

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.

Analyzing Swapped Logit Distillation via Bi-level Teacher Alignment

The paper presents an innovative approach to knowledge distillation (KD) by proposing a novel method termed Swapped Logit Distillation (SLD). The purpose of KD is to transfer knowledge from a large teacher model to a smaller student model effectively, optimizing the network size for scenarios requiring reduced computational resources. Traditionally, this knowledge transfer occurs via a straightforward relay of logits from the teacher to the student, potentially propagated with inaccuracies due to mispredicted logit information. Addressing this, SLD introduces a swap-based mechanism to refine this process effectively.

Key Contributions and Methodological Advances

The paper identifies two foundational issues in existing KD frameworks: (1) incorrect predictions due to the maximum confidence not aligning with the correct label and (2) uncertainty in determining the natural limit of probability as an optimal addition to the target. The authors propose addressing this through their SLD technique, structured as follows:

1. Swap Logit Mechanism: SLD intervenes by swapping logit positions between the incorrectly forecasted target and the non-target class holding the highest confidence value. This mechanism ensures that the target receives priority attention during prediction rectification.
2. Bi-level Teacher Setup: A significant methodological innovation involves utilizing a dual teacher structure. The SLD extends teaching responsibilities to both the primary teacher logits and modified student logits acting as a pseudo-teacher. This aligns with a bi-level alignment design, amplifying knowledge gains from diverse logit origins.
3. Loss Scheduling: To resolve conflicts in teacher-student logit learning and synchronization, the authors integrate a loss scheduling approach. Here, the learning from the pseudo-teacher is staggered until a suitable learning juncture, ensuring the fundamental student-teacher alignment precedes more complex logit interactions.

These methodological transformations allow SLD to consistently deliver superior classification results, outperforming existing state-of-the-art KD techniques across several image classification tasks. By addressing intrinsic logit misclassification error sources, the approach advances both the theoretical foundations and practical efficacy of KD.

Robust Experimental Validation

The experimental evaluation deploys SLD on standard datasets such as CIFAR-100 and ImageNet across varied neural network architectures (e.g., ResNets, VGGs). Noteworthy results confirm the superior performance of SLD over competitive methods such as classical KD, Mutual Learning, and various feature-based distillation techniques. The authors highlight that the student models can even outperform their teacher models—a distinct advantage facilitated by the robust pseudo-teacher contributions.

Moreover, ablation studies dissect key components like the loss scheduling mechanics and logit swap strategy, establishing their individual and collective contributions to improving accuracy. Critical improvements observed include reduction in training time, and no additional parameter overhead, indicating SLD's utility in resource-constrained environments.

Theoretical and Practical Implications

Theoretically, SLD enriches the discourse on logit distribution optimization by offering an approach where "dark knowledge" captured in logits undergoes structural rectification before transfer. This refinement is theoretically consistent with ideas advocated by Hinton et al., emphasizing logit-based soft label merits.

Practically, the paper supports SLD's deployment in diverse AI systems, especially those mandating lightweight models due to infrastructure limitations, e.g., mobile or IoT devices. The reduction in computational expense without sacrificing model prediction accuracy could pioneer more widespread KD application across these domains.

Future Directions

The promising direction suggested for future research involves extending SLD's application beyond image classification to more complex tasks, such as object detection and segmentation, where logit misclassifications can lead to more nuanced errors. Additionally, examining the performance scope of SLD within emerging KD paradigms like self-distillation or its potential integration with zero-shot learning frameworks remains an intriguing proposition.

In summary, this paper's presentation of Swapped Logit Distillation marks a substantive contribution to the optimization of knowledge transfer processes in neural networks, proposing robust solutions and empirically proving their merit across diverse machine learning paradigms.