arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2411.08937cs.CVcs.LG

双头知识蒸馏:通过辅助头增强logits的利用

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

Nanyang Technological University(南洋理工大学);Nanjing University of Aeronautics and Astronautics(南京航空航天大学);Southeast University(东南大学)

Penghui Yang, Chen-Chen Zong, Sheng-Jun Huang, Lei Feng, Bo An

AI总结:

本文提出双头知识蒸馏方法,通过分离分类头和骨干网络,解决传统知识蒸馏中logit级损失与概率级损失结合导致的分类头崩溃问题,提升模型性能。

摘要:

传统知识蒸馏侧重于对齐学生模型的预测概率与真实标签及教师模型的预测概率。然而,从logits到预测概率的转换会掩盖某些关键信息。为此,引入logit级损失函数作为补充,以利用logits的潜在信息。但实验发现,将logit级损失与概率级损失结合会导致性能下降。我们归因于分类头的崩溃,并通过神经崩溃理论进行理论分析。具体而言,两种损失函数的梯度在线性分类器中存在矛盾,但在骨干网络中无冲突。基于此,我们提出双头知识蒸馏方法,将线性分类器分为两个负责不同损失的头,从而保留两种损失对骨干网络的有益影响,同时消除对分类头的不良影响。大量实验验证了该方法能有效利用logits中的信息,优于现有方法。代码见:https://github.com/penghui-yang/DHKD。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