【论文笔记_知识蒸馏_2022】Spot-adaptive Knowledge Distillation

点自适应知识蒸馏

摘要

知识提取(KD)已经成为压缩深度神经网络的一个良好建立的范例。进行知识提炼的典型方式是在教师网络的监督下训练学生网络，以利用教师网络中一个或多个点(即，层)的知识。蒸馏点一旦指定，在整个蒸馏过程中，所有训练样本的蒸馏点都不会改变。在这项工作中，我们认为蒸馏点应该适应训练样本和蒸馏时期。因此，我们提出了一种新的提取策略，称为点自适应KD (SAKD ),在整个提取周期的每次训练迭代中，自适应地确定每个样本在教师网络中的提取点。由于SAKD实际上关注的是“从哪里蒸馏”，而不是大多数现有作品广泛研究的“蒸馏什么”，因此它可以无缝集成到现有的蒸馏方法中，以进一步提高它们的性能。对10个最先进的蒸馏器进行了广泛的实验，以证明SAKD在均相和非均相蒸馏设置下改善其蒸馏性能的有效性。代码可在https://github.com/zju-vipa/spot-adaptive-pytorch.获得。

1.介绍

…
现有的蒸馏方法，无论是单点蒸馏还是多点蒸馏，都有一个共同的特征:蒸馏点通常是手工设计的选择，这对于优化是低效的，尤其是对于具有数百或数千层的网络。一方面，如果蒸馏点设置得太稀疏，学生模型不能被老师充分监督。另一方面，如果提取点设置得太密集，例如每个可能的层或神经元，则学生模型学习可能被过度正则化，这也恶化了提取性能。此外，目前的方法采用全局蒸馏策略，即一旦指定，所有样品的蒸馏点都是固定的。潜在的假设是这样的蒸馏点对于整个数据分布是最佳的，这在许多情况下是不正确的。理想情况下，我们希望自动确定每个样品每个可能点的蒸馏点。

在本文中，我们提出了一种新的蒸馏策略，称为点自适应KD (SAKD ),使蒸馏点适应训练样本和蒸馏阶段。为此，我们首先将学生模型和教师模型合并成一个多路径路由网络，如图2所示。当数据流经网络时，有多条可行的路径让数据到达输出层。轻量级策略网络被设计成当数据到达网络中的分支点时，在最佳传播路径上对每个样本做出决策。如果数据被策略网络路由到教师模型的层，则表明学生模型中的对应层(缩写为学生层)还不能替换教师模型中的层(缩写为教师层)。因此，这些教师层中的知识应该提取到相应的学生层中，否则，如果数据被策略网络路由到一些学生层，则表明这些学生层是相应教师层的良好替代，产生更好或至少相当的性能。在这种情况下，在这些层中不允许蒸馏。因为策略网络是在路由网络之上设计的，并且与路由网络同时优化，所以它可以在学生模型的不同训练迭代中自动确定每个样本的最佳蒸馏点。

所提出的方法集中于“从哪里提取”，这与主要研究“提取什么”的当前文献大相径庭。因此，它可以与现有方法无缝结合，以进一步提高其蒸馏性能。特别地，所提出的方法自然地与同质蒸馏兼容，其中学生模型与教师模型具有相同风格的架构。然而，实验表明，所提出的方法在学生模型与教师模型差异很大的异构蒸馏设置下也惊人地工作良好。此外，虽然所提出的方法主要是为多点蒸馏而设计的，但它也可以通过动态地确定每个训练样本的蒸馏与否来提高单点蒸馏的性能。

简而言之，我们在这项工作中做出了以下三个主要贡献:
1.据我们所知，我们首先介绍了自适应蒸馏问题，其中蒸馏点应该适应不同的训练样本和不同的蒸馏阶段。

2.我们提出了一种新的点自适应蒸馏策略来自动确定蒸馏点，使蒸馏点适应训练样本和蒸馏阶段。

3.在各种实验环境下进行了大量的实验，以展示所提出的方法对于提高现有最先进蒸馏器的蒸馏性能的有效性。

这项工作的提示组织如下。第二节概述了知识提炼和路由网络的相关工作。第三节详细描述了提出的点自适应蒸馏。第四节提供了实验设置和结果。第五节得出了结论。

