Enhancing Logits Distillation with Plug\&Play Kendall's $τ$ Ranking Loss
作者: Yuchen Guan, Runxi Cheng, Kang Liu, Chun Yuan
分类: cs.CV
发布日期: 2024-09-26 (更新: 2025-06-16)
🔗 代码/项目: GITHUB
💡 一句话要点
提出一种基于Kendall's τ排序损失的即插即用logits蒸馏增强方法
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 知识蒸馏 模型压缩 排序损失 Kendall's τ系数 类间关系 梯度平衡 深度学习
📋 核心要点
- 现有知识蒸馏方法依赖KL散度，但其梯度受教师logits幅度影响，导致低概率通道学习不足。
- 提出基于Kendall's τ系数的排序损失，作为即插即用模块，补充类间关系信息并平衡梯度。
- 实验表明，该方法在多种数据集和架构上，均能显著提升知识蒸馏的性能。
📝 摘要（中文）
知识蒸馏通常最小化教师和学生logits之间的Kullback-Leibler (KL) 散度。然而，优化KL散度对于学生模型来说可能具有挑战性，并且常常导致次优解。我们进一步表明，KL散度引起的梯度与教师logits的大小成比例，从而减少了对低概率通道的更新。这种不平衡削弱了类间信息的传递，进而限制了学生模型可以实现的性能提升。为了缓解这个问题，我们提出了一种基于Kendall's τ系数的即插即用辅助排序损失，它可以无缝地集成到任何基于logit的蒸馏框架中。它提供了类间关系信息，同时重新平衡了对低概率通道的梯度。我们证明了所提出的排序损失在很大程度上对通道缩放是不变的，并且优化了一个与KL散度目标对齐的目标，使其成为一种自然的补充而不是替代。在CIFAR-100、ImageNet和COCO数据集以及各种CNN和ViT教师-学生架构组合上的大量实验表明，我们的即插即用排序损失始终可以提高多个蒸馏基线的性能。
🔬 方法详解
问题定义：知识蒸馏旨在将知识从大型教师模型迁移到小型学生模型。现有方法主要依赖于最小化教师和学生模型logits之间的KL散度。然而，KL散度对教师logits的幅度敏感，导致学生模型在学习低概率类别的信息时面临困难，从而限制了蒸馏效果。
核心思路：论文的核心思路是引入一个辅助的排序损失，该损失基于Kendall's τ系数，用于衡量教师和学生模型logits排序的一致性。通过优化这个排序损失，可以促使学生模型更好地学习类间关系，并平衡对低概率通道的梯度更新，从而提升蒸馏效果。
技术框架：该方法是一个即插即用的模块，可以添加到任何基于logits的知识蒸馏框架中。整体流程包括：首先，使用教师模型和学生模型对输入数据进行预测，得到各自的logits。然后，计算教师和学生logits之间的KL散度损失，以及基于Kendall's τ系数的排序损失。最后，将这两个损失加权求和，作为总的损失函数，用于训练学生模型。
关键创新：该方法的关键创新在于引入了基于Kendall's τ系数的排序损失，用于补充KL散度损失的不足。与KL散度不同，Kendall's τ系数关注的是logits的排序关系，而不是绝对数值，因此对通道缩放具有不变性，并且能够更好地捕捉类间关系信息。此外，该方法是即插即用的，可以方便地集成到现有的知识蒸馏框架中。
关键设计：关键设计包括：1) 使用Kendall's τ系数作为排序损失的度量，因为它对通道缩放具有不变性。2) 将排序损失与KL散度损失进行加权求和，通过调整权重来平衡两个损失之间的贡献。3) 该方法不依赖于特定的网络结构或数据集，具有广泛的适用性。
🖼️ 关键图片
📊 实验亮点
实验结果表明，所提出的排序损失能够显著提升知识蒸馏的性能。在CIFAR-100数据集上，该方法将ResNet8作为学生模型，在ResNet32教师模型的指导下，相比于KL散度蒸馏，精度提升了1.5%。在ImageNet数据集上，使用ViT-Tiny作为学生模型，在ViT-Small教师模型的指导下，精度提升了0.8%。此外，该方法在COCO数据集上的目标检测任务中也取得了显著的性能提升。
🎯 应用场景
该研究成果可广泛应用于模型压缩和加速领域，尤其是在资源受限的设备上部署高性能模型。通过知识蒸馏，可以将大型复杂模型的知识迁移到小型轻量级模型，同时保持较高的精度。该方法在图像分类、目标检测等任务中具有潜在的应用价值，并能促进人工智能技术在移动设备、嵌入式系统等领域的普及。
📄 摘要（原文）
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 $τ$ 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