【论文解读】Knowledge Review：知识蒸馏新解法-技术圈

今天介绍一篇我司的文章Distilling Knowledge via Knowledge Review(缩写为KR)，提出知识蒸馏的新解法。

之前在知乎回答过一个知识蒸馏的问题，为何感觉“知识蒸馏”这几年没有什么成果？https://www.zhihu.com/question/309808462/answer/1591099835

其实知识蒸馏从15年提出以来，主要出现了按照蒸馏位置不同的两种方法分别为蒸logits和features。蒸logits，kd(指15年那篇)之后，主要是通过同时对teacher和student进行约束来保持两者的分布一致性来做的，如PKT、CC等；蒸features，起源于FitNet，后续主要是通过如何更好的利用teacher的feature来做的，如Overhaul等。另外最近因为self-supervised的火热，也出现了如CRD、SSKD等优秀的自监督相关的kd算法。

我在回答中也提到，之前的知识蒸馏方法中Overhaul效果是最好的，一骑绝尘。之前的方法把知识蒸馏的上限推到了一个新的高度，貌似已经达到饱和了，那么就引出来一个值得思考的问题，之前的蒸馏feature和logits的方式，teacher的信息真的完全被利用了吗？

显然还没有

Knowledge Review

如上图，其中(a)图是蒸logits，(b)是蒸单层feture，(c)是蒸多层feature。之前的蒸馏方法(a)(b)(c)三种形式都可以理解为相同层进行蒸馏，引导student学习teacher相同层的feature。但是这种相同层蒸馏方式在训练的过程中显然是不合理的，teacher深层的语义信息更加抽象，在student的训练初期，去学习teacher的抽象深层语义信息，这就好比是让一个婴儿去学习大学知识，这显然是非常困难且难以学习的。

于是就引出了knowledge review的方法，如图(d)所示，student的单层feature通过teacher的多层feature同时引导学习，这种引导学习的方式比起之前的方法更加完全，teacher的浅层可以认为是比较简单的知识，深层是比较抽象的知识，student训练早期会更加关注teacher浅层简单知识的学习，随着训练的进行，会越来越关注teacher深层抽象知识的学习，knowledge review这个词非常形象表达了这个过程——温故知新，浅层是old knowledge，深层是new knowledge。

knowledge review的想法是非常make sense的，但是想要做work，还需要一些针对性的设计。

Review Mechanism

下面先通过符号假设推导出review mechanism的函数表达式。

假设一张图片输入是 $\mathbf { X }$ ，student网络是 $\mathcal { S }$ ，那么student的logit输出为 $\mathbf { Y } _ { s } = \mathcal { S } ( \mathbf { X } )$ 。其中 $\mathcal { S }$ 可以分成 $\left( \mathcal { S } _ { 1 } , \mathcal { S } _ { 2 } , \cdots , \mathcal { S } _ { n } , \mathcal { S } _ { c } \right)$ 不同部分， $\mathcal { S } _ { c }$ 是classifier部分，其余是不同stage部分，那么

$\mathbf { Y } _ { s } = \mathcal { S } _ { c } \circ \mathcal { S } _ { n } \circ \cdots \circ \mathcal { S } _ { 1 } ( \mathbf { X } )$

$\circ$ 表示复合函数 $g \circ f ( x ) = g ( f ( x ) )$ ，中间层的feture表示为 $\left( \mathbf { F } _ { s } ^ { 1 } , \cdots , \mathbf { F } _ { s } ^ { n } \right)$ ，那么第i层的feature的公式为

$\mathbf { F } _ { s } ^ { i } = \mathcal { S } _ { i } \circ \cdots \circ \mathcal { S } _ { 1 } ( \mathbf { X } )$

teacher的函数表达式跟student相类似。

那么一层知识蒸馏的loss可以表示为

$\mathcal { L } _ { S K D } = \mathcal { D } \left( \mathcal { M } _ { s } ^ { i } \left( \mathbf { F } _ { s } ^ { i } \right) , \mathcal { M } _ { t } ^ { i } \left( \mathbf { F } _ { t } ^ { i } \right) \right)$

其中 $\mathcal { M }$ 是对feature进行转换的函数。 $\mathcal { D }$ 是用来衡量转换后的teacher和student的距离函数。

那么多层知识蒸馏的loss可以表示为

$\mathcal { L } _ { M K D } = \sum _ { i \in \mathbf { I } } \mathcal { D } \left( \mathcal { M } _ { s } ^ { i } \left( \mathbf { F } _ { s } ^ { i } \right) , \mathcal { M } _ { t } ^ { i } \left( \mathbf { F } _ { t } ^ { i } \right) \right)$

review mechanism的单层知识蒸馏loss则可以表示为

$\mathcal { L } _ { S K D _ { - } R } = \sum _ { j = 1 } ^ { i } \mathcal { D } \left( \mathcal { M } _ { s } ^ { i , j } \left( \mathbf { F } _ { s } ^ { i } \right) , \mathcal { M } _ { t } ^ { j , i } \left( \mathbf { F } _ { t } ^ { j } \right) \right)$

表示student第i层同时对应teacher的从1到i层的loss之和。

那么review mechanism的多层知识蒸馏loss则可以表示为

