1.介绍

​ 知识蒸馏,近年来越来越受到研究界的关注。大型深度神经网络取得了显著的成功,特别是在大规模数据的真实场景中,因为当考虑新数据时,过度参数化提高了泛化性能。然而,由于移动设备的算力和存储受限,在这些设备中部署DNN模型是一个巨大的挑战。为了解决这个问题,Bucilua等人(2006)首次提出了模型压缩,将信息从大模型或集成模型转移到训练的小模型中,并且不会显著导致精度下降。从大模型中学习小模型的方法后来作为知识蒸馏正式推广。

​ 在知识蒸馏中,小型学生模型通常由大型教师模型监督。其主要思想是:学生模型模仿教师模型,并获得近似甚至更优越的性能表现。关键问题是如何将大教师模型转移到小的学生模型上。知识蒸馏系统基本由三个关键部分组成:知识、蒸馏算法、教师-学生结构。其一般框架如下:

image-20221105214440974

2.知识

​ 将知识分为三种形式:基于响应的(response-based)、基于特征的(feature-based)、基于关系的(relation-based)。

image-20221106204723627

2.1 基于响应的知识(response-based)

​ 基于响应的知识通常指教师模型最后一个输出层的神经响应。其主要思想是直接模拟教师模型的最终预测。给定一个logit(即一般神经网络最后一层的输出,后接一个softmax可以转换为概率值)作为最后一个全连接层的输出,基于响应的知识蒸馏loss可以表示为:。其中,表示logits的之间的loss,分别表示教师和学习的logits。典型的基于响应的知识蒸馏模型如下图:

image-20221106210608819

​ 基于响应的知识可以用于不同类型的模型预测。例如,在目标检测任务中的响应可能包含bounding box的偏移量的logits;在人类姿态估计任务中,教师模型的响应可能包括每个地标的热力图。最流行的基于响应的图像分类知识被称为软目标(soft target)。软目标是输入的类别的概率,可以通过softmax函数估计为:

是第i个类别的logit,是温度因子,控制每个软目标的重要性。软目标包含来自教师模型的暗信息知识(informative dark knowledge)。因此,软logits的蒸馏loss可以重写为:

​ 通常, 使用KL散度loss。(Kullback-Leibler divergence,衡量两个概率分布的相似性的度量指标)。优化该等式可以使学习的logits与教师的logits相匹配。下图为基准知识蒸馏的具体架构。

image-20221107154250742

​ 然而,基于响应的知识通常需要依赖最后一层的输出,无法解决来自教师模型的中间层面的监督,而这对于使用非常深的神经网络进行表示学习非常重要。由于logits实际上是类别概率分布,因此基于响应的知识蒸馏限制在监督学习。

2.2 基于特征的知识(feature-based)

​ DNN善于随着抽象程度的增加学习到不同级别的特征表示。这被称为表示学习(representation learning)。对于最后一层的输出和中间层的输出(特征图,feature map),都可以作为监督学生模型训练的知识。来自中间层的基于特征的知识是基于响应的知识的一个很好的扩展,特别是对于训练窄而深的网络。

​ 一般,基于特征的知识蒸馏loss可以表示为:

​ 其中,分别是教师模型和学生模型的中间层的特征图。变换函数当教师和学生模型的特征图大小不同时应用。衡量两个特征图的相似性,常用的有L1 norm、L2 noem、交叉熵等。下图为基于特征的知识蒸馏模型的通常架构。

image-20221107160526275

image-20221107170426947

2.3 基于关系的知识(relation-based)

​ 基于响应和基于特征的知识都使用了教师模型中特定层的输出,基于关系的知识进一步探索了不同层或数据样本的关系。

​ 一般,将基于特征图关系的关系知识蒸馏loss表述如下:

其中,分别表示教师模型和学生模型的特征图。分别是教师模型和学生模型的特征图组(pairs)。表示特征组的相似性函数。
下图为实例关系的知识蒸馏:

image-20221107164251384

​ 传统的知识迁移方法往往涉及到单个知识蒸馏,教师的软目标被直接蒸馏给学生。实际上,蒸馏的知识不仅包含特征信息,还包含数据样本之间的相互关系。这一节的关系涉及许多方面,作者的设计很灵活,建议看原论文更清楚。 image-20221107170335515

