A Comprehensive Overhaul of Feature Distillation
我们研究了实现网络压缩的特征蒸馏方法的设计方面,并提出了一种新的特征蒸馏方法,其中蒸馏损失的设计是为了使教师变换、学生变换、蒸馏特征位置和距离函数等各个方面产生协同作用。我们提出的蒸馏损失包括一个带有新设计的余量ReLU的特征变换、一个新的蒸馏特征位置和一个部分L2距离函数,以跳过对学生的压缩有不利影响的冗余信息。在ImageNet中,我们提出的方法与ResNet50的top-1误差达到21.65%,超过了教师网络ResNet152的性能。我们提出的方法在图像分类、物体检测和语义分割等各种任务中进行了评估,并在所有任务中取得了显著的性能改进。The code is available at bhheo.github.io/overhaul
1. Introduction
在许多使用神经网络的机器学习任务中经历了显著的进步,研究人员已经开始致力于网络压缩和增强。人们提出了几种方法,如模型修剪、模型量化和知识蒸馏,以使模型更小,成本效益更高。其中,知识蒸馏法正在被积极研究。知识蒸馏指的是在一个较大的网络(教师)的监督下,帮助较小的网络(学生)的训练过程的方法。与其他压缩方法不同,它可以缩小网络的规模,而不考虑教师和学生网络之间的结构差异。由于允许架构上的灵活性,知识提炼正在成为下一代网络压缩方法的雏形。
Hinton等人[8]提出了一种使用教师网络的softmax输出的知识提炼(KD)方法。由于两个输出的尺寸相同,这种方法可以应用于任何一对网络结构。然而,高性能教师网络的输出与ground truth没有明显区别。因此,只转移输出类似于用ground truth训练学生,使得输出提炼的性能受到限制。为了更好地利用教师网络中包含的信息,人们提出了几种方法来代替输出蒸馏的特征蒸馏。FitNets[22]提出了一种方法,鼓励学生网络模仿教师网络的隐藏特征值。尽管特征提炼是一种很有前途的方法,但FitNets的性能改进并不明显。
在FitNets之后,人们提出了如下的特征提炼的变体方法。[30, 28]中的方法将特征转化为具有降低维度的表示,并将其转移给学生。尽管降低了维度,但据报道,抽象的特征表示确实导致了性能的提高。最近提出了一些方法(FT[13],AB[7])来增加蒸馏中的传输信息量。FT[13]使用自动编码器将特征编码为一个 “因子”,以减轻信息的泄漏。
AB[7]侧重于激活一个网络,只转移特征的标志(the sign of features)。这两种方法通过增加转移的信息量显示出更好的提炼性能。然而,FT[13]和AB[7]对教师的特征值进行了变形,这给性能的提高留下了进一步的空间。
在本文中,我们提出了一种新的特征蒸馏损失,进一步提高了特征蒸馏的性能,这种损失是通过对教师transform、学生transform、蒸馏特征位置和距离函数等各种设计方面的调查而设计的。我们的方法旨在从特征中转移两个因素。第一个目标是ReLU之后的特征响应的大小,因为它承载了大部分的特征信息。第二个是每个神经元的激活状态。最近的研究[20, 7]表明,神经元的激活状态强烈地代表了一个网络的表现力,在提炼时应该考虑它。
图1. 蒸馏方法的性能。AT [30], FT [13], AB [7]和提议的方法在ImageNet上的表现。图中显示了用每种蒸馏方法训练的ResNet50的准确率(%)。请注意,ResNet152的准确率为78.31%,是作为教师使用的。
为此,我们提出了一个margin ReLU函数,将蒸馏特征位置改为ReLU的前面(change the distillation feature position to the front of ReLU),并使用部分L2距离函数来跳过不必要的信息的蒸馏。所提出的损失明显改善了特征蒸馏的性能。在我们的实验中,我们在不同的领域评估了我们提出的方法,包括分类(CIFAR[15],ImageNet[23]),对象检测(PASCAL VOC[2])和语义分割(PASCAL VOC)。如图1所示,在我们的实验中,所提出的方法显示出优于现有的最先进的方法甚至是教师模型的性能。
2. Motivation
图2. 特征提炼的一般训练方案。教师变换Tt、学生变换Ts和距离d的形式因方法不同而不同。
在本节中,我们研究了实现网络压缩的特征蒸馏方法的设计方面,并介绍了我们的方法与前述方法不同的新方面。首先,我们描述了特征蒸馏法中损失函数的一般形式。如图2所示,教师网络的特征表示为Ft,学生网络的特征为Fs。为了匹配特征维度,分别为Tt和Ts,我们对特征Ft和Fs进行变换。变换后的特征之间的距离d被用作损失函数 L d i s t i l l L_{distill} Ldistill。换句话说,特征提炼的损失函数被概括为
学生网络是通过最小化蒸馏损失 L d i s t i l l L_{distill} Ldistill来训练的。
理想的做法是设计蒸馏损失,以便在不遗漏教师任何重要信息的情况下转移所有特征信息。为了实现这一目标,我们旨在设计一种新的特征蒸馏损失,在这种损失中,所有重要的教师信息都被尽可能地转移,以提高蒸馏性能。为了达到这个目的ÿ