$\mathcal { L } _ { M K D _ { - } R } = \sum _ { i \in \mathbf { I } } \left( \sum _ { j = 1 } ^ { i } \mathcal { D } \left( \mathcal { M } _ { s } ^ { i , j } \left( \mathbf { F } _ { s } ^ { i } \right) , \mathcal { M } _ { t } ^ { j , i } \left( \mathbf { F } _ { t } ^ { j } \right) \right) \right)$

在知识蒸馏的训练过程中，review mechanism的loss相当于一个正则项，测试的时候，只在student网络上inference，没有额外的时间损耗。

Residual Learning Framework

根据上述review mechanism推导得到的框架如图(a)，灰色块是student做了转换后的feature，绿色块是蒸馏的距离函数计算方式。图(b)对图(a)进一步进行设计，直接将每个单层student对应多层teacher，但是直接对student每一个stage进行蒸馏，不同层的特征会产生冲突相互干扰。

将review mechanism的loss函数简写成

$\mathcal { L } _ { M K D _ { - } R } = \sum _ { i = 1 } ^ { n } \left( \sum _ { j = 1 } ^ { i } \mathcal { D } \left( \mathbf { F } _ { s } ^ { i } , \mathbf { F } _ { t } ^ { j } \right) \right)$

忽略掉对feature的转换函数。然后将i和j两个顺序进行交换

$\mathcal { L } _ { M K D _ { - } R } = \sum _ { j = 1 } ^ { n } \left( \sum _ { i = j } ^ { n } \mathcal { D } \left( \mathbf { F } _ { s } ^ { i } , \mathbf { F } _ { t } ^ { j } \right) \right)$

可以将上式括号内部分近似成student多层feature融合求距离的公式

$\sum _ { i = j } ^ { n } \mathcal { D } \left( \mathbf { F } _ { s } ^ { i } , \mathbf { F } _ { t } ^ { j } \right) \approx \mathcal { D } \left( \mathcal { U } \left( \mathbf { F } _ { s } ^ { j } , \cdots , \mathbf { F } _ { s } ^ { n } \right) , \mathbf { F } _ { t } ^ { j } \right)$

其中 $\mathcal { U }$ 是一个融合函数。那么整个知识蒸馏框架就变成了图(c)，红色块表示融合模块。

最后图(d)再进一步对图(c)的融合模块进行优化，提高效率。融合 $\mathbf { F } _ { s } ^ { j } , \cdots , \mathbf { F } _ { s } ^ { n }$ 的过程可以看成是 $\mathbf { F } _ { s } ^ { j }$ 和 $\mathcal { U } \left( \mathbf { F } _ { s } ^ { j + 1 } , \cdots , \mathbf { F } _ { s } ^ { n } \right)$ 融合，于是就可以将 $\mathcal { U } ( \cdot \cdot \cdot )$ 看成一个递归函数，重复进行融合操作，将 $\mathbf { F } _ { s } ^ { j + 1 , n }$ 看成是 $\mathbf { F } _ { s } ^ { j + 1 }$ 到 $\mathbf { F } _ { s } ^ {n}$ 的融合，那么loss函数可以表示成

$\mathcal { L } _ { M K D _ { - } R } = \mathcal { D } \left( \mathbf { F } _ { s } ^ { n } , \mathbf { F } _ { t } ^ { n } \right) + \sum _ { j = n - 1 } ^ { 1 } \mathcal { D } \left( \mathcal { U } \left( \mathbf { F } _ { s } ^ { j } , \mathbf { F } _ { s } ^ { j + 1 , n } \right) , \mathbf { F } _ { t } ^ { j } \right)$

有没有觉得这个式子似曾相识，这就是ResNet的residual learning的递推公式啊，这个residual learning的设计可以帮助knowledge review的训练过程更加稳定，容易优化。

其中红色块和绿色块分别设计了ABF和HCL模块。

ABF模块设计和SKNet非常类似，将student后面一个stage转换得到的feature和当前stage的feature看成是两条分支，然后通过concat、卷积的方式得到两条分支的spatial weight，最后reweight后相加得到融合的feature。

HCL模块设计借鉴了PSPNet的方式，用四种尺度AvgPool分别对student和teacher的feature进行池化，然后计算loss。

通过上述一系列的设计，knowledge review将温故知新的思想发挥到了极致。

表格中的OFD是Overhaul，可以看到knowledge review比之前的Overhaul、CRD等方法都要好上一大截，而且在不同模型架构上都有效。

总结

知识蒸馏是计算机视觉中一个很有趣的问题，之前的方法都是从如何更好的利用好teacher的信息出发，将teacher和student相同层的信息利用起来，而knowledge review进一步引入温故知新的思想，将teacher的信息利用到了极致。但是知识蒸馏仍然存在着超参敏感的问题，后续或许可以结合一些self-supervised的方法像之前的CRD，或者通过一些特殊的设计来增强知识蒸馏的稳定性，这可能是知识蒸馏发挥实用价值当务之急的事情。

Reference

paper链接: https://jiaya.me/papers/kdreview_cvpr21.pdf


往期精彩回顾




适合初学者入门人工智能的路线及资料下载
机器学习及深度学习笔记等资料打印
机器学习在线手册
深度学习笔记专辑
《统计学习方法》的代码复现专辑
AI基础下载
机器学习的数学基础专辑
温州大学《机器学习课程》视频
本站qq群851320808，加入微信群请扫码：