3.蒸馏机制

根据教师模型是否与学生模型同时更新,知识蒸馏的学习方案可分为离线(offline)蒸馏、在线(online)蒸馏、自蒸馏(self-distillation)。

image-20221107171110931

3.1 离线蒸馏

​ 大多数之前的知识蒸馏方法都是离线的。最初的知识蒸馏中,知识从预训练的教师模型转移到学生模型中,整个训练过程包括两个阶段:1)大型教师模型蒸馏前在训练样本训练;2)教师模型以logits或中间特征的形式提取知识,将其在蒸馏过程中指导学生模型的训练。教师的结构是预定义的,很少关注教师模型的结构及其与学生模型的关系。因此,离线方法主要关注知识迁移的不同部分,包括知识设计、特征匹配或分布匹配的loss函数。离线方法的优点是简单、易于实现。

​ 离线蒸馏方法通常采用单向的知识迁移和两阶段的训练程序。然而,训练时间长的、复杂的、高容量教师模型却无法避免,而在教师模型的指导下,离线蒸馏中的学生模型的训练通常是有效的。此外,教师与学生之间的能力差距始终存在,而且学生往往对教师有极大依赖。

3.2 在线蒸馏

​ 为了克服离线蒸馏的局限性,提出了在线蒸馏来进一步提高学生模型的性能,特别是在没有大容量高性能教师模型的情况下。在线蒸馏时,教师模型和学生模型同步更新,而整个知识蒸馏框架都是端到端可训练的。

​ 在线蒸馏是一种具有高效并行计算的单阶段端到端训练方案。然而,现有的在线方法(如相互学习)通常无法解决在线环境中的高容量教师,这使进一步探索在线环境中教师和学生模式之间的关系成为一个有趣的话题。

3.3 自蒸馏

​ 在自蒸馏中,教师和学生模型使用相同的网络,这可以看作是在线蒸馏的一个特例。例如论文(Zhang, L., Song, J., Gao, A., Chen, J., Bao, C. & Ma, K. (2019b).Be your own teacher: Improve the performance of convolutional eural networks via self distillation. In ICCV.)将网络深层的知识蒸馏到浅层部分。

​ 从人类师生学习的角度也可以直观地理解离线、在线和自蒸馏。离线蒸馏是指知识渊博的教师教授学生知识;在线蒸馏是指教师和学生一起学习;自我蒸馏是指学生自己学习知识。而且,就像人类学习的方式一样,这三种蒸馏由于其自身的优点,可以相互补充

4.教师-学生架构

​ 在知识蒸馏中,教师-学生架构是形成知识迁移的通用载体。换句话说,从教师到学生的知识获取和蒸馏的质量取决于如何设计教师和学生的网络结构。基于人类学习的习惯,希望学生模型能找到合适的教师。近年来,在蒸馏中,教师和学生的模型设置几乎是预先固定的,大小和结构不变,容易产生模型容量上的差距。两者之间模型的设置主要有以下关系:

image-20221107194852541

​ Hinton等人(Hinton, G., Vinyals, O. & Dean, J. (2015). Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531.)曾将知识蒸馏设计为压缩深度神经网络的集合。深度神经网络的复杂性主要来自于深度和宽度这两个维度。通常需要将知识从更深、更宽的神经网络转移到更浅、更窄的神经网络。学生网络的结构通常有以下选择:1)教师网络的简化版本,层数更少,每一层的通道数更少;2)保留教师网络的结构,学生网络为其量化版本;3)具有高效基本运算的小型网络;4)具有优化全局网络结构的小网络;5)与教师网络的结构相同。

​ 大网络和小网络之间容量的差距会阻碍知识的迁移。为了有效地将知识迁移到学生网络中,许多方法被提出来控制降低模型的复杂性。论文(Mirzadeh, S. I., Farajtabar, M., Li, A. & Ghasemzadeh, H. (2020). Improved knowledge distillation via teacher assistant. In AAAI.)引入了一个教师助手,来减少教师模型和学生模型之间的训练差距。另一篇工作通过残差学习进一步缩小差距,利用辅助结构来学习残差。还有其他一些方法关注减小教师模型与学生模型结构上的差异。例如将量化与知识蒸馏相结合,即学生模型是教师模型的量化版本。

