文 | 陀飞轮@知乎

今天介绍一篇我司的文章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的信息真的完全被利用了吗？

显然还没有。

01 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的函数表达式。

假设一张图片输入是 ，student网络是 ，那么student的logit输出为 。其中\mathcal { S }可以分成 不同部分， 是classifier部分，其余是不同stage部分，那么

表示复合函数 ，中间层的feture表示为 ，那么第i层的feature的公式为

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

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

其中 是对feature进行转换的函数。 是用来衡量转换后的teacher和student的距离函数。

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

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

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

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

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

Residual Learning Framework

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

将review mechanism的loss函数简写成

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

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

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

最后图(d)再进一步对图(c)的融合模块进行优化，提高效率。融合 的过程可以看成是 和 融合，于是就可以将 看成一个递归函数，重复进行融合操作，将 看成是 到 的融合，那么loss函数可以表示成

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

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

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

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

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

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

总结

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

后台回复关键词【入群】

加入卖萌屋NLP/IR/Rec与求职讨论群

后台回复关键词【顶会】

获取ACL、CIKM等各大顶会论文集！

 

[1]. Distilling Knowledge via Knowledge Review: https://jiaya.me/papers/kdreview_cvpr21.pdf