[AAAI22] Knowledge Distillation for Object Detection via Rank Mimicking and Prediction-guided Feature Imitation
Motivation
导致 Student 网络性能差的两个原因:
- 候选框 bbox 差别很大:如图所示,对于简单的目标 (图a),两个模型的最终留下的 bbox 是由同一个 anchor 回归的 (对于 anchor-free 方法是 anchor points)。但是,对于难一点的 bbox (图b) ,却是由不同的 anchor 回归。对于难例样本,两者的 anchor 置信度不同。
- Student 和 teacher 的特征和预测之间都存在着较大 gap。特征差异性 $(F_{dif} = F_{tea} - F_{stu}) $和预测差异性
(
P
d
i
f
=
P
t
e
a
−
P
s
t
u
)
(P_{dif} = P_{tea} - P_{stu})
(Pdif=Ptea−Pstu),是直接通过特征图相减得到的。关键是:二者 gap 位置是不一样的 (图中 inconsistency 处)。这表明这部分区域特征差异性很大,预测差异性却很小。故不需要将对应的特征差异性用来蒸馏。预测指导特征模仿的思路:用预测差异性来指导特征模仿学习。
Method
基于上述两个问题,作者提出了 Rank Mimicking (RM) 和 Prediction-guided Feature Imitation (PFI) 对一阶段的检测器进行蒸馏。前者将 teacher 对候选框的排序作为一种待蒸馏的知识;后者用预测差异性来指导特征差异性,可以提高 student 模型的检测结果。
Rank Mimicking
- 对于一个目标 j j j,假设有 N N N 个 anchor 负责对该个物体进行预测 (positive anchor)。对 student 和 teacher 网络,计算得到这 N N N 个 anchor 的预测的类别分数。
- 分别对 student 网络和 teacher 网络的 positive anchor 类别分数进行 softmax,得到其预测分布:
s i j ′ = exp ( s i j ) ∑ m = 1 N exp ( s m j ) t i j ′ = exp ( t i j ) ∑ m = 1 N exp ( t m j ) s_i^{j^\prime}=\frac{\text{exp}(s_i^j)}{\sum^N_{m=1}\text{exp}(s_m^j)} \ \ \ \ \ \ \ \ t_i^{j^\prime}=\frac{\text{exp}(t_i^j)}{\sum^N_{m=1}\text{exp}(t_m^j)} sij′=∑m=1Nexp(smj)exp(sij) tij′=∑m=1Nexp(tmj)exp(tij) - 最小化 KL 散度让 teacher 和 student 网络 positive anchor 分数的分布保持一致:
L R M = − 1 M ∑ j = 1 M ∑ i = 1 N t i j ′ l o g ( s i j ′ t i j ′ ) L_{RM} = -\frac{1}{M}\sum^M_{j=1}\sum^N_{i=1}t_i^{j^\prime}log(\frac{s_i^{j^\prime}}{t_i^{j^\prime}}) LRM=−M1j=1∑Mi=1∑Ntij′log(tij′sij′)
其中,M表示当前图像中的待检测目标格式。
- Teacher 模型与 student 模型对于所有 positive anchor 的预测类别分数具有同样的分布。
- 将所有 positive anchor 作为一个整体,利用KL散度去蒸馏其分布情况。
Prediction-guided Feature Imitation
现象:预测差异大的位置特征差异可能会很小,而预测差异小的位置特征差异可能会很大。
- 如果直接进行无差别的特征蒸馏,则特征差异较大、预测差异较小的区域会在反向传播中占据主导地位。这会导致特征蒸馏效果变差。
- 采用 teacher 和 student 网络预测类别分数差的
ℓ
2
\ell_2
ℓ2 距离,来表示其预测的差异性。其中
C
C
C 代表类别数,得到的
P
P
P 是一个
H
×
W
H\times W
H×W 的 mask (PFI mask)。某个点的值越大,则代表 S 和 T 的预测差异性越大。
P d i f = 1 C ∑ c = 1 C ∥ P s t u c − P t e a c ∥ 2 2 F d i f = 1 Q ∑ q = 1 Q ∥ F s t u q − F t e a q ∥ 2 2 L P F I = 1 L ∑ l = 1 L 1 H l ∗ W l ∥ P d i f ⊙ F d i f ∥ 2 2 L t o t a l = L t a s k + α L R M + β L P F I P_{dif}=\frac{1}{C}\sum^C_{c=1}\|P_{stu}^c-P^c_{tea}\|^2_2\\ F_{dif}=\frac{1}{Q}\sum^Q_{q=1}\|F_{stu}^q-F^q_{tea}\|^2_2\\ L_{PFI}=\frac{1}{L}\sum^L_{l=1}\frac{1}{H_l*W_l}\|P_{dif}\odot F_{dif}\|^2_2\\ \\ L_{total} = L_{task}+\alpha L_{RM}+\beta L_{PFI} Pdif=C1c=1∑C∥Pstuc−Pteac∥22Fdif=Q1q=1∑Q∥Fstuq−Fteaq∥22LPFI=L1l=1∑LHl∗Wl1∥Pdif⊙Fdif∥22Ltotal=Ltask+αLRM+βLPFI
Experiment