date: 2020/05/02

DADDY! DADDY! DO! - 鈴木雅之/鈴木愛理:继“嗷~辣妹~迷死他”之后又一次被大叔洗脑

大叔这声线加上小姐姐的和声,真是绝了!!


知识蒸馏(Knowledge Distillation, KD),按照字面意思,就是将某个经过训练网络的知识,蒸馏到另一个网络上去。换句话说,存在一个已经训练好的,具备知识的成熟教师网络(Teacher),用它来指导学生网络(Student)的学习过程。对学生网络而言则是一个模仿过程,模仿什么和如何模仿是知识蒸馏的两个核心问题。这个过程通常是通过复用教师网络权重,或者度量最终输出/隐含特征之间的差异作为辅助损失来实现的。


简单来说,知识蒸馏包含以下三个优势:

  1. 加快训练
    直观上理解,教师网络把学习方法教给了学生,使得学生网络能更快地收敛;
  2. 提高性能
    教师网络通常采用更加冗余复杂、表现更好的网络,因此可以指导学生网络得到一个恰当的输出分布或是隐含特征的分布。
    直观上理解,教师网络把自己的经验作为先验知识传授给了学生,使得学生网络能够继承先辈经验,从而有更好的表现(比如选择题的“三短一长选最长”);
  3. 迁移学习
    对于相似任务,教师网络提取的特征(尤其是浅层特征)可能有很好的指导意义,用大数据集训练过的教师网络来指导学生网络在小数据集上训练是有很大帮助的。
    直观上理解,就像是物理老师在教学生解数学题,尽管不是专门的数学老师,但由于物理中频繁使用到数学工具,因此经验丰富的物理老师依旧拥有较高的数学水平,可能就有资格教育学生解答数学问题。


知识蒸馏的初步思想最早于《Model compression (2006)》一文提出,直到《Do Deep Nets Really Need to be Deep? (NIPS2014)》才有比较成熟的知识蒸馏思路提出,次年Hinton大神终于在《Distilling the Knowledge in a Neural Network (2015)》给这类技术冠上了“知识蒸馏”这样一个名字并一直沿用至今。


早期知识蒸馏只是在模仿标签也即模型最终输出的一个分布,后来引入了隐含特征的模仿,如今两者通常是结合使用的。分类任务上标签容易模仿,但目标检测、语义分割等更高级的任务在模仿的时候就不得不多动点心思,目前通常是用教师网络来生成一些指导意义的“伪标签”,然后和真实标签分别与预测结果计算损失,然后加权求和作为总体损失来实现。作为模型压缩的一种形式,知识蒸馏也常常跟其他模型压缩手段如裁剪、量化进行“联动”,最简单的一种形式即是将压缩前、压缩后的网络分别视作教师网络、学生网络来开展恢复训练。


参考:


模仿标签

模仿标签可以分为两种,一种是学生纯粹的模仿老师,我称之为“单纯的模仿”;另一种是学生在模仿老师的同时,还在自己思考如何更好的解题,我称之为“有思考的模仿”。
用书法打个比方,前者就只是单纯的临摹,模仿名家的字体,后者则是在临摹的同时摸索出一套适合自己的书法,最终甚至自成一派。


单纯的模仿

Do Deep Nets Really Need to be Deep? (NIPS2014)》设计了一种Mimic Learning(请记住mimic这个词,它一直被后续的知识蒸馏论文所沿用),它在不使用真实标签的情况下,用教师网络的最终输出来直接指导学生网络的学习。此时loss即为

其中,

分别是教师网络和学生网络的输出;

可以是任意的距离函数,比如L1、L2、KL等;

论文在实验中指出L2应该是最好的距离函数,从后续的论文中我们也可以看到,大多数情况下都默认地使用了L2这个距离函数进行度量。


这篇论文也通过实验证明了“浅网络的知识容量并不比深网络小”这样一个核心论点,作者训练了一个较深的网络作为教师网络,然后又设计了一个较浅但参数规模与教师网络相近的学生网络,用教师网络的输出作为训练用的标签指导学生网络训练(注意这里没有使用真实标签),训练出来的学生网络并不亚于教师网络的表现。


有思考的模仿

