Knowledge Review（CVPR 2021）论文解析

paper：Distilling Knowledge via Knowledge Review

official implementation：https://github.com/dvlab-research/ReviewKD

前言

识蒸馏将知识从教师网络转移到学生网络，可以提高学生网络的性能，作为一种“模型压缩”的方法被广泛使用。15年第一篇蒸馏方法KD对logits进行蒸馏，FitNet提出对feature进行蒸馏，后续AT、PKT、CRD等多种蒸馏方法被提出，这些蒸馏方法大都专注于改进教师与学生同一层级特征的转换和损失函数上。

本文的创新点

本文从另一个新的角度来思考蒸馏问题，即教师和学生网络之间的连接路径，之前的方法大都使用教师网络的同一层级的信息来指导学生，本文认为这是整个蒸馏框架的瓶颈所在。本文提出了一种新的蒸馏方法，关键点在于用教师网络的低层特征来监督学生网络的高层特征，从而提高蒸馏的效果。

本文提出的方法叫做knowledge review，如图1(d)所示，这种review机制用之前（浅层）的特征来指导当前的特征，这意味学生必须总是检查以前学过的东西，以更新对“旧知识”的理解和背景。这和人类的学习曲线很类似，一个小孩只能理解被教知识的一部分，在成长的过程中，过去几年学过的知识越来越多的被记住和理解。

方法介绍

Review Mechanism

假设一张输入图片为 \(\mathbf{X}\)，学生网络表示为 \(\mathcal{S} \)，\(\mathbf{Y}_{s}=\mathcal{S}(\mathbf{X})\) 表示学生的输出logit。其中 \(\mathcal{S} \) 可以分成多个部分 \((\mathcal{S}_{1}, \mathcal{S}_{2},\cdot\cdot\cdot,\mathcal{S}_{n},\mathcal{S}_{c})\)，其中 \(\mathcal{S}_{c}\) 表示classifier部分，\(\mathcal{S}_{1},\cdot\cdot\cdot,\mathcal{S}_{n}\) 表示不同的stage部分，那么 \(\mathbf{Y}_{s}\) 可以表示为

\(\circ \) 表示函数的嵌套，\(g\circ f(x)=g(f(x))\)。中间层的特征表示为 \((\mathbf{F}^1_{s},\cdot \cdot \cdot ,\mathbf{F}^{n}_{s} )\)，第 \(i\) 层的特征计算如下

对于教师网络 \(\mathcal{T} \) 过程是类似的。单层的知识蒸馏可以表示如下

其中 \(\mathcal{M} \) 是转换函数，\(\mathcal{D}\) 是衡量教师和学生之间差异的距离函数。类似的，多层的知识蒸馏可以表示如下

其中 \(\mathbf{I}\) 存储了所有要进行蒸馏的层的特征。

本文提出的 \(review\) 机制是用之前的特征来指导当前的特征。 review机制下单层的蒸馏表示如下

尽管看上去和多层蒸馏很像，但实际上有本质的区别。这里学生的特征固定为 \(\mathbf{F}^{i}_{s}\)，然后使用教师的前 \(i\) 层特征来指导 \(\mathbf{F}^{i}_{s}\)。当review机制和多层知识蒸馏结合时，损失函数表示如下

Residual Learning Framework

基于上述review机制得到的结构如图2(a)所示，推广到多层如图2(b)所示，但是由于不同stage之间巨大的信息差异这种策略不是最优的，并且过程复杂计算量大，比如如果一个网络有 \(n\) 个stage那么损失函数就需要计算 \(n(n+1)/2\) 对特征。为了使过程更加简洁优雅，我们可以将图2(b)的对应的式(6)重新表示如下

其中为了简洁省去了转换函数，然后转换 \(i,j\) 两个求和的顺序，如下

当固定住 \(j\)，式(9)计算的是教师特征 \(\mathbf{F}^{j}_{t}\) 和学生特征 \(\mathbf{F}^{j}_{s}-\mathbf{F}^{n}_{s}\) 之间的距离，然后我们将距离的和近似为融合特征的距离，如下

其中 \(\mathcal{U} \) 是融合特征的模块，得到的结构如图2(c)所示。然后如图2(d)所示，作者用递归的方式对融合的计算进行进一步优化，\(\mathbf{F}^{j}_{s},\cdot\cdot\cdot,\mathbf{F}^{n}_{s}\) 的融合可以看成是 \(\mathbf{F}^{j}_{s}\) 和 \(\mathcal{U}(\mathbf{F}^{j+1}_{s},\cdot\cdot\cdot,\mathbf{F}^{n}_{s})\) 的融合，其中 \(\mathcal{U} (\cdot,\cdot)\) 看作一个递归函数，\(\mathbf{F}^{j+1}_{s}\) 到 \(\mathbf{F}^{n}_{s}\) 的融合定义为 \(\mathbf{F}^{j+1,n}_{s}\)，损失函数可以表示为

图2(d)的结构非常优雅并且利用残差学习的概念简化了蒸馏过程，例如，学生stage4的特征和stage3的特征结合到一起来模拟教师stage3的特征，因此学生stage4的特征学习了教师和学生stage3之间的残差，这种残差信息很有可能是教师产生更高质量结果的关因素。

ABF and HCL

受SENET的启发，作者设计了一个基于注意力的融合模块，如图3(a)所示，高层特征首先resize成和低层特征一样的shape，然后concat一起生成两个 \(H\times W\) 的特征图，然后分别与原本的两个特征图相乘最后再相加得到最终结果。

HCL如图3(b)所示，我们通常用 \(\mathcal{L}_{2}\) 距离作为两个特征图之间的损失函数，但都是用在同一level的特征之间，但本文中不同层级的特征聚集到一起然后向teacher学习，直接利用全局 \(\mathcal{L}_{2}\) 距离不足以传递不同层级之间的复合信息，受PSPNET的启发，作者提出了HCL，利用spatial pyramid pooling将知识的传递划分到不同层级的context信息中。