2.相关工作

A.知识蒸馏
…

图2.建议方法的示意图。上面是策略网络，下面是路由网络。在这个例子中，有六个可能的蒸馏点。策略网络的输入是来自教师和学生模型的特征。政策网络的输出决定了蒸馏的地点。老师在蒸馏过程中被固定。

B.路由网络
路由网络[28]-[30]是一种具有高度模块化的神经网络，这是鼓励任务分解、降低模型复杂性和提高模型泛化能力所需的关键属性。路由网络通常由两个可训练组件组成:一组功能模块和一个策略代理。在神经网络设置中，功能模块由子网络实现，并用作处理输入数据的候选模块。对于每个样本，策略代理从这些候选中选择功能模块的子集，将它们组装成完整的模型，并将组装的模型应用于任务预测的输入数据。已经提出了几种算法来优化策略模块，包括遗传算法[31]、多代理强化学习[28]、重新参数化策略[29]等。

路由网络与条件计算[32]、[33]、专家混合[34]以及它们的现代注意力基础[35]和稀疏[36]变体密切相关。他们已经成功地应用于几个领域，如多任务学习[28]，迁移学习[37]和语言建模[30]。在这项工作中，借助于路由网络，我们提出了一种新的蒸馏策略来自动确定网络中的蒸馏点。

3.点自适应知识蒸馏

在这一节中，我们介绍了提出的点自适应KD。图2示出了所提出的方法的概述。整个模型由两个主要部分组成:多路径路由网络和轻量级策略网络。多路径路由网络由教师模型和学生模型组成，如果必要的话，具有使它们的特征相互适应的适应层。当数据到达路由网络中的分支点时，策略网络用于对数据流路径上的每个样本做出路由决定。

所提出的蒸馏方法的总体思想是自动确定是否在候选蒸馏点进行蒸馏，如图2所示。如果样本被策略网络路由到某些教师层，则表明对应的学生层还不能替换这些教师层。因此，这些教师层中的知识应该被提取到相应的学生层中。否则，如果数据通过策略网络传输到一些学生层，则表明这些学生层是相应教师层的良好替代，产生更好的或至少相当的性能。在这种情况下，不再需要蒸馏。蒸馏的最终目标是使策略网络逐渐选择学生层来路由数据，这意味着学生模型是教师网络的良好替代品。

A.多路径路由网络
不失一般性，假设用于视觉分类的典型卷积神经网络(CNN)由用于表示学习的几个卷积块、用于向量化特征图的全连接层和用于进行概率预测的softmax层组成。每个卷积块由几个卷积层组成，每个卷积层之后是一个非线性激活层和一个批量归一化层。一般来说，在每个块之后，使用汇集层或卷积层将特征图缩小2倍或更多。形式上，我们用以下方式表示教师CNN的功能：

学生CNN的基本功能为：

其中S表示softmax函数，F表示线性函数。B_i表示第i个块下面的函数，上标s和t分别表示学生模型和教师模型。符号⚪表示功能组合操作。多路径路由网络由教师网络和学生网络组成，它们的中间层相互连接。然而，在教师和学生特征之间可能存在维度不匹配。我们采用由1 × 1卷积层实现的适配层[6]来对齐它们的特征。多路径路由网络的基本功能表示为：

其中w和w_i是由策略网络产生的特征融合权重，以[0，1]为界。当特征融合权重取{0，1}的离散值时，网络实际上变成了一个组合网络，其层由教师层和学生层交织而成。请注意，在等式中。(4)为了简单起见，我们省略了用于特征对准的适配层。利用路由网络M，我们的最终目标是获得一个孤立的学生模型S，它尽可能好地执行感兴趣的任务。

