近些年深度神经网络改变了人类的生活, 无论是在工业、科研、日常生活中都有着很高的应用需求, 然而这种深度神经网络模型对设备的要求苛刻, 通常需要大量的计算量和内存, 解决这一问题目前比较受欢迎的方法是知识蒸馏(knowledge distillation, KD)^[1].

知识蒸馏作为一种模型压缩的技术, 通过将复杂模型(教师模型)的知识传递给简单模型(学生模型)来实现, 目的是在花费最小化代价的情况下尽量使其效果向复杂模型不断靠近, 以满足大多数设备的性能. 自vanilla-KD提出之后基于响应的(response-based)^[2]知识蒸馏走近人们的视野, 随后基于特征的(feature-based)^[3–7]、基于关系的(relation-based)^[8–10]的知识蒸馏也被相继提出, 目前受到大家关注的方法大多数都是基于中间特征层的, 这些方法忽略掉了logit层的语义信息更高的特点, 解耦知识蒸馏(decoupled knowledge distillation, DKD)^[11]的提出使得基于响应的知识蒸馏重新回到了SOTA行列.

解耦知识蒸馏(DKD)是一种基于响应的方法, 传统的知识蒸馏logit部分存在高度耦合, 这是导致logit蒸馏局限性的根本原因, 该方法通过将logit中的信息拆分为目标类(target class)和非目标类(non-target class)两部分, 来达到解耦的效果, 将经典的KD改为TCKD和NCKD, 从实验结果来看, TCKD只是起到了转移训练样本的“困难”知识, NCKD中转移了大量的暗知识, 这一部分才是logit蒸馏的主角, 将学生模型输出的logit不断地向教师模型的logit靠近, 然而, 这种方法忽略了一个问题, 当师生网络模型差距较大的时候, 即使通过上述方法, 将经典KD进行解耦, 效果也无法达到理想状态^[12], 因为教师网络模型和学生网络模型之间的表征差距过大, 学生模型的特征无法对齐教师模型的特征.

为了解决上述的问题, 本文从一个新的角度出发, 认为解决这种体量差距较大的知识蒸馏的关键点是在于如何摒弃掉无用的噪声, 学生模型可以视作一个低配版的教师模型, 由于其体量较小, 学习能力相对较弱, 在训练过程中不能准确地分辨出哪些是需要学习的知识, 这种对噪声的提炼会导致训练效果退化. 本文采用扩散模型^[13,14]来执行去噪模块, 先消除掉学生模型内部的噪声信息, 再进行蒸馏. 在这种方法下的学生模型更加贴合教师模型.

本文主要贡献如下.

1)本文提出了一个噪声自适应模块, 能够更精确的指定学生特征的噪声水平.

2)使用了一个由ResNet^[15]中两个瓶颈块组成的轻量级扩散模型, 从实际情况出发搭建更加符合知识蒸馏的扩散模型.

3)解决了以前知识蒸馏中师生特征无法对齐的问题, 在利用解耦知识蒸馏在logits上可以获取低纬度语义信息的优势, 来优化模型效果.

1 相关工作 1.1 知识蒸馏

知识蒸馏这一理论最早出现是在Hinton等人^[1]所发表的vanilla-KD论文中, 采用的是一种师生模型网络, 将预训练的教师模型的知识通过网络传递给学生模型, 输出一般为模型的预测和中间特征. 在这个过程中学生模型的输出为$ {F}^{(s)} $, 教师模型的输出为$ {F}^{(t)} $, 最后得到的知识蒸馏损失为:

$ {L}_{\rm kd}: =d({F}^{(t)}{, }{F}^{(s)}) $

(1)

其中, $ d $为距离函数, 通过该函数测量师生模型输出间的差距, 例如在概率分布任务中使用的Kullback-Leibler(KL)散度, 一般分类任务中的交叉熵(cross entropy), 中间特征和回归输出使用的均方误差(MSE).

1.2 扩散模型

扩散模型是一种对图像增强和图像恢复的非线性的处理方法, 正向扩散过程遵循马尔可夫链的概念, 通过控制时间步t不断向图像中添加噪声, 然后利用预测和去噪逆转这一过程. 用$ y $表示样本, 下角标为时间步$t \in \left\{ {0, 1, 2, \cdots , T} \right\}$, 扩散模型的过程可表示为:

$ q\left( {{y_t}\mid{y_0}} \right) = \mathcal{N}\left( {{y_t};\sqrt {{{\bar \alpha }_t}} {y_0}, \left( {1 - {{\bar \alpha }_t}} \right)I} \right) $

(2)

其中, 定义${\alpha _t} = 1 - {\beta _t}, {\bar \alpha _t} = \prod _{s = 1}^t{\alpha _s}$, ${\beta _t}$是一个超参数序列, 为每一时间步中应添加的噪声量, 通常会随着时间的增长而增长. 当${\textit{z}}$是从一个均值为$\mu $方差为${\sigma ^2}$的高斯分布中采样得到的^[16], 可以引入一个服从标准高斯分布的随机变量$\mathcal{E}$:

