今天介绍一篇我司的文章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

【论文解读】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,还需要一些针对性的设计。

02

Review Mechanism

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

假设一张图片输入是   ,student网络是   ,那么student的logit输出为   。其中   可以分成   不同部分, 

 表示复合函数   ,中间层的feture表示为 

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

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

其中   是对feature进行转换的函数。 

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

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

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

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

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

03

Residual Learning Framework

【论文解读】Knowledge Review:知识蒸馏新解法_计算机视觉_02

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

将review mechanism的loss函数简写成

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

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

其中 

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

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

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

【论文解读】Knowledge Review:知识蒸馏新解法_机器学习_03

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

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

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

【论文解读】Knowledge Review:知识蒸馏新解法_人工智能_04

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

04

总结

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

Reference

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