多任务下的特征分布蒸馏算法研究

葛海波，周 婷，黄朝锋，李 强

西安邮电大学 电子工程学院，西安 710000

近年来，卷积神经网络（convolutional neural network，CNN[1]引起了业内学者广泛的关注并得到了快速的发展，在图像分类[2，4]、目标检测[5，7]等任务中的表现取得了巨大的进步，但高性能的深度学习网络常常是计算密集型和参数密集型的，这一特点限制了网络在低资源设备上的应用。因此，各种模型压缩技术应运而生[8-9]。

知识蒸馏（KD）[10]因其在压缩率和精度保留方面的优异表现引起研究人员们的关注，继而诞生了FitNets算法使学生网络尽量学习教师网络的隐藏特征值[11]，该算法在性能方面的表现并不突出，但提供了特征提取这个很有价值的蒸馏思想，为知识蒸馏的发展打开了新思路。文献[12]引入注意力机制（AT），利用模型的注意力图指导训练，使简单模型生成的注意力图与复杂模型相似，以此来提升学生模型的学习效果。文献[7]提出包含特征交互和特征融合两个模块的注意力特征交互蒸馏（AFID），该方法利用交互教学机制完成两个网络之间的发送、接收和传递反馈，促进学生网络的学习效果。文献[13]在多教师模型的基础上，利用多组策略传递各个教师网络的中间特征（AMTML-KD），使学生网络自适应地学习不同等级的特征知识。文献[14]使用对抗训练的方式在线互相学习特征图的分布，同时结合logit蒸馏进一步提高在分类任务上的准确率（AFD）。

上述蒸馏算法在图像分类任务上的精度和压缩效果都有很好的表现，但让学生网络直接学习教师网络的中间层特征信息并不能很好地训练学生网络的拟合能力，导致其在目标检测、文本检索和语义分割等任务上的表现并不突出且算法并不支持在不同架构、不同维度的网络之间转移知识，对知识蒸馏技术的发展和应用带来一定的局限性。

针对上述问题，本文提出一种适用于多任务的基于特征分布蒸馏算法，主要有以下几点工作：

（1）引入互信息对输入图像的特征分布建模，利用条件概率分布代替互信息中的联合概率密度分布，使得到的特征分布结果描述更准确。

（2）在知识蒸馏中将教师网络的特征分布作为学生网络的学习知识，并在损失函数的设计中引入最大平均差异（MMD），以最小化教师网络和学生网络特征分布之间的距离，使学生网络更好地拟合教师网络的特征分布情况。

（3）在知识蒸馏的基础上使用toeplitz 矩阵实现学生网络的全连接层权重参数共享，进一步实现模型的压缩与加速。

（4）在图像分类、目标检测和语义分割等不同的任务上进行了实验，其在CIFAR-100、ImageNet 和VOC 数据集上都取得了很好的蒸馏效果。

1 特征分布知识蒸馏

以FitNets 蒸馏为代表的特征蒸馏算法中，通常是将每个位置空间的激活作为一个特征，每个滤波器中的二维激活映射作为神经元表示某个空间样本，仅仅反映了CNN是如何解释该样本的，忽略了CNN在处理该样本时更关注哪部分以及更强调使用哪种类型的激活模式。因此，直接匹配教师网络的特征图并不能很好地锻炼学生网络的特征拟合能力，因为它忽略了空间中样本的特征分布情况。

从上述问题出发，本文提出基于特征分布蒸馏算法，通过条件互信息分别对教师网络和学生网络的特征分布进行建模，利用MMD 度量两者之间的距离，使学生网络更好地拟合教师网络的特征分布情况，结合MMD匹配损失和目标蒸馏损失函数来设计本文算法的损失函数。算法原理图如图1所示，其中灰色框内表示算法损失函数，两个正方形框分别表示教师与学生网络在经过最小化MMD匹配前后的特征分布状态，图中的三角形和圆圈分别对应教师网络与学生网络的特征分布，对比两种情况下的特征分布情况，可以直观地看出经过特征分布蒸馏后的学生网络的特征分布与教师网络的特征分布更为相近。因此，从理论上分析，特征分布蒸馏算法可以使学生网络更好地拟合教师网络的特征分布，从而提高学生网络的特征提取能力。

图1 特征分布知识蒸馏原理图Fig.1 Principle diagram of characteristics of distributed knowledge distillation

1.1 目标蒸馏