$ {\textit{z}} \sim \mathcal{N}\left( {{\textit{z}}, \mu , {\sigma ^2}I} \right) \to {\textit{z}} = \mu + \sigma \cdot \mathcal{E}, \mathcal{E} \sim \mathcal{N}\left( {0, I} \right) $

(3)

因此可以将${y_t}$表示为${y_0}$与噪声$\mathcal{E}$的线性组合:

$ {y_t} = \sqrt {{{\bar \alpha }_t}{y_0}} + \sqrt {1 - {{\bar \alpha }_t}\mathcal{E}} $

(4)

在逆向扩散过程中需要训练一个神经网络${\Phi _\theta }\left( {{y_t}, t} \right)$, 通过缩小$\mathcal{E}$与${\Phi _\theta }\left( {{y_t}, t} \right)$的${L_2}$损失来预测${y_t} \cdots {y_0}$中的噪声, 即:

$ {\mathcal{L}_{\rm diff}}: = \left\| {{\mathcal{E}_t} - {\Phi _\theta }\left( {{y_t}, t} \right)} \right\|_2^2 $

(5)

在推理过程中, 初始噪声为${Y_t}$, 通过训练好的${\Phi _\theta }\left( {{y_t}, t} \right)$, 来完成对数据样本${Y_0}$的迭代去噪过程:

$ {p_\theta }\left( {{y_{t - 1}}\mid{y_t}} \right): = \mathcal{N}\left( {{y_{t - 1}};{\Phi _\theta }\left( {{y_t}, t} \right), \sigma _t^2I} \right) $

(6)

其中, $\sigma _t^2$为DDIM中的过渡方差.

本文受到上述方法的启发, 利用扩散模型对学生模型进行去噪, 如图1 DDKD架构图所示, 学生模型的特征在初始阶段可视作带有噪声的教师特征, 然后利用教师特征训练扩散模型, 最后用训练好的扩散模型对学生特征进行去噪, 在得到无噪声的学生特征后进行蒸馏.

2 本文方法

在本节中将介绍本文所提出的基于扩散模型的解耦知识蒸馏, 首先KD中的特征对齐任务会被转换为扩散模型的去噪过程, 通过对齐师生模型的特征来获得更有效的蒸馏.

为了更好地提高计算效率, 本文引入了特征自编码器^[17]来降低特征映射的维数, 从而简化了扩散过程. 此外, 本文还提出了自适应噪声匹配模块, 以提高学生特征的去噪性能. 在最后将特征的高语义信息进行蒸馏.

2.1 基于扩散模型的解耦知识蒸馏

一般来说不同的模型有着不同的架构, 在特征的提取上也会有不同的关注点, 即使是在相同的数据集上训练也是如此. 教师模型体量较大, 在特征提取上比学生模型的效果好, 当教师模型与学生模型之间体量相差越大的时候, 他们的差距就会越加明显. 所以解决知识蒸馏问题的关键在于如何缩小这种师生差距. 在以前的论文中^[18]也调查过教师模型和学生模型之间的差异. 教师的预测概率分布比学生的预测概率分布更加清晰, 而且教师模型的预测错误答案的方差也比学生模型要小^[19], 以上这些说明, 教师模型的输出相对于学生模型来说要更加显著. 学生模型的预测比教师模型的预测包含了更多的噪声.

由于教师模型和学生模型之间的差距, 这些噪声无法仅凭通过简单的模仿蒸馏中的教师模型来消除. 但可以将教师模型和学生模型关注点更多地放在有价值的信息上面, 本文受到扩散模型的启发^[20], 将学生模型视为教师模型的带噪版本, 先用教师特征去训练一个扩散模型, 然后再用它对学生特征进行去噪.

对于一个样本在蒸馏过程中教师特征表示为${F^{\left( {\rm Tea} \right)}}$学生特征表示为${F^{\left( {\rm Stu} \right)}}$, 在前向噪声过程$q\left( {F_t^{\left( {\rm Tea} \right)}\mid{F^{\left( {\rm Tea} \right)}}} \right)$式(2)中使用${F^{\left( {\rm Tea} \right)}}$来训练$ \mathcal{L}_{\text {diff }} $式(5)的扩散模型.

然后将学生特征输入到学习扩散模型的迭代去噪过程中, 即式(6)中的$ p_\theta\left(F_{t-1}^{({\rm Stu})} \mid F_t^{({\rm Stu})}\right)$, 其中${F^{\left( {\rm Stu} \right)}}$为该过程的初始噪声特征. 经过这个去噪过程, 我们得到一个去噪的学生特征${F^{\left( {\rm Stu} \right)}}$, 用它来计算式(1)中与原始教师特征${F^{\left( {\rm Tea} \right)}}$的KD损失.