​ 通过替换基础组件也是有效的方法,例如将普通卷积替换为在移动和嵌入式设备上更高效的深度可分离卷积。受神经架构搜索的启发,通过搜索基于基础构建块的全局结构,小型神经网络的性能进一步提高。此外,动态搜索知识迁移机制的想法也出现在知识蒸馏中,例如,使用强化学习以数据驱动的方式自动去除冗余层,并寻找给定教师网络的最优学生网络。

5.蒸馏算法

​ 一个简单但有效的知识迁移方法是直接匹配基于响应的、基于特征的或教师模型和学生模型之间的特征空间中的表示分布。这小节回顾用于知识迁移的典型蒸馏算法。

5.1 对抗性蒸馏(Adversarial Distillation)

​ 生成对抗网络(GAN)中的鉴别器估计样本来自训练数据分布的概率,而生成器试图使用生成的数据样本欺骗鉴别器的概率预测。受此启发,许多对抗性的知识蒸馏方法已经被提出,以使教师和学生网络能够更好地理解真实的数据分布。

image-20221107212712495

如上图,对抗性学习的蒸馏方法可以分为三类:

a)训练一个对抗性生成器生成合成的数据,将其直接作为训练集或用于增强训练集。

其loss可以表示为:

其中,分别是教师模型和学生模型的输出;表示给定随机输入向量z的生成器G生成的训练样本;是蒸馏损失,以迫使预测的概率分布与真实概率分布之间匹配,蒸馏损失函数通常采用交叉熵或KL散度。

b)使用鉴别器,利用logits或特征来分辨样本来自教师或是学生模型

代表性方法如论文(Wang, Y., Xu, C., Xu, C. & Tao, D. (2018f). Adversarial learning of portable student networks. In AAAI.),其loss可以表示为:

其中,G是一个学生网络,是生成对抗网络使用的损失函数,使学生和教师之间的输出尽可能相似。

c)在线方式进行,在每次迭代中,教师和学生共同进行优化

利用知识蒸馏压缩GAN,小GAN学生网络通过知识迁移模仿大GAN教师网络

从上述对抗性蒸馏方法中,可以得出三个主要结论:1)GAN是通过教师知识迁移来提高学生学习能力的有效工具;2)联合GAN和知识蒸馏可以为知识蒸馏的性能生成有价值的数据,克服了不可用和不可访问的数据的限制;3)知识蒸馏可以用来压缩GAN

5.2 多教师蒸馏(Multi-teacher Distillation)

​ 不同的教师架构可以为学生网络提供他们自己有用的知识。在训练一个教师网络期间,多个教师网络可以单独或整体地用于蒸馏。在一个典型的师生框架中,教师通常是一个大的模型或一个大的模型的集合。要迁移来自多个教师的知识,最简单的方法是使用来自所有教师的平均响应作为监督信息。

image-20221108143929293

​ 多个教师网络通常使用logits和特征表示作为知识。除了来自所有教师的平均logits,还有其他的变体。文献(Chen, X., Su, J., & Zhang, J. (2019b). A two-teacher tramework for knowledge distillation. In ISNN.)使用了两个教师网络,其中一名教师将基于响应的知识迁移给学生,另一名将基于特征的知识迁移给学生。文献(Fukuda, T., Suzuki, M., Kurata, G., Thomas, S., Cui, J. & Ramabhadran,B. (2017). Effificient knowledge distillation from an ensemble of teachers. In Interspeech.))在每次迭代中从教师网络池中随机选择一名教师。

5.3 跨模态蒸馏(Cross-Modal Distillation)