知识蒸馏与迁移学习的思想较为相似，但知识蒸馏更加强调知识的迁移而非权重的迁移。目标知识蒸馏主要侧重于学习大型网络的输出知识，如逻辑单元和类概率，类概率是由Softmax激活函数转化的逻辑单元，其公式如式（1）所示：

其中，Zi表示第i类的逻辑单元值，Pi是第i类的类概率，k表示不同种类的数量，T表示温度系数，当T=1时则表示输出的类概率，当T=∞时，其为输出的逻辑单元[10]。

目标蒸馏时的损失函数如式（2）所示：

其中，N代表小批量的数量，LCE表示交叉熵，δ表示Softmax函数，yi表示样本i的真实标签，ZS∈RC，ZT∈RC分别为C 类任务上两个网络的Softmax 输出。该损失函数主要分为两部分：当T=∞时，学生网络和教师网络分类预测的交叉熵；当T=1 时，真实标签与学生网络的分类预测交叉熵。

1.2 特征分布建模

1.2.1 预备知识

知识蒸馏是将两个模型之间传递的知识表示为集合T，T={T1,T2,…,TN}，称该集合为转移集。设x=f(t)，y=g(t,W)，x、y分别用来表示教师模型和学生模型的输出，其中W为学生模型的参数。在知识蒸馏过程中，g(·)中的参数W用来学习并模拟f(·)的表示。

互信息是利用两个随机数据点间的相互关系描述特征空间的几何形状，同时也表示两个随机变量间相互依赖的程度，决定着联合分布P(x,y)和分解的边缘分布乘积P(x)P(y) 之间相似程度。设随机变量为X、Y，两者之间的联合概率密度可表示为ρ(x,y)=P(X=x,Y=y)，则X与Y之间的互信息I(X,Y)的计算公式如式（3）所示：

其中，P(x,y)表示随机变量X、Y的联合概率密度，P(x)和P(y)分别表示X和Y的边缘概率密度。

1.2.2 条件互信息

本文采用互信息对教师网络和学生网络的特征分布情况进行建模，为更好地拟合特征分布情况，互信息中利用核密度估计近似联合密度概率，教师网络和学生网络中的联合概率密度分别表示为式（4）和（5）所示：

其中，K(a,b;2δ2t)表示一个宽度为δt的对称核，a、b分别表示向量。

虽然使用联合概率密度分布来建模数据的几何分布并进行知识蒸馏可以解决传统知识蒸馏存在的许多问题，但考虑到样本局部区域内的数据存在相似性这一特点，为更准确地描述样本局部区域，本文提出条件互信息的概念。利用条件概率分布代替原有的联合概率密度函数，通过条件互信息对特征分布情况进行建模，教师模型和学生模型中的条件概率分布可分别表示为式（6）、（7）：

为验证条件互信息的有效性，分别对基于条件互信息和传统互信息的特征提取做了对比实验，结果如图2所示，其中图2（a）为输入原始图像，图2（b）为基于传统互信息的特征提取，图2（c）为基于条件互信息的特征提取。图2（b）、（c）中高亮的部分代表中心特征，可以看到经过条件互信息处理后的模型对样本中局部区域的关注度更高，更能捕捉到中心特征附近与之相似的特征信息。

图2 条件互信息与传统互信息的特征提取对比图Fig.2 Comparison of feature extraction between conditional mutual information and traditional mutual information

由于经特征分布蒸馏算法训练后的学生网络在提取特征时会更聚焦于样本的局部区域，更容易捕捉到样本的特征区间，这一特点在图像处理任务上都有很大的优势。因此，在第3章的实验中本文将该算法在图像分类、目标检测和语义分割任务上进行了实验，验证该算法在处理不同任务时的性能表现。同时，由于特征分布蒸馏算法是通过对教师网络的特征分布进行建模并以此为蒸馏的知识指导学生网络的学习，而不是让学生网络学习教师网络的中间层特征图。因此，特征分布算法并不受限于只能在相同网络架构间进行蒸馏，对此本文在第3章中设置了实验进行验证。

1.3 损失函数的设计

其中，k(·,·)是核函数，将样本向量映射到高阶特征空间中。综上所述，最小化MMD相当于最小化α和β之间的距离，利用这一思想设计MMD匹配损失函数。

由于本文所提的知识蒸馏算法是利用教师网络的特征分布和Softmax 层输出概率来指导学生网络的学习。因此，损失函数的设计可分为以下两部分：MMD匹配损失函数以及目标蒸馏的交叉熵损失，则损失函数的定义为：

