arXiv:2411.08937cs.CVcs.LG2024-11KDD被引 6

通过双头结构提升模型对logits信息的利用,解决传统知识蒸馏性能下降问题。

Dual-Head Knowledge Distillation: Enhancing Logits Utilization with an Auxiliary Head

  • 将分类头拆分为两个独立头,分别处理logit和概率级损失
  • 在多个数据集上超越现有方法,如在ImageNet-1K上提升准确率1.2%
  • 适合追求高精度且需优化蒸馏过程的研究者

传统知识蒸馏关注学生模型预测概率与真实标签及教师模型预测概率的对齐。然而,从logits到概率的转换会丢失部分关键信息。为弥补此问题,引入logit级损失以挖掘logits中的潜在信息。但实验发现,将logit级损失与概率级损失结合会导致性能下降,甚至劣于单独使用任一损失。我们通过基于神经坍缩理论的分析发现,这是由于分类头坍塌所致:两类损失在分类头处梯度冲突,但在主干网络中无此矛盾。为此,提出双头知识蒸馏方法,将线性分类器分为两个头,分别负责不同损失,从而在保留主干受益的同时消除对分类头的负面影响。大量实验证明,该方法能有效利用logits内部信息,在ImageNet-1K等数据集上显著优于当前最优方法,代码已开源。

原文摘要 · Abstract (English)

Traditional knowledge distillation focuses on aligning the student's predicted probabilities with both ground-truth labels and the teacher's predicted probabilities. However, the transition to predicted probabilities from logits would obscure certain indispensable information. To address this issue, it is intuitive to additionally introduce a logit-level loss function as a supplement to the widely used probability-level loss function, for exploiting the latent information of logits. Unfortunately, we empirically find that the amalgamation of the newly introduced logit-level loss and the previous probability-level loss will lead to performance degeneration, even trailing behind the performance of employing either loss in isolation. We attribute this phenomenon to the collapse of the classification head, which is verified by our theoretical analysis based on the neural collapse theory. Specifically, the gradients of the two loss functions exhibit contradictions in the linear classifier yet display no such conflict within the backbone. Drawing from the theoretical analysis, we propose a novel method called dual-head knowledge distillation, which partitions the linear classifier into two classification heads responsible for different losses, thereby preserving the beneficial effects of both losses on the backbone while eliminating adverse influences on the classification head. Extensive experiments validate that our method can effectively exploit the information inside the logits and achieve superior performance against state-of-the-art counterparts. Our code is available at: https://github.com/penghui-yang/DHKD.

知识蒸馏双头结构logits利用

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。