​ 在训练或测试时一些模态的数据或标签可能不可用,因此需要在不同模态间知识迁移。在教师模型预先训练的一种模态(如RGB图像)上,有大量注释良好的数据样本,(Gupta, S., Hoffman, J. & Malik, J. (2016). Cross modal distillation for supervision transfer. In CVPR.)将知识从教师模型迁移到学生模型,使用新的未标记输入模态,如深度图像(depth image)和光流(optical flow)。具体来说,所提出的方法依赖于涉及两种模态的未标记成对样本,即RGB和深度图像。然后将教师从RGB图像中获得的特征用于对学生的监督训练。成对样本背后的思想是通过成对样本迁移标注(annotation)或标签信息,并已广泛应用于跨模态应用。成对样本的示例还有1)在人类动作识别模型中,RGB视频和骨骼序列;2)在视觉问题回答方法中,将图像-问题-回答作为输入的三线性交互教师模型中的知识迁移到将图像-问题作为输入的双线性输入学生模型中。

跨模态蒸馏的框架如下:

image-20221108155101574

跨模态总结如下:

ResK表示基于响应的知识,FeaK表示基于特征的知识,RelK表示基于关系的知识

image-20221108155621495

5.4 基于图的蒸馏(Graph-Based Distillation)

​ 大多数知识蒸馏算法侧重于将个体实例知识从教师传递给学生,而最近提出了一些方法使用图来探索数据内关系。这些基于图的蒸馏方法的主要思想是1)使用图作为教师知识的载体;2)使用图来控制教师知识的传递。基于图的知识可以归类为基于关系的知识。

image-20221108160410882

1)使用图作为教师知识的载体

文献(Zhang, C. & Peng, Y. (2018). Better and faster: knowledge transfer from multiple self-supervised learning tasks via graph distillation for video classifification. In IJCAI)中,每个顶点表示一个自监督的教师,利用logits和中间特征构造两个图,将多个自监督的教师的知识转移给学校。

2)使用图来控制知识迁移

文献(Luo, Z., Hsieh, J. T., Jiang, L., Carlos Niebles, J.& Fei-Fei, L. (2018).Graph distillation for action detection with privileged modalities.In ECCV.)将模态差异纳入来自源领域的特权信息,特权信息。引入了一种有向图来探讨不同模态之间的关系。每个顶点表示一个模态,边表示一个模态和另一个模态之间的连接强度。

5.5 基于注意力的蒸馏(Attention-Based Distillation)

​ 由于注意力可以很好地反映卷积神经网络的神经元激活,因此利用注意力机制在知识蒸馏中提高学生网络的性能。注意力迁移的核心是定义特征嵌入到神经网络各层中的注意力映射。

5.6 无数据的蒸馏(Data-Free Distillation)

​ 无数据蒸馏的方法提出的背景是克服由于隐私性、合法性、安全性和保密问题导致的数据缺失。“data free”表明并没有训练数据,数据是新生成或综合产生的。新生的数据可以利用GAN来产生。合成数据通常由预先训练过的教师模型的特征表示生成

image-20221108164844505

5.7 量化蒸馏(Quantized Distillation)

​ 量化的目的是将高精度(如float32)表示的网络转换为低精度(如int8)以降低计算复杂度,而知识蒸馏的目的是训练一个小模型,以产生与复杂模型相当的性能。二者可以结合。

image-20221108165525385

5.8 终身蒸馏(Lifelong Distillation)

终身学习,包括持续学习、集成学习和元学习,旨在以类似于人类的方式学习,积累了之前学到的知识,并将学到的知识转移到未来的学习中

5.9 基于神经架构搜索的蒸馏(NAS-Based Distillation)

​ 神经架构搜索的目的是自动确认深度神经模型,并自适应地学习适当的神经结构。在知识蒸馏中,知识迁移的成功不仅取决于教师的知识,还取决于学生的体系结构。然而,在大教师模型和小学生模型之间可能存在着能力差距,这使得学生很难从老师那里学习得很好。为了解决这一问题,采用神经结构搜索寻找合适的学生结构。

6.总结

知识蒸馏的关键在于:1)在教师网络中学到什么样的知识(知识的类型);2)师生网络架构;3)如何迁移知识(蒸馏算法)

参考文献:

GOU J, YU B, MAYBANK S J, et al. Knowledge Distillation: A Survey[J]. International Journal of Computer Vision, 2021, 129(6): 1789–1819.