其中，H(·)为标准的交叉熵损失，ytrue表示真标签，PS为学生网络的输出概率，ω为正则项损失函数的权重。

2 参数共享

由于现有方法大多数侧重于蒸馏技术的提高，通常会忽略学生网络作为一个完整的神经网络本身具有的可压缩空间，因此本文在知识蒸馏的基础上结合参数共享方法。

参数共享主要是针对学生网络的全连接层，将全连接层的权重参数表示为m×n维的矩阵W，其中n、m分别表示输入层和输出层中神经元的数量，图3（a）表示输入层和输出层之间的神经元连接，图3（b）表示输入层与输出层间的权重参数矩阵。输入层第i个神经元与输出层神经元之间的连接可表示为一个n维向量，用ωi表示ωi=(ωi1,ωi2,…,ωin)，其中ωij表示连接输入层第i个神经元和输出层第j个神经元的权重参数。本文利用toeplitz矩阵实现网络全连接层的参数共享。toeplitz矩阵简称为T型矩阵，矩阵形式如公式（11）所示：

图3 输入层与输出层神经元连接和权重参数矩阵Fig.3 Input layer and output layer neurons connected and weighting parameter matrix

在toeplitz 矩阵中只有2n－1 个独立的元素。经过toeplitz 矩阵规划后的网络在输出层与输入层神经元之间的连接和权重参数矩阵如图4（a）、（b）所示。

图4 toeplitz矩阵规划后输入层与输出层连接和权重矩阵Fig.4 Connection between input layer and output layer and weight matrix after toeplitz matrix planning

利用toeplitz 矩阵完成全连接层的参数共享，网络反向传播过程中的梯度求取公式将会发生变化。全连接层权重参数可表示为行向量ωrow=(ω11,ω12,…,ω1n)和列向量ωcol=(ω11,ω21,…,ωm1)，可见行向量与列向量都有权重参数ω11。x为输入神经元，x=(x1,x2,…,xn)T，y为输出神经元y=(y1,y2,…,ym)T，E为模型误差，∂E∂y为误差经过反向传播返回到当前层的值，∂E∂y和y都是m维列向量，在对权重参数梯度进行求取时，将输入的n维x扩展为n+m-1维，具体是：当j ＞n时，xj=0 即x=(x1,x2,…,xn,0,0,0,…,0)T。行向量ωrow的梯度公式如式（12）～（15）所示：

列向量ωcol的梯度公式如式（16）～（19），ω11的梯度由公式（12）求得公式（16）：

其中，LSR(w,g)表示将向量ω逻辑右移g位，×代表两个矩阵相乘，T表示矩阵的转置。

3 实验验证

为证明本文所提特征分布蒸馏算法的有效性，首先在图像分类任务上对该算法进行了实验验证；其次，在目标检测和语义分割任务上进行了实验，验证该算法在不同图像处理任务上的优势；最后在每个图像处理任务中分别设置两组教师网络和学生网络的蒸馏实验，验证该算法在不同网络架构间的蒸馏效果。

3.1 图像分类

由于CIFAR-100 是在图像分类任务上运用最广泛的数据集，为保证与其他知识蒸馏算法的比较的公平性，在图像分类任务中选择CIFAR-100 数据集进行实验。实验是由PyTorch实现，最小批次设置为64，epochs设置为500，初始学习率设为0.1，然后分别在第200 轮次和300 轮次时将学习率调整为0.01 和0.001。蒸馏过程中的参数设置，logits蒸馏中软化系数T为4.0。

CIFAR-100数据集总共有60 000张尺寸为32×32的彩色图像，包含有100 个类，每一类有600 张图像，其中500张图像用于训练，100张用于测试，分别采用ResNet-152 和ResNet-50 作为教师模型和学生模型进行实验。首先设计消融实验分别验证算法中特征分布蒸馏和参数共享的有效性，其结果如表1 所示，可以看出经过特征分布蒸馏算法训练后的学生网络的分类准确率（Accuracy）有了明显的提升，甚至优于教师网络，在压缩率（Params）方面，学生网络的参数量远远少于教师网络。在知识蒸馏的基础上对模型进行参数共享处理，实验表明经过参数共享处理后的网络实现了在精度仅减少0.1%～0.2%（可接受范围内）的情况下参数量降低至教师网络的33.87%，节省了模型对存储空间的需求。同时将具有代表性的知识蒸馏算法KD、FitNets、AT、AFID、AMIMI-KD、AFD与本文所提的特征分布蒸馏算法在相同的参数和环境设置下做了对比实验，实验结果如表2所示。

