CrossKD: Cross-Head Knowledge Distillation for Dense Object Detection

Introduction

论文题目 ：CrossKD: Cross-Head Knowledge Distillation for Dense Object Detection

论文地址 ：https://arxiv.org/pdf/2306.11369.pdf

论文出处 ：In 2023 ’ ICCV (under review)

代码实现 ：https://github.com/jbwang1997/CrossKD

Idea

目标检测知识蒸馏受到越来越多的关注，最近几年发现在目标检测领域中特征蒸馏(Feature Imitation)比逻辑蒸馏(Prediction Mimicking)更加高效一些。

作者发现GT和Distill Target(教师的预测输出)之间存在很大的差异，这一点一直都被之前的工作忽视了，然而作者认为这一点正是阻碍逻辑蒸馏Predicted Mimicking达到更高性能的原因。

本文提出一种简单且有效的蒸馏机制CrossKD，它直接将Student Head的中间特征送入到Teacher Head，所得跨头(Cross-Head)预测将被用于最小化与老师模型预测的之间差异。这样的蒸馏机制缓解了Student Head从GT与Teacher Head预测处接收差异过大的信息，进而极大的改善了学生模型的检测性能。

看代码发现目前只适配了One-Stage的检测器，不知道Two-Stage检测器能不能work。

Detail

Cross-Head Knowledge Distillation

下图就是CrossKD的具体网络结构：

在分类分支和回归分支都采用了CrossKD，并且Cross的方案都是一样的

图上还画了反向传播的路径，Cross Head就是按原路径返回

要注意的是教师这边是已经训练好了的，所有参数都已经frozen固定住了，这也是为什么不进行反向传播

$\mathcal{L}_{\mathrm{CKD}}=\frac{1}{|\mathcal{S}|}\sum_{r\in\mathcal{R}}\mathcal{S}(r)\mathcal{D}_{\mathrm{pred}}(\hat{\boldsymbol{p}}^{s}(r),\boldsymbol{p}^{t}(r))$

$\mathcal{S}$ 这个系数，在分类分支是1，在回归分支，前景部分是1，背景部分是0.

本质是CrossKD是一种逻辑蒸馏(Prediction Mimicking)

Optimization Objectives

训练的总损失就是检测损失和蒸馏损失的加权和：

$\begin{aligned}\mathcal{L}&=\mathcal{L}_{\mathrm{cls}}(\boldsymbol{p}_{\mathrm{cls}}^s,\boldsymbol{p}_{\mathrm{cls}}^{gt})+\mathcal{L}_{\mathrm{reg}}(\boldsymbol{p}_{\mathrm{reg}}^s,\boldsymbol{p}_{\mathrm{reg}}^{gt})\\&+\mathcal{L}_{\mathrm{CKD}}^{\mathrm{cls}}(\boldsymbol{\hat{p}}_{\mathrm{cls}}^s,\boldsymbol{p}_{\mathrm{cls}}^t)+\mathcal{L}_{\mathrm{CKD}}^{\mathrm{reg}}(\boldsymbol{\hat{p}}_{\mathrm{reg}}^s,\boldsymbol{p}_{\mathrm{reg}}^t),\end{aligned}$