[论文阅读]CWD——基于稠密预测的通道式知识蒸馏

CWD

基于稠密预测的通道式知识蒸馏
论文网址：CWD

论文总结

这篇论文「Channel-wise Knowledge Distillation for Dense Prediction」提出了一种针对密集预测任务的基于通道的知识蒸馏方法。下面我详细讲解该方法的主要内容:

1. 问题背景
   在计算机视觉任务中,图像分类只需要预测整张图像的类别,而密集预测需要对每个像素或对象进行预测,输出更丰富的结果,如语义分割、目标检测等。直接应用分类任务中的知识蒸馏方法于密集预测任务效果不佳。已有的方法通过建模空间位置之间（指的是图像中的像素位置）的关系来传递结构化知识。
2. 方法概述
   论文提出了一种通道级的知识蒸馏方法。主要分为两个步骤:
   (1) 对特征图的每个通道进行softmax标准化,得到一个概率分布（表示了该通道中每个位置的相对重要性或响应强度）。
   (2) 计算教师网络和学生网络相应通道概率分布之间的asymmetric KL散度作为损失,使学生网络在前景显著区域模仿教师网络。
3. 具体实现
   对特征图或logits的每个通道,对H×W个位置的激活值进行softmax计算,得到概率分布表示每个位置的相对重要性。
   然后计算这个分布与教师网络中相应通道分布的asymmetric KL距离,重点对齐前景显著区域。
4. 优点
   相比于空间位置建模的方法,该通道距离有以下优势:

• 更简单有效,提升显著且计算量小
• 自动区分前后景,无需手动设计
• 易于迁移到不同任务与网络结构

5. 实验
   在多个数据集上验证了该方法的有效性,包括Cityscapes语义分割、VOC语义分割和COCO目标检测,均取得SOTA的蒸馏效果。

综上,该论文提出一种简单的通道级知识蒸馏方法,只需计算每个通道的概率距离,就可以有效地对齐前景特征,实现高效的密集预测网络蒸馏。

摘要