表1 消融实验Table 1 Ablation experiment

表2 CIFAR-100数据集上不同蒸馏算法间的性能对比Table 2 Performance comparison of different distillation algorithms on CIFAR-100 dataset

由表2的实验结果可以看出，特征分布蒸馏算法在分类任务上的准确率达到了78.92%，相较于全监督教师网络提高了0.78%；相较于全监督学生网络提高了8.89%；相较于KD、FitNets、AT、AFID、AMIMI-KD、AFD算法分别提高了1.58%、6.9%、0.52%、1.92%、2.09%和0.85%。可以看出特征分布蒸馏算法在图像分类任务上具有很好的性能表现，同时该实验也证明了利用图像在教师网络中的特征分布情况来指导学生网络学习的效果明显优于利用特征图来指导学生网络学习的效果。

上述实验可以看出特征分布蒸馏算法优秀的分类准确率和高压缩比，显示了该算法的压缩潜力。为此，本文采用更大更复杂的数据集ImageNet进一步验证该算法的收敛速度和鲁棒性。ImageNet 数据集包含有120 万张训练图像和5 万张测试图像，共分为1 000 个类。与CIFAR-100数据集相比，ImageNet数据集具有更丰富的种类，且图像规模更大（平均469×387）。首先以ResNet-152和ResNet-50作为教师模型和学生模型测试学生网络在经过特征分布蒸馏算法前后的Kappa系数，Kappa 系数是在分类任务中评价分类准确率的一个重要评价指标，其值越高表示分类准确率越高，最高值为1.0。实验结果如图5所示，可以看到无论是否经过本文蒸馏算法进行蒸馏，Kappa 系数均呈现上升趋势，但ResNet-50未经特征分布蒸馏时的训练数据集Kappa系数始终低于经过特征分布蒸馏后的Kappa 系数。同时未经过特征分布蒸馏的预测训练集收敛速度也比较慢，且波动幅度也比较大，经特征分布蒸馏后的网络在预测训练集上的波动幅度小，并且在迭代次数为400次时便完成了收敛，开始趋于稳定。实验证明，在经过特征分布蒸馏后的网络加快了图像分类的收敛速度，提升了图像分类准确率。

图5 ResNet-50在特征分布蒸馏前后的Kappa系数Fig.5 Kappa coefficient of ResNet-50 before and after distillation of characteristic distribution

此外，本文采用Top-1 error和Top-5 error作为评价指标验证该算法在不同网络结构间的蒸馏能力，第一组仍然选用ResNet-152和ResNet-50作为教师模型和学生模型，第二组选用ResNet-50 作为教师网络，MobileNet作为学生网络。为保证实验的公平性，两个教师网络均提前在PyTorch 库中预先训练，实验结果如表3 所示。可以看出，本文的算法在ImageNet 数据集上仍有优秀的表现。

表3 ImageNet数据集上的图像分类性能表现Table 3 Image classification performance on ImageNet dataset

在相同网络架构间进行蒸馏时，特征分布蒸馏算法的Top-1 error 和Top-5 error 相较于全监督教师网络分别减少了0.6%和4.7%；相较于全监督学生网络分别减少了9.56%和20.45%。可以看出经特征分布蒸馏后的网络性能均优于全监督教师网络和全监督学生网络。

在不同的网络架构间进行蒸馏时，特征分布蒸馏算法的Top-1 error 和Top-5 error 相较于全监督教师网络是增加的，但相较于全监督学生网络而言分别减少了14.45%和20.6%。可以看出即使受不同网络架构的影响，特征分布蒸馏算法在不同网络架构间的蒸馏效果仍是优于全监督学生网络的。

3.2 目标检测

本文将特征分布蒸馏算法应用于目前流行的高速检测器SSD 上[15]，所有模型都是用VOC2007 训练集进行训练，其中主干网络使用ImageNet 数据集进行与训练，将没有经过蒸馏训练的SSD 作为基准，分别以ResNet50 和ResNet18 作为教师网络和学生网络，为验证本文所提算在目标检测任务上的性能优势，将特征分布算法与AT、AFID、FitNets 和AFD 算法在目标检测任务上的性能进行了比较，图6 表示几种算法分别在IOU=0.5 和IOU=0.7 时的PR 比曲线，从比较结果来看，经过特征分布蒸馏后的网络在目标检测任务上的准确率和召回率两个性能指标均优于其他蒸馏算法，这是因为特征分布蒸馏算法中利用条件互信息对样本的特征分布情况进行建模，使其在蒸馏过程中更加聚焦于样本的局部特征信息，对边界框的回归自然更加准确。