完全模仿教师网络的输出并不那么合理,Hinton大神在Distilling the Knowledge in a Neural Network (2015)》中提出同时利用真实标签和教师网络输出的软标签进行训练的一种训练形式,并从此给这一类方法冠上了“知识蒸馏”的名字。这种方式简单有效,一直被广泛使用,常见的深度学习框架中的softmax函数也能看到相应的参数设置,因此使用起来非常方便。


具体来说,

损失包含硬标签损失和软标签损失两部分,

硬标签即真实标签,其损失即传统的分类损失;

软标签则是教师网络的softmax输出标签,这里做softmax之前还对特征值分别除以温度,它代表了对标签的软化程度,通常取越大,标签的概率分布就越软。

同时引入一个权重来权衡两个标签损失的重要程度,越大对老师的信赖程度就越高,通常训练初期可以给一个比较大,而后期逐渐降低,直观点理解就是“师父领进门,修行看个人”啦~

image.png

(KD计算示意,图源GJCfight


除了提前训练一个教师网络之外,《Apprentice: Using Knowledge Distillation Techniques To Improve Low-Precision Network Accuracy (2017)》指出同时对大网络和小网络进行联合训练也是一种可选的方案,比如将损失函数设计为

image.png

其中,

是损失函数;

分别是教师网络和学生网络在softmax前的特征向量;

分别是教师网络和学生网络在softmax后的预测概率;

是真实标签;

是三个人为定义的权重


模仿隐含特征

除了模仿最终结果,模仿一些隐含的中间结果也是有价值的。该思路最早是由FitNets提出,而且称之为“暗示训练(Hint Traing)”,在原先模仿标签的知识蒸馏的基础上又提出模仿隐含特征的知识蒸馏形式,当然两者不是对立了,往往我们都是将两者结合起来设计损失,在模仿隐含特征的同时,也对最终标签进行模仿。因为涉及隐含特征的学习,因此也有机会构造由浅入深、由易到难的学习策略,也即我们常说的课程学习(Curriculum Learning)。


隐含特征就比标签有意思多了,毕竟它在定义上就非常抽象模糊,有相当大的发挥空间。


模仿中间结果

论文:《FitNets: Hints for Thin Deep Nets (ICLR2015)

image.png


暗示训练(Hint Training)指的是建立教师网络和学生网络之间的隐含联系,比如取教师网络、学生网络的某一层输出特征图来建立联系,通过约束迫使学生网络模仿教师网络产生类似的隐含特征。为了配对特征图的尺寸,可以引入一个简单的回归模块,将特征图Resize到教师网络隐含特征图相同的shape上去,方便计算损失(如L2损失)。暗示训练有明显的正则化效果,而且模仿的特征图越深,学生网络的训练自由度就越低,正则化力度就越强。


论文提出的知识蒸馏分两个阶段,第一阶段为暗示训练,只训练部分网络结构,让学生网络先去模仿如何产生相近的隐含特征;第二阶段则为普通的、对标签的模仿。通常用浅宽的易训练网络来指导深窄的难训练网络学习,但随着BN层的普及,深窄网络其实不像以前那么难训练了,这篇论文更大的价值是为知识蒸馏开创了模仿隐含特征这样一种思路。


模仿注意力

论文:《Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer (ICLR2017)

这里的注意力主要是空间上的注意力,采用两种注意力图(激活注意力图和梯度注意力图),并建立教师网络与学生网络直接注意力图的约束,鼓励学生网络模仿教师网络在提取特征时的注意力分布和对输入的敏感程度(梯度)的分布。


激活注意力图

为了得到激活注意力图,需要找到一种聚合函数,使得特征图能沿着通道方向拍扁

有三种简单的聚合方案:

  1. 绝对值之和
  2. 绝对值的p次方之和
  3. 最大值的p次方之和


image.png

从上图不同列的激活注意力图可以看出,随着p的将越来越关注注意力更大的区域而抑制注意力不那么强的区域,则干脆只关心注意力最强的通道;

从上图不同行的激活注意力图可以看出,表现更好的网络的激活注意力图更加准确、细致地落在核心目标区域上。


image.png

(不同层级的注意力分布,以为例

再看看不同层级的注意力分布,可以看到浅层更关注细碎的信息,中间层则开始关注一些局部信息(比如眼睛),而到了深层则关注一整块区域的整体信息。


从上边两图可以看出,激活注意力图这一隐含特征是有指导意义的,因此可以设计出类似如下的模仿注意力的训练框架(注意注意力图的尺寸要保持一致,如果不一致可以考虑一些插值的手段)。

image.png


梯度注意力图

除了激活注意力图,梯度注意力图也有较强的指导意义,你可以把它看作是针对输入的一张敏感度图。引入梯度注意力图的知识蒸馏训练,需要多做一次反向传播。


第一次反向传播先采集教师网络、学生网络分别在网络输入上的梯度信息

然后度量梯度信息上的距离,来作为辅助损失计算到总损失中去

用新的损失进行反向传播并更新参数。


除此之外,还可以额外引入一个“翻转不变性”的约束,来作为一种正则化策略

其中是对输入进行水平翻转/垂直翻转的结果。


模仿思路

论文:《A Gift From Knowledge Distillation: Fast Optimization, Network Minimization and Transfer Learning (CVPR2017)


除了模仿某个中间位置的隐含特征之外,还可以建立跨层的联系,模仿某一个部分的“解题思路”,取网络的某一部分(若干卷积层堆叠),用其输入和输出融合构造出矩阵作为“解题思路”,然后计算教师网络与学生网络之间“解题思路”的L2损失作为辅助损失


首先定义第层到第层之间的一个“解题思路”FSP(Flow of Solution Process)为

其中,

是第层的输出特征图;

是第层的输出特征图;

注意两张特征图的分辨率是一致的,如果不一致可以考虑插值;

那么可以得到一个FSP矩阵,它就代表了从第层到第层之间的一个“解题思路”。


通过建立学生网络和教师网络之间,隐含的“解题思路”的约束,就可以让学生网络在训练的过程中去模仿教师网络的思路。

image.png


模仿决策边界

论文:《Knowledge Distillation with Adversarial Samples Supporting Decision Boundary (AAAI2019)

通过自对抗攻击技术,将原始样本移动到教师网络的决策边界附近(将这种样本称为Boundary Supporting Sample, BSS),指导学生网络在决策边界附近输出相似的概率分布,从而让学生网络来模仿教师网络的决策边界。

image.png


生成对抗样本

image.png

对于一个原始样本(Base sample),我们希望通过分类器用梯度下降法迭代更新输入,从而取得对于各个类别的对抗样本,记对于第个类别在第次迭代之后产生的对抗样本为,取初始状态


为分类器输出的一系列分类概率,分别记原始类别、目标类别的分类概率为,我们希望将样本移动到决策边界附近,那么可以定义损失函数为


迭代的终止条件(满足其一即终止迭代):

  1. 成功搜索到对抗样本,即
  2. 在搜索某一类别的对抗样本时误入其他类别的决策范围,搜索失败,即
  3. 决策边界太远以至于长时间迭代后仍未找到对抗样本,搜索失败,即


包含BSS的KD

总体来说,将损失函数设计为

其中,

分别是第个原始样本、第个样本在第个类别的对抗样本、第真实标签;

是损失函数,比如交叉熵

是分类函数,比如softmax函数;

的原始分类损失;

在标签上的知识蒸馏损失,也就是Hinton那套;

上,教师网络与学生网络之间的预测差异;

是权衡损失重要程度的超参;

的采样概率(或者说权重),

其中,

是教师网络将第个原始样本预测为第个类别的概率;

是第个原始样本的真实标签


由于生成对抗样本需要付出不小的计算代价,所以要有选择的选取部分原始样本来生成对抗样本,选择规则如下:

假设某个训练的batch为,挑选其中个原始样本来生成对抗样本,

  1. 找出教师和学生网络都能正确预测的样本,得到样本集合
  2. 如果,那么对中所有样本生成对抗样本;
  3. 如果,那么挑选教师和学生预测置信度相差较远的个样本来生成对抗样本,可以用欧式距离来衡量


image.png

(教师网络为ResNet26 with ACC=92.55%)

image.png

(教师网络为ResNet32 with TOP1=48.04% and TOP5=73.22%,学生网络为ResNet8)

image.png

(教师网络为ResNet42 with TOP1=56.10% and TOP5=78.71%,学生网络为ResNet10)


决策边界的相似性指标

为了知道学生网络有没有好好模仿教师网络的决策边界,作者设计了两个相似性度量指标:

  1. 尺度相似性(Magnitude Similarity, MagSim)
  2. 角度相似性(Angle Similarity, AngSim)


image.png


自蒸馏

论文:《Be Your Own Teacher: Improve the Performance of Convolutional Neural Networks via Self Distillation (ICCV2019)

这篇论文有点意思,它不需要单独训练一个教师网络,而是用网络自己的深层结构来指导前层结构的学习。


image.png

如图所示,思路其实也不复杂,将网络分为几个Section(比如ResNet按Stage划分Section,ResBlock1、ResBlock1/2、ResBlock1/2/3、ResBlock1/2/3/4分别组成深浅不一的Section),然后用最深的一个Section(即完整的网络)来指导较浅的几个Section进行训练。


较浅的Section分别引出后经过一个BottleNeck和一个全连接层,然后直接得到一个独立的分类器,论文没有介绍具体的配置,但为了能够正常的做Hint Training,按道理BottleNeck需要做下采样和通道扩增。


损失函数包含三个部分,并由超参调节各部分权重,

  1. 各个分类器与真实标签的交叉熵
  2. 最深层分类器与较浅层分类器输出的KL散度
  3. 最深层Section与较浅层Section的输出特征图的L2距离


新增的BottleNeck、全连接、Softmax等只在训练中使用,训练完成后只保留原始的网络。更灵活的是,如果需要获得一个更小的网络,可以保留对应的BottleNeck、全连接、Softmax而把原始网络后续的Section去掉。


image.png

(ImageNet上的实验结果)


image.png

(ImageNet:ResNet50自蒸馏 vs. ResNet152蒸馏ResNet50)


从上图可以看到,ResNet50自蒸馏训练的时间仅比ResNet50普通训练长一点,远远比训练一个ResNet152然后用其蒸馏一个ResNet50快的多,此外准确率也有明显的提升。


作者在论文中也对自蒸馏的有效性做出了一些解释:

  1. 自蒸馏可以使模型收敛到具有泛化特性的平坦极小值
    image.png
    比较冗余的网络更容易收敛到比较平坦的极小值点,而平坦的极小值点的泛化能力比陡峭的极小值点更强
  2. 自蒸馏可以避免梯度消失问题
    image.png
    训练结束后,自蒸馏模型相比于普通模型,平均梯度的数值尺度更大,这在浅层结构上表现的更为明显
  3. 自蒸馏可以使深层结构提取到更具区分度的特征
    image.png


不过,论文在arxiv上发布已经一年过去了,作者声称要开源的代码至今都没有踪影,倒是找到一份上个月(2020.04)有人复现的简单结果——

image.png

https://github.com/luanyunteng/pytorch-be-your-own-teacher

(CIFAR100上的实验,原文声称的结果和复现结果,Baseline为77.09%)


复现的作者并不能得到论文里的实验结果(当然可能也正常,毕竟论文作者没有公开完整的训练参数,通过微调能达到论文里的实验结果也是完全有可能的),而且相差0.78%,差距还是比较大的。考虑到论文作者没有主动开源自己的代码,我觉得论文里所描述的提升还是要存疑的。


权重复用

论文:《Net2Net: Accelerating Learning via Knowledge Transfer (ICLR2016)


复用预训练网络的权重是CNN中非常常用的做法,我们经常会用在ImageNet上预训练过的分类网络来作为backbone,然后在VOC、COCO等目标检测或是语义分割的更高级任务上进行进一步的训练。


但复用权重要求模型结构是一致的,而该论文提出一种部分复用和映射的方法,从而完成浅窄网络往深宽网络的一个权重复用方案;这其实也不算知识蒸馏,因为是往大网络扩展,不大可能用小网络来指导大网络的学习,就跟论文的标题一样,它其实是一种知识的迁移,甚至可以按照实际的业务需求扩大模型容量来做进一步的学习,达成一种终身学习(Lifelong Learning)的训练模式。


这种权重复用的训练下,可以达到从头训练相同的表现,但训练速度明显加快(图表就不贴啦,感兴趣的话去看原文)


宽度扩展

image.png

考虑一个全连接网络,为了给隐含层增加一个神经元,我们可以从已有神经元中随机复制权重(比如图中h[3]就复制了h[2]的权重c和d),然后对复制和被复制的神经元的输出进行简单的缩放(实际是上对下一层神经元权重的一个缩放),来控制下一层神经元的数值尺度。

用数学语言可以描述为


推广到卷积网络,则是通过增加卷积层的滤波器来增加宽度,滤波器的权重也是从已有滤波器中随机复制,然后为后继卷积层插入相应的卷积核并做出恰当的缩放操作来控制数值尺度。(有时间的话可以画张图放在这)


深度扩展

image.png

深度扩展则比较粗暴。

考虑一个全链接网络,可以加入一层新的全连接层(其宽度必须是上一层保持一致),其权重初始化为单位矩阵,也就是说将上一层各个神经元的输出作为本层对应神经元的输出。

扩展到卷积网络,则将新增的卷积层卷积核都初始化为中心为1,其他为0的形式即可。


KD目标检测

前述的知识蒸馏都是在分类任务上进行研究的,其实多数知识蒸馏的研究都是如此,但也不乏有些面向高级任务(如检测、分割)的论文,考虑到这些任务通常更有实际意义,他们在工程价值上可能会高些。


二阶段目标检测

论文:《Learning Efficient Object Detection Models with Knowledge Distillation (NIPS2017)


以Faster-RCNN为例,网络分为RPN和RCN两部分


image.png


其中,

分别是RCN和RPN的batch size;

是人为设置的权重


对于分类损失,

其中,

是交叉熵,而是加权交叉熵

负样本取正样本取以平衡正负样本(嗯?论文是不是说反了)

是人为设置的权重。


对于回归损失,

其中,

是人为设置的间隔,只有学生网络回归误差明显比教师网络大时,才接受来自教师网络的指导;

是人为设置的权重,论文中默认取


对于隐含层损失,直接使用L2损失或L1损失

适应层直接用一个1x1卷积来完成,并且论文中强调了适应层的重要性,直接比较隐含特征效果并不太好。


image.png

(不同backbone下的实验结果)


image.png

(用高分辨率输入的网络指导小分辨率输入的网络)


一阶段目标检测

论文:《Distilling Object Detectors With Fine-Grained Feature Imitation (CVPR2019)

image.png

作者并不是直接模仿目标检测的输出,而是模仿其中的隐含特征,但这个特征非常接近输出,也就是分类/回归分支前的最后一个特征图。同时也模仿FitNets的Hint Training,安排了一个简单的适应层进行Feature Adaptation,目的是为了将学生网络的特征图shape映射到跟教师网络相匹配的shape上,方便后续的距离计算,并且通过实验指出即使shape相同也不建议直接计算距离(效果变差),留下一个适应层进行缓冲会比较好。


考虑到目标检测中存在大量的负样本,让学生网络学习去模仿负样本的分类似乎不大行,所以我们需要对样本做一下筛选。

  1. 用真实框跟特征图上所有锚框计算IoU,得到一张的IoU图,其中分别是特征图的宽、高和每组分配的锚框数量;
  2. 用阈值去筛选空间上的锚框,得到一张掩膜

    其中,是手工设置的超参,通常取
    这里的思路跟裁剪的sensitivity有点像,因为IoU的变化范围比较大,手工挑选一个合适的常数阈值比较困难,所以通过设置一个对IoU尺度敏感的阈值在调参上会比较方便


image.png

(锚框筛选示意)

上图中,(a)是真实框,(b)(c)(d)分别是的筛选锚框,(e)(f)分别是的筛选锚框。


记学生网络准备去模仿的特征图为,教师网络中被模仿的特征图为,学生网络新增的适应层为,那么我们可以用L2距离构造出蒸馏的辅助损失

其中,,也即学生网络中“响应”的锚框组数。

与主损失加权得到总的一个损失


image.png

(实验结果,其中res101h代表在res101的基础上网络宽度减半,*-I表示运用了本文提出的方法)


KD裁剪、KD量化挖坑待填


实际应用问题