根据经验来说logits在语义水平上要比深层特征更高级, 所以本文中提到的方法是在整个模型架构的最后一部分使用解耦知识蒸馏, 将去噪后的学生输出${F^{\left( {\rm Stu} \right)}}$再进行logits部分的解耦. DDKD的架构如图1所示.

2.2 具有线性自编码器的扩散模型

由于教师特征尺寸相对较大, 模型在去噪过程会耗费大量的计算资源, 需要转发T次(在本文方法中使用T = 5)噪声预测网络${\Phi _\theta }$来去噪学生特征和训练1次带有教师特征的噪声预测网络. 当教师特征的维度较大时, 这T+1次转发会导致计算成本很高. 考虑到这个问题, 本文从ResNet中提取了一种由两个瓶颈块叠加到一起的轻量级扩散模型. 然后, 遵循着扩散模型的方法, 并提出使用线性自编码器模块压缩通道数量. 将压缩后的特征作为扩散模型的输入. 线性自编码器由两个卷积神经网络组成, 一个是编码器用于减少通道数量, 另一个是解码器用来重构教师特征. 编码器的输出特性被用来训练扩散模型和学生模型(式(2)).

自编码器只使用重构损失进行训练, 重构损失是原始教师${F^{\left( {\rm Tea} \right)}}$与重构教师特征$\tilde F_{\rm ae}^{\left( {\rm Tea} \right)}$之间的均方误差, 即:

$ {\mathcal{L}_{\rm ae}}: = \left\| {{{\tilde F}^{\left( {\rm Tea} \right)}} - {F^{\left( {\rm Tea} \right)}}} \right\|_2^2 $

(7)

用于训练扩散模型的潜在教师特征$ {Z^{({\rm Tea})}} $与扩散模型是分离的, 并且扩散模型没有向后梯度.

本文还使用卷积层将学生特征投影到与教师潜在特征$ Z^{(\text {Tea})} $相同的维度, 表示为$ Z^{({\rm Stu})} $. 然后将其传递给扩散模型执行反向去噪过程如式(6), 并生成去噪的学生特征$ Z^{({\rm Stu})} $. 然后用它们来计算KD损失并监督学生, 即:

$ {\mathcal{L}_{\rm DDKD}}: = d\left( {{{\hat Z}^{\left( {\rm Stu} \right)}}, {Z^{\left( {\rm Tea} \right)}}} \right) $

(8)

本文使用简单的MSE损失和KL散度损失作为距离函数$d$来计算去噪的学生特征和教师特征的差异, 然后利用基准的方法, 在logits部分将经典的KD损失分解为两个部分, 一个是目标类和非目标类的二分类预测, 另一个是非目标类的多分类预测.

2.3 自适应噪声匹配

如前文所述, 本文将学生特征视为教师特征的嘈杂版本. 然而, 表示教师和学生特征之间差距的噪声水平是未知的, 并且可能因不同的训练样本而变化. 因此, 不能直接确定应该从哪个初始时间步长开始扩散过程. 为了解决这个问题, 本文引入了一个自适应噪声匹配模块, 将学生特征的噪声水平与预定义的噪声水平相匹配.

如图2所示, 将构建一个简单的卷积模块来学习一个融合了学生输出和高斯噪声的权值γ, 这有助于学生输出能更快匹配与初始时间步长t的噪声特征的相同噪声水平. 因此, 去噪过程中的初始噪声特征变为:

$ Z_T^{\left( {\rm Stu} \right)} = \gamma {Z^{\left( {\rm Stu} \right)}} + \left( {1 - \gamma } \right) \in T $

(9)

这种噪声适应可以自然地通过KD损失${L_{\rm kd}}$进行优化, 因为在去噪过程中, 当学生特征与适当的噪声水平匹配时, 这时去噪后的学生特征与教师特征的差异才会最小.

2.4 总体损失函数

DDKD的整体损失函数由原始任务损失、优化扩散模型的扩散损失、学习自编码器的重建损失以及对教师特征和去噪学生特征进行蒸馏的KD损失组成, 即:

$ {\mathcal{L}_{\rm train}} = {\mathcal{L}_{\rm task}} + {\lambda _1}{\mathcal{L}_{\rm diff}} + {\lambda _2}{\mathcal{L}_{\rm ae}} + {\lambda _3}{\mathcal{L}_{\rm DDKD}} $

(10)

其中, ${\lambda _1}, {\lambda _2}, {\lambda _3}$是用来平衡损失的损失权值. 在所有实验中, 简单地设${\lambda _1} = {\lambda _2} = 1$.

3 实验数据分析