B.决策网络
我们采用一个决策网络来对通过路由网络的数据流路径上的每个样本进行决策。这里我们简单地采用一个轻量级的全连接层来实现策略网络。策略网络的输入是来自教师和学生模型的级联特征。策略网络的输出是N + 1个二维路由向量，其中N + 1表示分支点的数量，即候选蒸馏点的数量。每个路由向量都是一个概率分布，从中我们随机抽取一个分类值来决定路由网络中一个分支点的数据流路径。由于离散化，采样操作是不可微的。为了实现采样操作的可区分性，我们利用Gumbel-Softmax [39]来实现策略网络。形式上，假设对于路由网络中的第I个分支点，对应的路由向量是aij = {ai1，ai2}，其中元素ai1存储概率值，该概率值表示第i个块中的教师层将被用于处理输入数据的可能性。在正向传播过程中，策略根据分类分布做出离散决策:
这里wⁱ是二维的独热向量。“one hot”是返回一个one-hot vector的函数，其中只有指定的元素为1，所有其他元素为0。ε∈ R2是一个向量，其中的元素是从Gumbel分布(0，1)中提取的同分布样本，以添加少量噪声，从而避免argmax运算总是选择具有最高概率值的元素。

为了实现离散采样函数的可微性，我们使用Gumbel-Softmax技巧在反向传播期间放松wi，如下所示：

其中，τ是温度，控制近似后分布的尖锐程度。请注意，对于每个矢量wi，由于，我们只需使用wi和1wi来表示教师层和学生层的路由决策，如等式1所示。(4).

C.点自适应蒸馏
提出的点自适应提取是通过同时训练路由网络和策略网络来进行的。从策略网络和路由网络的角度来看，训练所提出的网络都是非静态的，因为最佳路由策略取决于模块参数，反之亦然。在这项工作中，以端到端的方式同时训练多路径路由网络和策略网络。总体目标是：

其中L_student是目标和仅来自学生模型的预测之间广泛使用的交叉熵损失[40]。L_KL是教师和学生预测之间的KL散度[41]，这也是Hinton等人提出的传统蒸馏损失[1]。LKD是将现有的知识蒸馏损失强加于中间层，如FitNets提出的损失[6]，注意力转移[7]等。由于所提出的方法关注于“在哪里蒸馏”而不是当前文献主要研究的“蒸馏什么”,所以它可以与大多数现有的蒸馏方法相结合。学生、LKL和LKD被用来使学生模型与教师模式的效果相近。L_routing是路由网络中目标和预测之间的交叉熵损失。β₁、β₂和β₃是权衡这些损失项的三个超参数。

注意，在整个训练阶段，教师模型的预训练参数保持固定。可训练参数仅包括学生模型、适配层和策略网络的参数。由于策略网络和适配层仅参与计算L_routing，因此它们的参数在L_routing的监督下被训练。学生网络和政策网络形成一个循环，其中学生模型的输出进入政策网络，政策网络的输出再次进入学生网络。为了稳定学生网络的训练，我们不再将梯度从策略网络反向传播到学生网络。在训练过程的早期阶段，由于教师模型被很好地预先训练，样本更可能被策略网络传递到教师层。在这种情况下，知识提炼发生在所有候选提炼点。随着训练的进行，学生模型逐渐在不同的层次上不同程度地掌握了老师的知识。在这种情况下，策略网络可以为教师层和学生层交织的每个样本规划一条路径。因此，在某些层自适应地进行知识提炼，以推动只涉及学生层的最佳策略。在下面的部分中，我们提供了详细的优化算法来更清楚地描述所提出的蒸馏方法。

D.最优化算法
…

4.实验

这里我们首先简要描述一下主要的实验设置，这对于理解我们在这项工作中的实验是必要的。然后通过基准测试比较，验证了该方法的优越性。最后，我们进行了广泛的烧蚀研究，以说明所提出的方法的有效性。
…

5.结论

在这项工作中，我们认为知识提炼策略不应该在整个训练过程和所有训练样本中固定不变。因此，我们提出了点自适应蒸馏策略，以自动决定每个样品在不同蒸馏阶段的蒸馏点。所提出的蒸馏策略首先将教师网络和学生网络合并成多路径路由网络。然后，引入策略网络来对通过网络的数据流路径上的每个样本进行决策。提取决策是基于来自策略网络的路由决策做出的。我们在均相和非均相蒸馏环境下评估了所提出的方法。为了验证该方法的有效性，还进行了大量的烧蚀研究。实验结果表明，所提出的自适应蒸馏策略持续改善了各种现有蒸馏方法的性能。

对于未来可能的工作，可以研究所提出的方法以应用于更广泛的场景，例如，在原始训练数据不可用的情况下的无数据提取，或者在用趋势变压器替换模型架构的情况下的基于变压器的提取。