知识蒸馏（KD）已被证明是一种简单有效的训练密集预测模型的工具。轻量级学生网络是通过从大型教师网络转移过来的额外监督来训练的。以往大多数针对密集预测任务的KD变体是在空间域对齐来自学生和教师网络的特征图，通常是通过标准化每个空间位置上的激活值，并最小化点和/或成对的差异。与以往的方法不同，本文提出了将每个通道的特征图归一化来得到软概率图。通过简单地最小化两个网络的通道概率图之间的Kullback-Leibler (KL)发散，蒸馏过程更关注每个通道最显著的区域，这对于密集预测任务很有价值。
本文对几个密集预测任务进行了实验，包括语义分割和目标检测。 实验表明，本文提出的方法比现有的蒸馏方法性能好得多，并且在训练过程中需要更少的计算代价。 特别地，在COCO数据集上对MAP中的RetinaNet检测器(resnet50主干）进行了3.4%的改进，在Cityscapes数据集上对mIoU中的PSPNET(RESNet18主干）进行了5.81%的改进。

引言

稠密预测任务是计算机视觉中的一组基本任务，包括语义分割和目标检测。 这些任务要求在像素级学习复杂场景理解的强特征表示。 因此，最先进的模型通常需要很高的计算成本，使得它们难以部署到移动设备上。 因此，为密集预测任务设计的轻量级网络引起了人们的广泛关注。 此外，在以往的研究中，利用知识蒸馏(KD)对轻量级网络进行了有效的训练。 一个轻量级的网络是在一个大型教师网络的监督下训练的，可以实现更好的性能。 开创性的工作[Distilling the knowledge in a neural network,Fitnets: Hints for thin deep nets.]被提出并得到很好的研究，主要用于图像分类任务。
密集预测任务是以像素为单位的预测问题，比图像级分类更具挑战性。 以往的研究[ Structured knowledge distillation for dense prediction,Mimicking very efficient network for object detection]发现，将分类中的KD方法[Distilling the knowledge in a neural network, Fitnets]直接转移到语义分割中可能不会得到令人满意的结果。 严格对齐教师和学生网络之间的逐点分类分数或特征映射可能会实施过于严格的约束，并导致次优解。
逐点分类分数指的是语义分割任务中,对每个像素点得到的类别预测分数。例如对一个包含C个类别的图像进行语义分割,每个像素点都会得到一个C维向量,表示每个类别的预测分数,这个C维向量就是逐点分类分数。
最近的研究[Structured knowledge distillation for dense prediction,Structured knowledge distillation for semantic segmentation,Inter-region affinity distillation for road marking segmentation.]注重加强不同空间位置之间的相关性。 如下图所示，每个空间位置上的特征值（包括最终logits和内部特征映射）被规范化。 然后，通过聚合不同空间位置的子集来处理某些任务特定的关系，例如成对关系和类间关系。 这些方法在获取空间结构信息方面可能比点对齐更有效，提高了学生网络的性能。 然而，特征图中的每个空间位置对知识转移的贡献是相等的，这可能会从教师网络中带来冗余信息。

在本工作中，提出了一种新的基于通道的知识蒸馏方法，即对密集预测任务的每个通道中的激活映射进行归一化，如下图所示。然后，最小化归一化通道激活映射的不对称Kullback-Leibler (KL)发散，它转化为教师和学生网络之间的每个通道的分布。

最近的一些作品利用了通道中包含的知识。通道蒸馏[Channel distillation]提出将每个通道中的激活转移为一个聚合标量，这可能有助于图像级分类，但空间聚合丢失了所有的空间信息，不适合密集预测。其他的工作如MGD、channel exchangin和CSC都显示了通道信息的重要性。MGD将教师通道与学生通道进行匹配，并将其作为一个作业问题来解决。channel exchanging利用融合模块在不同模式的子网之间动态交换channel。
本文表明，对每个通道进行简单的归一化操作可以极大地改善基线空间蒸馏。 提出的通道式蒸馏简单，易于应用于各种任务和网络结构。 主要贡献总结如下：

• 与现有的空间蒸馏方法不同，本文提出了一种适用于密集预测任务的新的通道蒸馏范式。方法简单而有效。
• 在语义分割和目标检测方面，所提出的通道蒸馏方法明显优于现有的KD方法。
• 本文在四个具有不同网络结构的基准数据集上对语义分割和目标检测任务进行了一致的改进，证明了本文的方法是通用的。 鉴于该方法的简单性和有效性，该方法可以作为密集预测任务的强基线KD方法。

related work

多数关于知识蒸馏的研究都集中在图像分类任务上。本文在这里主要研究将知识蒸馏技术有效地应用到密集预测任务的方法,而不仅仅是简单地将图像分类中的像素级蒸馏技术应用到密集预测任务中。
知识蒸馏用于语义分割：在文献中,一个局部相似度图被构建来最小化教师网络和学生网络之间分割边界信息的差异,其中使用从中心像素到8邻域像素的欧式距离作为知识传递。Liu等人提出了两种方法来捕获像素之间的结构信息,包括像素对之间的相似性以及被判别器捕获的整体相关性。文献注重传递同一语义类别的不同像素之间的类内特征变化,其中一组cosine距离被构建来传递结构知识。
知识蒸馏用于目标检测：许多方法发现区分前景和背景区域在目标检测任务的蒸馏中是非常重要的。有文献强制使学生网络中的特征图与教师网络中的特征图在RPN内通过L2损失相似,并发现直接应用像素级损失可能会损害目标检测的性能。 Wang等人提出只提取目标锚点位置附近的细粒度特征。Zhang和Ma生成了蒸馏的注意力掩码来区分前景和背景,并获得了很好的结果。
关于通道级知识的工作： 一些最近的工作也关注每个通道所包含的知识。文献计算每个通道的激活值均值,并对每个通道进行加权差异整合,这对图像分类任务可能有帮助,但是空间聚合会失去所有的空间信息,因此不适合密集预测任务。其他一些工作展示了通道级信息的重要性。总体来说,通道级别的知识蒸馏技术在密集预测任务中的研究还比较少。

本文方法

空间蒸馏

• ATTENTION TRANSFER: 使用注意力掩码来聚合特征图,主要思路是高层语义信息集中在激活值较高的区域,所以可以用激活值的聚合(如求和)来表示关键区域,然后让学生网注意力区域逼近教师网。
• PIXELWISE: 对最后输出层逐像素的分类结果进行对齐,让学生网每个像素的预测分布逼近教师网。
• LOCAL AFFINITY: 计算每个像素与局部邻域(如8邻域)像素的关系,比如欧式距离,然后让学生网的局部亲和力矩阵逼近教师网,以传递局部结构信息。
• PAIRWISE AFFINITY: 计算任意两个像素之间的相似度或关系,构建全局的亲和力矩阵,让学生网逼近教师网的亲和力矩阵,以传递全局结构信息。
• 以上方法主要在最后输出层进行,也可以在中间特征层进行,目的是传递不同粒度的结构知识。

这些方法的简要概述如下。 注意力转移(AT)使用注意力掩模将特征映射压缩到单个通道中进行蒸馏。 像素级损失直接与点状类概率对齐。 局部亲和力由中心像素与其8个邻域像素之间的距离计算。 利用成对亲和力来传递像素对之间的相似性。 计算每个像素的特征与其对应的分类原型之间的相似性来传递结构知识。 [Structured knowledge distillation for dense prediction, Structured knowledge distillation for semantic segmentation]中的整体损失使用对抗方案来对齐来自两个网络的特征映射之间的高阶关系。 注意，后四项考虑了像素之间的相关性。 表1所示的现有KD方法都是空间蒸馏方法。 所有这些方法都将空间位置的N个通道激活值作为操作的特征向量。

Channel-wise Distillation

为了更好地利用每个渠道中的知识，本文建议在教师和学生网络之间柔和地调整相应渠道的激活。 为此，首先将通道的激活转换成概率分布，这样我们就可以概率距离度量来测量差异，例如KL散度。 不同通道的激活倾向于编码输入图像的场景类别的显著性。 此外，一个训练有素的语义分割教师网络显示了每个通道的明确类别特定掩码的激活图–这是预期的。 本文提出了一个新颖的通道式蒸馏范式来指导学生从训练有素的教师那里学习知识。
将教师和学生网络表示为T和S，来自T和S的特征图分别为YT和YS。 通道蒸馏损失可按一般形式表示：

在本例中，φ(·)用于将特征值转换为如下所示的概率分布：

其中c=1,2,…,c对通道进行索引； i是索引通道的空间位置。 T是一个超参数（温度）。 如果使用一个更大的t，概率会变得更软，这意味着本文对每个通道关注更宽的空间区域。 通过应用Softmax归一化，消除了大型网络和紧凑型网络之间大小尺度的影响。 正如[In defense of feature mimicking for knowledge distillation]所观察到的，这种归一化在KD中是有帮助的。 如果教师和学生之间的通道数目不匹配，则使用1×1卷积层来上采样学生网络的通道数目。 评估来自教师网络和学生网络的渠道分布之间的差异。 本文使用KL散度：

KL散度是一个非对称度量。 由上式可以看出，如果φ(yt c,i）很大，那么φ(ys c,i）应该和φ(yt c,i）一样大才能使KL散度最小。 否则，如果φ(yt c,i）很小，则KL散度对最小化φ(ys c,i）的关注较少。 因此，学生网络在前台显著区的特征分布趋于相似，而教师网络的背景区对应的特征对学习的影响较小。 本文假设KL的这种非对称性有利于密集预测任务的KD学习。

结论

在本文中，提出了一个新的通道蒸馏密集预测任务。 与以往的空间蒸馏方法不同，将每个通道的激活归一化为一个概率图。 然后，应用非对称KL散度来最小化师生网络之间的差异。 实验结果表明，在四个具有不同网络主干的公共基准数据集上，所提出的蒸馏方法在语义分割和目标检测方面都优于现有的蒸馏方法。
此外，本文的消融实验证明了通道蒸馏的效率和有效性，它可以进一步补充空间蒸馏方法。 本文希望所提出的简单有效的蒸馏方法可以作为一个强大的基线，为许多其他密集预测任务有效地训练紧凑网络，包括实例分割、深度估计等。