为了验证本文方法的有效性, 我们在ImageNet^[21]和CIFAR-100^[22]上进行实验, 首先会介绍本文的实验细节和相关数据集, 在实验部分会介绍我们的消融实验, 来展示本文方法的有效性.

3.1 数据集介绍

ImageNet数据集, 是目前最大的图像识别数据库, 在分类、定位和检测任务中使用较多, 其中包含了14197122张图像, 共有21841个类别, 图像覆盖了生活中大多数能见场景, 包含了更多与图像无关的噪声和变化, 可以更好地测试模型的鲁棒性.

在本文中我们使用的是ImageNet的子数据集, 具有相同的效果, 训练集有128167张图像, 包含1000个类别, 每个类别包含约1300张图片, 验证集有50000张图像, 每个类别包含50张图像, 测试集有100000张图像, 每类100张图像.

CIFAR-100是一个常用的图像数据集, 用于图像分类任务和计算机视觉研究. 图像内容丰富多样, 涵盖了各种日常物体和动物, 可以更好地检验算法的泛化能力, 该数据集总共包含60000张彩色图像, 其中训练集包含50000张图像, 测试集包含10000张图像. 每张图像的尺寸为32×32像素, 分为100个类别. 每个类别包含500张训练图像和100张测试图像.

3.2 实验细节

为了验证本文方法的可行性, 将在ResNet^[15]、ShuffleNet^[23]、MobileNet^[24]、Wide ResNet (WRN)^[25]和VGG^[26]网络中进行对比.

对于ImageNet数据集, 本文的训练策略将Epoch设为100, Batchsize设为256, 学习率(LR)初始值设为0.1, 每30个Epoch再衰减为当前的0.1, 优化器(optimizer)使用的是随机梯度下降算法(SGD), 超参数权重衰退(weight decay)为0.0001.

对于CIFAR-100 数据集, 本文的训练策略将Epoch设为240, Batchsize设为64, 学习率(LR)初始值设为0.05, 优化器选择随机梯度下降算法, 权重衰退为0.0005, 学习率会在从第150个Epoch之后, 每30个Epoch衰减为当前大小的0.1.

3.3 消融实验

为了验证不同方法对实验效果的影响, 本文进行了消融实验, 在基线设置上, 使用了ResNet-18和MobileNet V1作为学生模型, 教师模型分别为ResNet-34和ResNet-50网络. 表1为在ImageNet数据集上的实验结果. 可以发现vanilla-KD的效果无论是在同构网络还是异构网络中, 在TOP-1和TOP-5上的准确率都是最低的, 当使用DKD方法时, 准确率会有明显提升, 在MobileNet V1和ResNet-50这种设置下更为明显, 这种方法解决了传统知识蒸馏中的目标类与非目标类的高耦合问题, 比传统知识蒸馏TOP-1的准确率提高了1.37%. 本文的方法相对于目前最先进的方法(DKD)在ResNet-18和ResNet-34这种网络设置下TOP-1的准确率提升了0.64%, 在MobileNet V1和ResNet-50设置下TOP-1的准确率提升了1.82%, 这是因为在蒸馏前利用扩散模型对学生模型进行了去噪过程, 进行了特征对齐. 这个过程减少了学生与教师模型的差距.

3.4 对比实验

为了验证本文方法的有效性, 本文还与其他方法进行对比, 如表2、表3所示, 分别为同构网络和异构网络准确率对比的结果.

从结果上来看, 本文方法在异构网络中有更为显著的效果, 考虑到其他方法是通过添加中间网络, 将知识过渡给学生模型, 本文提出了一个更为可靠的方法从根本上解决师生的差距问题, 在蒸馏开始前利用轻量级的扩散模型解决了在训练前师生特征对齐的问题, 然后再利用解耦知识蒸馏的方法来解决传统知识蒸馏中目标类与非目标类高度耦合的问题.

在CIFAR-100数据集上, 同系列网络架构设置的情况下最好的结果比基准方法提升了0.61%, 准确率提高了2.6个百分点. 在不同系列的网络架构设置下, 最好的结果比基准方法提升了0.71%, 准确率提高了3.4个百分点.

4 结束语

本文研究了教师和学生在知识蒸馏方面的差异. 从使用到的基准方法来看, 解耦知识蒸馏在logits部分将其分为目标类和非目标类确实有很大的提升, 但本文从开始就将学生与教师特征对齐, 从本质上缩小了师生之间的差距, 为了减少差异, 提高蒸馏性能, 本文从一个新的角度出发, 提出用扩散模型显式地消除学生特征中的噪声. 在此基础上, 进一步引入了一个带有线性自编码器的轻量级扩散模型来降低该方法的计算成本, 并引入了一个自适应噪声匹配模块来将学生特征与正确的噪声水平相匹配, 从而提高了去噪性能. 在图像分类任务上的大量实验验证了本文的有效性和泛化性.