Enhancing Logits Distillation with Plug\&Play Kendall's $\tau$ Ranking Loss

TL;DR

This paper introduces a plug-and-play Kendall's tau ranking loss to improve logits distillation by balancing gradient updates, leading to better transfer of inter-class information and enhanced student model performance.

Contribution

The authors propose a novel auxiliary ranking loss based on Kendall's tau that complements existing distillation methods and addresses gradient imbalance issues.

Findings

Consistently improves performance across datasets and architectures.

Balances gradients towards low-probability channels effectively.

Enhances inter-class relational information transfer.

Abstract

Knowledge distillation typically minimizes the Kullback-Leibler (KL) divergence between teacher and student logits. However, optimizing the KL divergence can be challenging for the student and often leads to sub-optimal solutions. We further show that gradients induced by KL divergence scale with the magnitude of the teacher logits, thereby diminishing updates on low-probability channels. This imbalance weakens the transfer of inter-class information and in turn limits the performance improvements achievable by the student. To mitigate this issue, we propose a plug-and-play auxiliary ranking loss based on Kendall's $\tau$ coefficient that can be seamlessly integrated into any logit-based distillation framework. It supplies inter-class relational information while rebalancing gradients toward low-probability channels. We demonstrate that the proposed ranking loss is largely invariant to channel scaling and optimizes an objective aligned with that of KL divergence, making it a natural complement rather than a replacement. Extensive experiments on CIFAR-100, ImageNet, and COCO datasets, as well as various CNN and ViT teacher-student architecture combinations, demonstrate that our plug-and-play ranking loss consistently boosts the performance of multiple distillation baselines. Code is available at https://github.com/OvernighTea/RankingLoss-KD