【专栏:前沿进展】在青源LIVE第31期中,旷视研究院张祥雨团队的张培圳研究员深入浅出地为我们介绍了其团队被 NeurIPS 2021 录用的论文「用于物体检测的实例条件知识蒸馏」。本期报告首先简要回顾了用于图像识别、物体检测的知识蒸馏技术的发展历程,介绍了论文的研究动机、模型架构、训练技巧以及实验结果。
本文整理自青源 LIVE 第 31 期,视频回放链接:https://event.baai.ac.cn/activities/204
主讲:张培圳
整理:熊宇轩
审校:李梦佳
01
背景知识
早在 NIPS 2014 上,Hinton 等人在论文「Distilling the knowledge in a Neural Network」中就开始研究用于图像分类的知识蒸馏技术。根据知识源的不同,我们可以将用于图像分类的知识蒸馏技术分为「基于 Logits 响应」、「基于中间特征图激活」和「基于关系建模」的知识蒸馏。
如上图所示,基于响应的知识蒸馏往往会在弱监督场景下计算预训练教师网络和学生网络用于分类的 Logits 的 KL 散度;基于激活的知识蒸馏会计算教师网络和学生网络中间特征图的 MSE;基于关系建模的知识蒸馏则会针对教师网络和学生网络对每一个 mini-batch 中图片经过全局平均池化后的特征矩阵计算 MSE 等距离。
面向物体检测任务的知识蒸馏技术大多属于基于中间特征图激活的知识蒸馏。近年来,一些面向工程化的工作也会尝试系统地使用上述知识蒸馏方式的组合(例如,FBKD、GID、DeFeat、G-DetKD,等)
相较之下,分类模型需要处理的图像往往包含单个物体,待分类的物体位于图像的中央;而检测模型需要处理的图像中往往包含多个物体,这些物体从属于各种类别、分布于图像的各个位置,图像中还包含大量冗余的背景信息。
在针对中间特征图进行蒸馏时,蒸馏特征的位置是十分关键的,我们通常会利用以下三类拓扑信息:
(1)前景先验:给定训练图片中物体的边界框 BBOX,我们将其映射到中间特征图上,对该区域进行知识蒸馏;对于基于 anchor 的物体检测任务,重点对与 groundtruth 重叠比例较高的 anchor 区域进行蒸馏;对于多阶段物体检测任务,针对 proposal 进行蒸馏;
(2)前景和背景先验:将前景、背景分开蒸馏/在全图上使用注意力机制;
(3)基于度量的区域:根据设计出的度量方式找出有效的蒸馏信息(例如,教师网络与学生网络的差异)。
如上图所示,蓝色和棕色的实线方框代表真实的物体边界框;黑色的方框代表物体附近的 anchor;蓝色、棕色、红色的虚线方框代表多阶段检测模型提出的 proposals,可以通过非极大抑制等手段对其进行修正。
在考虑背景先验时,我们可以通过二值掩膜将前景和背景分隔开,进而分别进行蒸馏。此外,我们可以在全图中分别在空间和通道上使用注意力机制,我们可以分别对教师和学生网络计算空间维度和通道维度上的注意力图,进而对这些注意力两两进行知识蒸馏。
在利用教师网络、学生网络评分差异进行知识蒸馏时,假设教师网络和学生网络预测头的预测特征图的维度相同,图像上的某个像素点邻近于某个 anchor 的预测框,模型会计算出该预测框从属于某一类的置信度以及物体回归的位置编码信息。在典型的情况下,教师网络可以很好地预测出类别和边框位置信息,而学生网络对于某些类别物体的预测性能较差。因此,我们可以基于度量学习的方式,重点对学生网络学习性能较差的类别进行蒸馏。
以 ADL 模型为例,其关键之处在于计算自适应蒸馏权重(ADW)。如上图所示,红色圆圈代表教师网络和学生网络在某个空间位置上预测评分的 KL 散度。此外,在 GID 模型中,作者定义了一种通用实例(General Instance)代表教师网络或学生网络的 proposal,这两种网络对proposal 会产生预测得分的差异,GI Score 为针对所有类的得分差异的最大值,该值越大说明不确定性越强。
基于前景/背景的建模方法非常直接,符合人类的认知习惯,但是这些方法获得的知识可能是次优的。基于度量的方法相对来说较为灵活,但是往往并不直观,并且存在多种可能的度量方法。
02
用于物体检测的实例条件知识蒸馏
为了解决上述问题,旷视研究院张祥雨组提出了「用于物体检测的实例条件知识蒸馏」,在保证模型具有直观性、可解释性的同时,兼具了基于度量的蒸馏方法的灵活性。模型的整体架构如上图所示。
具体而言,该方法利用人类标注的实例作为条件,来查询教师网络的特征,从而检索出有用的知识,用于蒸馏学生网络的特征图。
如上图所示,左侧为 Query 的特征构造过程,右侧为抽象的建模过程。从抽象建模过程开始,我们可以将图中物体的 one-hot 类别信息和实例的位置嵌入通过多层感知机编码到 Query 中。这样一来,每个实例的特征都可以被表示为一个低维向量。
我们利用上述实例标注在教师网络的特征图上提取有效信息。每个特征图是一个三维(H*W*D)的张量,我们可以将其在空间维度上展平,得到特征 (指展平后教师特征的个数)。
接着,我们通过「键-值」形式的注意力,将物体标注和保留了空间-拓扑信息的特征图像素之间的关系实例化,从而得到了「实例感知注意力掩模」。这样一来,我们就可以衡量物体和图像位置的关系。其中,基于 Query、Key、Value 的自注意力的计算范式与原始的 Transformer 架构中的计算范式相同。
如上图所示,我们获得的实例条件知识为。其中,为注意力掩模,为教师网络的特征图,i 为图中第 i 个物体的索引,j 为多头注意力机制中的第 j 个注意力头。
为了保证掩模与待检测物体的位置相关,我们设计了面向检测任务的辅助任务引导对掩模的学习。具体而言,我们利用注意力机制的输出作为识别、定位辅助任务的监督信号。值得注意的是,辅助任务不会更新学生网络的参数。在识别任务中,我们首先根据数据集先验采样出一些伪样本。对于识别出的真实样本,我们进一步对其位置进行定位。
如上图所示,在公式(10)中,辅助损失中的第一项是针对识别任务的二值交叉熵损失,第二项为对真实样本的定位回归损失。其中,为实例查询与教师特征图之间注意力机制的输出,为由多层感知机和 Sigmoid 函数构成的分类头,当查询的标注物体为真时,反之其值为 0。
然而,注意力机制的输出依赖于注意力掩模和教师网络特征图,而注意力掩模又依赖于标注。因此,通过公式(10)惩罚注意力机制的输出和之间的分类误差,往往会使网络学习到一个平凡解,从而使「实例感知注意力掩模」失效。
为此,我们引入了信息丢失(information dropping)技术,破坏实例标注的信息。对于回归任务,我们对物体边界框的中心进行抖动,并且对尺度信息进行量化。这样一来,我们就会强迫模型从教师网络特征图上学习到更多的信息。
因为信息丢失和定位回归损失之间有联系,信息丢失能强迫从教师网络特征图上学习到更多的信息。我们的定位任务不是像一般检测器那样去回归物体框,而正是去回归抖动的物体框中心和抖动前的框中心的这个抖动距离。因为 本身是基于抖动后的框信息编码的,它如果想把上面这个任务做好,降低其损失,必须从教师特征图上面去感知到这个真实物体中心。
整体而言,本文所提出的框架的损失函数由蒸馏损失、学生网络的检测损失、辅助任务损失构成,其中为真实样本个数。
03
实验结果
与现有的知识蒸馏方法相比,本文提出的方法在 COCO 数据集上取得了较优的性能。在使用 RetinaNet 作为物体检测网络,采用多尺度训练时,本文提出的方法可以将学生网络(ResNet-50)的物体检测 AP 从 37.4 提升到 40.7,这一数字甚至超过了教师网络(ResNet-101)的 40.4 的 AP。在训练过程中进行的,涉及物体Groundtruth信息的实例建模为学生从单纯的教师网络特征图信息以外获得更好训练,提供了可能。
在其它的检测和分割模型上,本文提出的知识蒸馏方法也取得了优秀的性能表现。如上图所示,继承策略(Inheriting Strategy)是一种非常有用的训练技巧:当教师网络和学生网络的Neck部分和检测头部分的架构相同时,学生网络就可以继承上面两个部分的参数进行良好的初始化,从而进一步提升性能或快速收敛。
在使用 EfficientNet-B0、MobileNet-V2 等面向移动设备的主干网络时,本文提出的知识蒸馏方法可以大大提升物体检测的 AP。
Transformer 解码器的注意力头在原图中的激活可视化结果如上图所示。实验结果表明,注意力头可以有效注意到实例周围的信息。
04
讨论
在整体的计算过程中,计算量主要体现在教师网络的预训练和学生网络本身的训练开销,而本文加入的辅助任务和知识蒸馏的计算开销相对较小。在知识迁移方面,本文提出的方法可以直观地检索图像中有意义的内容,引入了类似基于度量的蒸馏方法的灵活特性,具有一定的可解释性。本文采用的继承策略训练技巧可以有效地加快收敛。
05
结语
在本文中,旷视研究院提出了一种面向物体检测的新型知识蒸馏框架,该框架可以学习特征解码器模块从而检索有用的知识,设计了辅助任务来优化解码器。作者设计了基于实例的条件,从而适应目标检测任务,其中 Query 为标注的实例信息,Key 和 Value 为教师网络的表征知识。
欢迎点击阅读原文参与文章讨论。