定位蒸馏
定位蒸馏:Localization Distillation,简称LD
论文地址:
Localization Distillation for Dense Object Detection
开源代码地址:
https://github.com/HikariTJU/LD
MMDetection:https://github.com/open-mmlab/mmdetection/tree/master/configs/ld
方法概括:
把用于分类head的KD(知识蒸馏),用于目标检测的定位head,即有了LD(Localization Distillation)
head
head是获取网络输出内容的网络,利用之前提取的特征,head利用这些特征,做出预测。
做法:先把bbox的4个logits输出值离散化成4n个logits输出值,之后与分类KD完全一致。
意义:LD使得logit mimicking首次战胜了Feature imitation。分类知识与定位知识的蒸馏应分而治之、因地制宜
一、LD的诞生背景
1. bbox分布与定位模糊性 -- 痛点
关于bbox分布建模,它主要来源于GFocalV1 (NeurIPS 2020)与Offset-bin (CVPR 2020)这两篇论文。
bbox的表示通常是4个数值,一种如FCOS中的点到上下左右四条边的距离(tblr),还有一种是anchor-based检测器中使用的偏移量,即anchor box到GT box的映射(encoded xywh)。
GFocalV1针对tblr形式的bbox建模出了bbox分布,Offset-bin则是针对encoded xywh形式建模出了bbox分布,它们共同之处就在于尝试将bbox回归看成一个分类问题。并且这带来的好处是可以建模出bbox的定位模糊性。
那么用n个概率值去描述一条边,可以显示出模型对一个位置的定位模糊估计,越尖锐的分布说明这个位置几乎没有模糊性(比如大象的上边界),越平坦的分布说明这个位置有很强的模糊性(大象的下边界)。当然不光是bbox分布的平坦度,形状上还可分为单峰型,双峰型,甚至多峰型。
2. 知识蒸馏
Knowledge Distillation (KD)最早是针对图像分类而设计的。核心思想是:用一个预训练的大模型(teacher)去指导一个小模型(student)的学习。
网络模型输出n个logits z,z经过一个带有温度 τ
的softmax变换 S(⋅,τ)
,便得到了概率分布 p=S(z,τ)
.
那么只要去最小化student的概率分布 \(p_s\) 与teacher的概率分布 \(p_T\)的KL loss
,就可以将大模型的知识迁移至小模型。这里温度 τ>1 对图像分类任务是常见的选取,即软化分类知识。
二、定位蒸馏
看到这里,LD的idea已经不言自明,bbox的一条边是n个logits z,一个bbox就有4个 z,每个 z 作用上一个带有温度 τ 的softmax函数,把定位知识软化一下,随后一样的KL loss,让student的bbox分布去拟合teacher的bbox分布。
为验证LD的有效性,我们先在检测器的positive location上进行,也就是哪里进行了bbox regression,哪里就进行LD。大模型teacher是一个已经训练了24 epochs的高精度检测器(比如ResNet-101),小模型student可以是ResNet-50。在COCO数据集上,我们只需要稍微调一调温度 τ 就可以在GFocalV1的基础上提高1.0AP,特别是AP75的提升最为显著,说明LD确实显著改善了定位精度。
由以上可以看到,定位蒸馏LD与分类蒸馏KD从公式上来看是完全一致的,都是针对head的输出logits做知识传递,这为目标检测知识蒸馏提供了一个统一的logit mimicking框架。
三、分类KD的低效性与Feature imitation
以往许多工作指出了分类KD的蒸馏效率低下(涨点低),这主要有两个方面:
-
1、在不同的数据集中,类别数量会变化,较少的类别可能给student提供不了很多有用的信息。
-
2、一直以来logit mimicking都只能在分类head上操作,而无法在定位head上操作,这自然忽视了定位知识传递的重要性。
基于这两个原因,人们将视线转向了另一个很有前景的知识蒸馏方法,Feature imitation。该方法主要受到FitNet的启发,一言以概之,就是不光分类head上要做logit mimicking,中间隐藏层(特征图)也要让student去拟合teacher,通过最小化L2 loss来完成。
形成了如下的目标检测知识蒸馏框架:
分类、定位和特征图:
其中分类head上是logit mimicking(分类KD),特征图上是Feature imitation (teacher与student特征图之间的L2 loss),定位head上是伪bbox regression,即把teacher预测框当成额外的回归目标。
Feature imitation在师生的特征图上施加监督,最常见的做法是先将student的特征图尺寸与teacher特征图对齐,之后再选择一些感兴趣的区域作为蒸馏区域,例如FitNet(ICLR 2015)在全图上蒸馏;Fine-Grained(CVPR 2019)在一些anchor box的location上蒸馏;还有DeFeat(CVPR 2021)在GT box内部用小loss weight,在GT box外部用大loss weight;亦或者是GI imitation(CVPR 2021)的动态蒸馏区域,但无论选择何种区域,最后都是在蒸馏区域上计算二者的L2 loss.
那么Feature imitation有什么好处呢?
在multi-task learning框架下,特征图相当于树根,下游的各个head相当于树的叶子。那么特征图显然包含了所有叶子所需要的知识。进行Feature imitation自然就会同时传递分类知识与定位知识,而分类KD却无法传递定位知识。
Featrue imitation有什么弊端?
答案自然还是它会在蒸馏区域中的每个location上同时传递分类知识与定位知识。
前后一对比,乍看之下岂不矛盾?让我来解释一下。
分类知识与定位知识的分布是不同的。这一点在以往的工作中有提到,例如Sibling Head (CVPR 2020)。
两种知识的分布不同,自然就导致了并不是在一个location上同时传递分类知识与定位知识都有利。很有可能某些区域仅对分类知识传递有利,也有可能某些区域仅对定位知识传递有利。换言之,我们需要分而治之、因地制宜地传递知识。这显然就是Feature imitation无能为力的事情了,因为它只会传递混合知识。
于是我们利用multi-task learning天然地把知识解耦成不同类型,这就允许我们在一个区域中有选择性地进行知识蒸馏。为此,我们引入了一个有价值定位区域 VLR (Valuable Localization Region)的概念,来帮助我们进行分而治之的蒸馏。
与以往的Feature imitation方法不同,我们的蒸馏分为两个区域:
-
Main distillation region (主蒸馏区域):即检测器的positive location,通过label assignment获得。
-
VLR:与一般的label assignment做法类似,但区域更大,包含了Main region,但去掉了Main region。VLR可以视为是Main region的向外扩张。
于是我们来探究一下在这两个区域上进行分类的KD与定位的LD会有什么效果。
这个表格是本文的一大精髓所在,有几个有趣的现象。
-
1、可以看到在Main区域上,KD与LD均有效,说明在这个区域上传递分类或定位知识都有好处。但明显LD提升更大,说明定位知识的传递更有利于性能的提升。
-
2、VLR上LD也有效,这也是为什么我们把这样的区域命名为有价值定位区域。但VLR KD的加入却损害了性能。
于是就得到了本文的logit mimicking策略,Main KD + Main LD + VLR LD
有了这个logit mimicking策略,我们可以与Feature imitation对比一下
这个实验首次展示了logit mimicking的巨大潜力,也说明了logit mimicking多年以来的蒸馏效率低下的原因是缺少有效的定位知识传递。当引入了LD后,弥补了这一缺陷,logit mimicking居然可以超过Feature imitation。
当然上表也展示了,最优的蒸馏策略依然是logit mimicking与Feature imitation都用上。只是可以注意到的是,在有了logit mimicking之后,各个Feature imitation的性能差异也不是很明显了,选择哪个蒸馏区域都差不了多少。
下图展示的是student与teacher的平均分类误差与定位误差。
可以看到一些代表性Feature imitation方法(如Fine-Grained,GI)确实可以同时降低分类误差与定位误差。在仅用上LD时,定位误差得到明显下降,但分类误差无法降低也是可以理解的。而在进一步用上了KD之后,两个误差都得到了明显下降。
下图则是在两个FPN层级上的定位误差可视化