图6 网络在经过不同算法蒸馏后的PR曲线Fig.6 PR curve after distillation by different algorithms

同时分别以ResNet18 和MobileNet 为学生网络，平均精度均值（mAP）作为评价指标进行实验，比较了几种不同的算法在目标检测任务中对不同网络架构的蒸馏效果。结果如表4所示。

表4 VOC2007测试集上的目标检测性能比较Table 4 Target detection performance comparison on VOC2007 testset

实验验证了本文所提的特征分布蒸馏可以很好地应用于目标检测任务上，同时该算法在压缩率方面也有很好的表现，这是在其他对比算法中没有实现的。

在相同网络架构间蒸馏时，虽然特征分布蒸馏算法相较于全监督教师网络的平均精度值降低了3.3%，但模型压缩率达到了教师网络的44.1%；相较于全监督学生网络而言平均精度值提高了4.2%，同时特在此基础上压缩了模型的大小；相较于FitNets、AT、AFID、AFD几种对比算法的平均精度值分别提高了2.4%、2.1%、1.04%和1.85%。

在不同网络架构间蒸馏时，特征分布蒸馏算法的平均精度值相较于全监督学生网络提高了3.85%，且优于其他几种对比算法。同时，可以发现其他几种对比算法的平均精度值与全监督学生网络相比均是降低的，证明了特征分布蒸馏在不同网络架构间蒸馏的有效性。

3.3 语义分割

本文对语义分割任务进行了知识蒸馏，实验中以基于ResNet101 的DeepLabV3+作为教师网络，并以基于ResNet18和MobileNetV2的DeepLabV3+分别作为学生网络，以均交并比（mIoU）和像素准确率（PA）作为本文算法在语义分割任务上的性能指标，其实验结果如表5所示，可以看到本文所提的特征分布蒸馏算法大幅度提高了ResNet18 和MobileNet 模型的性能表现，同时该算法的压缩效果优于另外几种对比算法。

表5 VOC2012测试集上的语义分割性能比较Table 5 Semantic segmentation performance comparison on VOC2012 testset

在相同网络架构间进行蒸馏时，特征分布蒸馏算法的均交并比和像素准确率相较于全监督教师网络的分别提高了4.34%和0.5%，相较于全监督学生网络分别提高了14.01%和4.28%；而其他特征蒸馏算法的均交并比和像素准确率相较于教师网络均是减小的，相较于全监督学生网络，两种指标的增长幅度都很小。这是因为其他特征蒸馏算法中，采用教师网络的特征图指导学生网络的学习，这种方法只传递给学生网络一个特征提取的结果，并不能很好地锻炼学生网络的特征拟合能力。而本文的特征分布蒸馏算法指导学生网络拟合教师网络的特征分布情况，很好地锻炼了学生网络的特征拟合能力从而提高了学生网络的特征提取能力，使学生网络能更好地应用于不同的图像处理任务中。

在不同的网络架构间进行蒸馏时，特征分布蒸馏算法的均交并比和像素准确率相较于全监督教师网络是减少的，但是相较于全监督学生网络及其他几种对比蒸馏算法而言是提高的，说明特征分布蒸馏算法虽然突破了不同网络架构间蒸馏的限制，但还是受其影响的。

4 总结

本文提出的特征分布知识蒸馏通过MMD 匹配教师网络与学生网络之间的特征分布情况，不需要对任何超参数进行特定的调整，同时通过利用条件互信息对数据样本的建模，使学生网络拟合教师网络的特征分布，更好地训练了学生网络的特征拟合能力，使其可以很好地应用于多种图像处理任务中，并且可以实现在不同维度空间内传递教师网络的知识。其次在特征蒸馏的基础上结合参数共享算法实现对学生网络的压缩，进一步节省了模型的资源需求。本文所提的特征分布蒸馏算法在图像分类、目标检测和语义分割任务场景下的性能和压缩率均优于其他几种对比算法。后续研究中会进一步研究不同网络架构间的蒸馏，提高不同网络架构间的蒸馏效果。