特征蒸馏方法实现网络模型压缩 A Comprehensive Overhaul of Feature Distillation
论文题目:
A Comprehensive Overhaul of Feature Distillation
https://openaccess.thecvf.com/content_ICCV_2019/papers/Heo_A_Comprehensive_Overhaul_of_Feature_Distillation_ICCV_2019_paper.pdf
论文代码:
https://github.com/clovaai/overhaul-distillation
摘要
为了实现网络模型压缩,本文提出了一个新的特征蒸馏方法,该方法设计了新的蒸馏损失函数,可以使得教师网络转移(teacher transform),学生网络转移(student transform),特征蒸馏位置(distillation feature position)以及距离函数(distance function)协同作用。具体来说,蒸馏损失包括一个新设计的边缘ReLU的特征变换,一个新的特征蒸馏位置,和一个部分L2距离函数来跳过冗余信息,防止对学生网络的压缩产生不利影响。在ImageNet数据集上,基于ResNet50,本文提出的方法top-1 error为21.65%,超越了教师网络ResNet152的性能。实验证明,本文提出的方法在诸如图像分类,目标检测,语义分割等视觉任务上都有提升。
1. 引言
深度神经网络已经取得了很大的进步,近几年对网络模型压缩和增强的研究越来越多,诸如模型剪枝,模型量化,知识蒸馏等方法都被证明可以构建一个更小,成本更低的网络模型。在这其中,知识蒸馏的研究是一个热点,知识蒸馏是指在一个大的,准确率高的教师网络的监督下,训练一个轻量级的学生网络并期望学生网络拥有教师网络的能力。与其他压缩方法不同的是,知识蒸馏不需要对教师网络或者学生网络的结构进行限制,这为网络的架构设计带来了很大的便利性,是网络模型压缩的一大发展方向。
知识蒸馏起源于机器学习大佬Hinton的一篇论文《Distilling the knowledge in a neural network》,在这篇论文中,首次提到了知识蒸馏(Knowledge distillation KD)的概念,将蒸馏应用到教师网络经过softmax之后的输出上。这种方法具有很强的适用性,只要教师网络和学生网络的输出维度相同即可。但是,性能强大的教师网络的输出与ground-truth的差别并不大,因此,单纯的利用教师网络的输出作为学生网络的监督与直接利用ground-truth对学生网络进行训练的效果差别不大,使得输出的蒸馏表现受限。为了更好地利用教师网络中包含的丰富信息,很多方法开始着眼于对特征进行蒸馏而不是对输出蒸馏。FitNets是其中的代表方法,鼓励学生网络学习模仿教师网络的隐藏特征值,但是这种方法效果提升并不明显。在FitNets之后,还有很多针对特征蒸馏的方法提出,如将教师网络特征转换为降维的表示形式并将其传递给学生网络,使用自动编码器将功能编码成“因子”,以减少信息泄漏等,上述方法都提升了性能,但是会造成教师网络的特征值发生形变,留下了进一步提升的空间。
在本文中,作者还是聚焦在提升特征蒸馏的表现上,通过对教师网络迁移,学生网络迁移,特征蒸馏位置,距离度量方式等的协同考虑,提出了一种新的特征蒸馏损失函数。本文方法旨在从教师网络的特征中转移两个因素:第一个是经过ReLU激活之后的特征响应的大小,因为它包含了大部分的特征信息,第二个是每个神经元的激活状态。近来的研究表明,神经元的激活状态与网络的表达具有非常强的关联系性,应该被纳入到蒸馏的范围中来。因此,本文提出了一个margin ReLU激活函数,并且利用一个局部的L2正则化进行距离度量,以此来跳过对非必要信息的蒸馏。新提出的蒸馏损失函数显著提升了特征蒸馏表现,如下图,在ImageNet的实验上表明,本文的方法超越了SOTA甚至高于教师网络的表现。
2. 动机
首先回归一下常规特征蒸馏的过程。如下图, 表示教师网络的特征, 表示学生网络的特征, 和 分别代表教师网络和学生网络的特征转移形式, 表示转化后特征的距离,被用来作为蒸馏的损失函数,用公式可以表示为:
通过最小化上述蒸馏损失来训练学生网络。
蒸馏损失函数的设计十分关键,必须不能遗漏教师网络的所有重要特征信息,为了进一步提升蒸馏表现,本文设计了全新的特征蒸馏损失,确保教师网络的所有重要特征都尽可能的进行迁移。主要从四个方面进行考虑:
教师网络转换(Teacher transform)
教师网络的转换 将隐藏特征转换成容易迁移的特征形式。这是特征蒸馏中非常重要的部分也是蒸馏过程中特征信息损失的主要原因。此前提出的一些方法,将教师网络的高维特征信息进行降维处理,造成了严重的信息损失,还有的方法引入了压缩率或者以二值化的形式表征原始特征,造成特征信息的变形,与原始信息有很大的差异。由于教师网络特征既包括不利的信息,也包括有益的信息,因此区分它们并避免遗漏有益的信息是很重要的。
本文利用了一个新的ReLU激活函数,margin ReLU,通过这个激活函数,有益的特征信息不经过任何变换直接保留,而无用的特征信息被抑制,通过这样的设置,能够最大限度的保留有益的特征信息。
学生网络转换(Student transform)
通常来说,学生网络的转换 利用与 相同的形式,因此会造成与 类似的不足即有益特征信息的丢失。由于教师网络的特征维度通常较高,因此学生网络通常需要通过1x1卷积进行升维来进行维度匹配,在这种情况下,学生网络的特征维度不仅不会降低还会增加,这种教师网络和学生网络间非对称式的转换方式更有利于有益信息保存。
特征蒸馏位置(Distillation feature position)
除了转换方式之外,在网络结构中的哪些位置进行蒸馏同样需要细心设计。在FitNets上用任意中间层的末端作为蒸馏点,提升有限。这里需要跟margin ReLU一起考虑,在margin ReLU中,有益特征信息被保留而无用特征信息被抑制,因此,蒸馏必须在激活函数之前进行,在本文中,作者设计了在ReLU之前的蒸馏损失函数,称为pre-ReLU。所有的特征值和位置信息都会被保存在pre-ReLU中,不会发生改变。
距离度量函数(Distance function)
大多数的蒸馏方法采用L1或者L2函数进行距离度量。但是在本文提出的方法中,需要根据新提出的教师网络转换方式,蒸馏位置等设计引入一个适合的距离度量。所有特征值的转移都可能对学生网络产生不利影响。为了解决这个问题,提出了一个新的距离函数,称为局部L2距离,它被设计用来跳过在无用区域上提取到的特征信息。
3. 方法
蒸馏位置
在神经网络中,激活函数能够其非线性表达能力,是神经网络成功的关键一环。在各式各样的激活函数类型中,ReLU及其各种变种在神经网络中有广泛应用。ReLU对于正值是一个线性映射的关系,对于负值采取清零的方式。这种对正负值的区分处理,有助于防止无用信息的进一步传播。ReLU在知识蒸馏中也有用武之地,作者定义了网络的最小单元,比如ResNet中的残差块,VGG中的Conv-ReLU结构,称为一个block,常规的蒸馏位置通常是在每一个block的后面。在本文中,蒸馏的位置位于第一个ReLU和每一个block之间,这样能让学生网络达到保存教师网络特征信息的目的。如下图所示,对于普通神经网络的block和残差网络的block,本文提出的方法与其他方法的不同之处在于,蒸馏发生在ReLU之前还是之后,而与那些使用了预激活的网络进行对比,两者之间的差异更大,因为预激活的网络在每一个block后没有ReLU,因此本文的方法需要在下一个block中寻找ReLU。虽然蒸馏位置根据体系结构的不同可能会比较复杂,但是实验证明它对性能有很大的影响。能显著提高学生网络的表现。
损失函数
上文已经解释了关于教师网络转换 ,学生网络转换 和距离度量 。由于教师网络的特征值 是ReLU之前的值,正的值是教师网络利用的信息,而负的值则不是。如果教师网络的特征值是正的,学生网络必须产生与教师网络相同的特征值。反之,如果教师网络的特征值为负,则学生网络产生的特征值应小于零,使神经元处于相同的激活状态。作者提出了margin-ReLU,具体形式为:
其中,m是小于0的边缘值(margin value),通过引入margin-ReLU,教师网络的转换可以保留有益的正的特征值。下图是margin-ReLU与常规ReLU的对比
可以看到,边际值m被定义为负响应的逐通道期望值,边际ReLU使用与输入的每个通道对应的值,公式表示为:
期望值可以在训练过程中直接计算,也可以使用前BN层的参数进行计算。margin-ReLU被用作教师网络转换 ,通过该操作产生学生网络的目标特征值。而学生网络的转换 则包括1x1卷积层和BN层。
现在再来看距离度量d。在教师网络的特征中,正响应实际上是用于网络训练的,这意味着教师网络的正响应应该以其确切的值来传递。然而,负响应却不是。假定教师网络此时输出一个负响应,如果此时学生网络输出的响应值高于目标值,那么必须要减小,但是如果此时学生网络的输出值低于目标值,就没有必要再增加,因为响应值在经过激活函数后会被抑制。因此,定义一个局部的L2函数,公式如下:
其中,T是教师网络的特征位置,S是学生网络的特征位置。
综上,设 表示margin-ReLU作为教师网络的转移操作, 表示1x1卷积和BN层作为学生网络的转移操作, 表示局部的L2正则化作为距离度量,总的蒸馏损失函数可以表示为:
有了蒸馏损失,总损失需要加上任务自身的损失,因此,总损失函数为:
BN层
最后看一下BN层在知识蒸馏中的作用。BN层在神经网络中有广泛应用,能够让训练更稳定。在蒸馏网络中,通常,学生网络的特征是逐批进行归一化的。因此,教师网络的特征也必须以同样的方式归一化。也就是说,教师网络应该在训练时进行批处理归一化。为此,作者在1x1的卷积层之后附加一个BN层,在学生网络的设置上也是如此。
4. 实验
CIFAR-100上的实验结果
为了验证知识蒸馏的普适性,作者在几个不同网络结构的教师网络和学生网络上进行训练,下图是不同网络结构的具体参数。可以看到,实验主要在WideResNet中进行,这是因为WideResNet的层数和深度易于修改。
实验结果如下表所示,可以看到本文提出的方法超越了SOTA,在(a)组的架构中,经过本文提出的方法进行蒸馏的网络,表现甚至超越了教师网络。而在(e)z组的架构中,错误率为17.8%,这一数据比所有基于WideResNet的网络都好,证明本文方法不仅适用于小网络,对于大网络同样有很好的表现。
ImageNet上的实验结果
CIFAR数据集上的输入图像尺寸为32x32,在现实世界中没有代表性,因此,本文在ImageNet中进行实验,该数据集的平均尺寸是469x387,能够用来评估知识蒸馏在大图像上的效果。整个实验在两组网络中进行。第一组是将ResNet152蒸馏得到ResNet50,第二组是用ResNet50蒸馏MobileNet。实验结果表明,本文提出的方法取得了显著提升,尤其是在第一组实验中,学生网络的表现超过了教师网络。
教师网络和学生网络输出相似性实验
知识蒸馏利用教师网络的监督,让学生网络的输出尽可能与教师网络相近。输出蒸馏的目的很直观,即如果一个学生网络的输出与教师网络的输出相似,那么他的表现也会相似。然而,在特征蒸馏的情况下,有必要研究学生网络的输出是如何变化的。为了观察学生网络模仿教师网络的能力,作者测量了在一致的设置下教师网络和学生网络的输出的相似性。在CIFAR100验证集上,测量了学生网络和教师网络的KL散度以及与Ground-Truth的交叉熵,结果如下所示,可以看到,与其他网络相比,本文提出的方法显著降低了散度值,表示学生网络与教师网络的输出非常相似。
5. 结论
本文提出了一种新的知识蒸馏方法,对现有特征蒸馏方法的各个方面进行了综合研究,提出了一个新的蒸馏损失函数对特征蒸馏进行监督。蒸馏损失函数包括教师网络转换操作(利用margin ReLU实现),学生网络转换操作(利用1x1卷积和BN层实现),距离度量方式(利用局部L2函数实现)。通过实验证明,在不同计算机视觉任务中都取得了很好的提升,超越了SOTA,证明了有效性。
点一下 在看 会更好看哦