【知识蒸馏】 DistPro: Searching A Fast Knowledge Distillation Process via Meta Optimization

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

一、摘要

最近的知识蒸馏 (KD) 研究表明，不同的手动设计方案显着影响学习的结果。然而，在 KD 中，自动搜索最佳蒸馏方案尚未得到很好的探索。在本文中，我们提出了 DistPro，这是一个新颖的框架，它通过可微元学习搜索最佳 KD 过程。具体来说，给定一对学生和教师网络，DistPro 首先将一组丰富的KD连接从教师的传输层设置为学生的接收层，同时还提出了各种变换来比较特征图及其蒸馏途径。然后，连接和转换选择（路径）的每个组合都与随机加权过程相关联，该过程表明其在蒸馏过程中的每一步的重要性。在搜索阶段，可以通过我们提出的两级元优化策略有效地学习该过程。在蒸馏阶段，DistPro 采用学习到的知识蒸馏过程，显着提高了学生的准确性，尤其是在需要更快的训练时。最后，我们发现学习到的过程可以在类似的任务和网络之间进行泛化。在我们的实验中，DistPro 在流行数据集的不同数量的学习时期（即 CIFAR100 和 ImageNet）下产生了最先进的 (SoTA) 准确度，这证明了我们框架的有效性。

二、相关工作

1.KD领域

知识评论和 L2T-ww与我们的工作最相关。前者研究了一些路径的重要性，并提出了一种具有新颖连接模式的知识评论机制，即它在几个常见的比较基准测试中提供了 SoTA 结果。后者为少镜头知识转移的几个中间路径学习一个固定的权重。
如图1所示，DistPro找到一个蒸馏过程。因此，我们扩展了搜索空间。此外，对于密集预测任务，一项相关工作是IFVD，它提出了一种类内特征变异比较(IFVD)。DistPro 可以自由地扩展到密集预测任务，并且在与 IFVD 结合后也获得了额外的好处。

图 1：教师和学生模型之间学习过程的蒸馏方法的比较。(a) 知识评论通过枚举不同的配置提出了固定的采样路径。(b) L2T-ww采用元学习框架来学习每个选定路径的浮点权重。(c）我们的框架为每个路径学习蒸馏过程。

2.KD/超参数的元学习

为了自动学习 KD 方案，我们研究了我们可能采用的其他领域的广泛有效的元学习方法。例如，L2L 建议通过基于 RNN 的元网络来学习超参数调度方案。法国schi et al.al 提出了一种没有外部元网络的基于梯度的方法。后来，这些元学习思想也被用于小样本学习的任务（例如学习重新加权）、学习跨任务知识转移（例如学习迁移 ）和神经架构搜索（NAS）（例如 DARTS）。通过这些方法具有相似的框架，虽然必须具有基本的嵌入领域知识和任务感知调整以使其工作至关重要。在我们的案例中，受这些方法的启发，我们主要利用基于梯度的策略，因为它对 KD 方案学习的效率，而且还首先提出在学习的重要性因素之外使用学习过程。

3.NAS 的知识蒸馏

例如，Liu et al.al试图找到最适合提取给定教师的学生模型，而 Yao et al. 建议根据一定的蒸馏损失搜索学生和教师模型的架构。尽管与我们的场景不同，即固定的学生-教师架构，但它引发了另一个重要问题，即如何在某些资源约束下找到架构和 KD 方案的联合空间内的帕累托最优，我们希望这可以激发未来的研究。

三、创新点方法

在本节中，我们通过首先设置具有构建搜索空间的中间特征的 KD 路径来详细解释 DistPro，同时建立我们的 KD 算法的符号和定义。通过推导出梯度，用于生成我们对 KD 算法的过程。最后，提出了整体DistPro算法。

3.1 The Distillation Process

现在我们将能够构建路径，将教师的任何特征层与具有适当转换的学生的任何层连接起来。然而，如第 2 节所述，并非所有这些路径都是有益的因此，通过为其分配重要性因子来找出每个可能路径的重要性。与现有工作不同，这里的重要性因素是一个过程。随机权重过程 $A_i$= { ${\alpha }_t^i$}${}_{t=1}^T$ 与路径 $i$相关联，其中$T$是总学习步骤。在这里, ${\alpha }_t^i$ 描述了不同学习步骤 $t$ 的重要性因子。

扫描二维码关注公众号，回复： 14653499 查看本文章

令 $D_{Train}$ := {($X_i,y_i$)}${}_{i=1}^{|D_{train}|}$ 和 $D_{val}$$:=(X_i,y_i)$${}_{i=1}^{|D_{val}|}$分别是训练集和验证集，其中$y_i$ 是样本$X_i$ 的标签。我们假设对于每对特征图 $F_j^s$和$F_i^t$ ，我们有 N 个候选预定义的变换 $M_{i,j,1},$$M_{i,j,2},......$$M_{i,jN},$构建了我们的搜索空间，如图 2 所示

图 2：搜索空间的图示。从学生和老师中选择了两组特征图。对于每对特征图 $F_j^s$和$F_i^t$，我们插入N个候选变换块$M_{i,j,1},$$M_{i,j,2},......$$M_{i,jN},$得到N个损失项。对于每个损失项，分配一个蒸馏过程 $A^{i,j,k}={\alpha }_{t=1}^{i,j,k}$。

我们现在定义由学生网络和 $A$ 的优化组成的搜索目标。学生网络在训练集上进行训练，损失编码来自真实标签和神经网络的监督。具体来说，用 $\theta$ 表示学生的参数和变换，学习步骤 t 训练集上的损失定义如下，

其中${\alpha }_t^{i,j,k}$∈$R_{\ge 0}^{L_{te}\times L_s\times N}$是训练步骤 t 的重要因素，${\delta }_{label}$是测量预测标签和标签之间差异的距离函数。这里，${\alpha }_t$ =$[{\alpha }_t^{0,0,0,0},\cdot \cdot \cdot \cdot ]$控制所有知识蒸馏路径在训练步骤 $t$ 中的重要性。为了数值稳定性和避免噪声梯度，应用参数为$g=1$的有偏 softmax 归一化来计算${\alpha }_t$，这在我们的实验中从各种归一化策略中得到验证。这是 NAS 等各种任务元学习中常用的技巧或few-shot学习。有关归一化方法讨论的更多详细信息，请参见实验部分（第 4.4 节）。

接下来，我们的目标是找到一个最优的采样过程A ，产生最好的KD结果，其中验证集通常用于评估在未见输入上训练的学生的性能。为此，在[26]之后，采用验证损失来评价A的质量，其定义为:

最后，我们将 $A$ 和网络参数$\theta$上的双层优化问题表述为：

其中 $\theta \ast (A)$ 是用 $A$ 定义的损失过程训练的学生网络的参数，即$L_{train}(\Theta ,A)$ = ${L_{train}({\theta }_t,{\alpha }_t)}_{t=1}^T$。优化的最终目标是找到一个 $A$ 使得验证集上的损失最小化。
请注意，在 NAS 中已经提出了类似的问题，要求固定架构，但我们的问题不同且更难，如果我们强制 A 中的所有值相同，则可以将其简化为它们。然而，同样，我们能够在链式规则之后应用基于梯度的方法来解决这个问题。

3.2 Learning the Process A

从公式3，直接解决问题是棘手的。因此，我们提出了两个假设来简化问题。

• 首先，平滑假设这意味着该过程的下一步应该接近之前的步骤。这通常用于具有学习率调度器的随机梯度下降 (SGD) 的 DNN 训练。
• 其次，类似的学习过程假设，这意味着在给定教师/学生对的不同蒸馏训练中，在相同的步骤$t$，学生将具有相似的参数 ${\theta }_t$。这个假设允许我们使用单个训练过程搜索 $A$，当批量大小相对较大（例如 ImageNet 为 512）时，我们的实验也成立。

基于这些假设，我们让${\alpha }_{t+1}={\alpha }_t+\gamma \Delta ({\alpha }_t|{\theta }_t)$，其中 $\gamma$ 控制其变化率很小。然后，我们能够采用贪心策略将Eq3的原始问题分解为一系列单步优化，可以定义为:

其中 $\xi$ 是训练学生网络内部优化的学习率。在实践中，内部优化可以通过更复杂的基于梯度的方法来解决，例如具有动量的梯度下降。在这些情况下，方程式 4 必须相应地修改，但以下所有分析仍然适用。
接下来，我们将链式法则应用于方程式 4并得到

然而，上述表达式包含二阶导数，计算量仍然很大。接下来，我们用[26]中介绍的有限差分近似这个二阶导数。设 $\epsilon$ 为一个小的正标量，定义符号${\theta }^{\pm }=\theta \pm \epsilon {\nabla }_{\theta }{L}_{val}({\theta }_t+1)$。然后，我们有

最后，我们设置初始值 ${\alpha }_0=1$并启动贪心学习过程。一旦学习，我们将所有计算的${\alpha }_t$与$T$步推送到序列中，该序列用作学习随机蒸馏过程 $A$ 的样本，可用于在相同或相似的任务中用 KD 重新训练学生网络。

3.3 Acceleration and adopting A for KD.

现在让我们仔细看看近似值。为了评估方程式中的表达式5，必须计算以下过程：

1. 首先，计算${\theta }_t+1$需要学生的前向和后向传递，以及教师的前向传递。
2. 然后，计算 ${\theta }^{\pm }$ 需要学生的另一个前向和后向传递。
3. 最后，计算 ${\nabla }_{\alpha }{L}_{train}({\theta }^{\pm },{\alpha }_t)$ 需要学生的两次前向传递。

总之，评估方程式中的近似梯度5 需要教师的一次前向传递，总共有四个前向传递和学生的两次后向传递。这是一个耗时的过程，尤其是当所需的 KD 学习 epoch 很大时，例如 ≥ 100。
在元学习 NAS 文献中，为了避免二阶近似，研究人员通常仅基于训练损失采用一阶训练。但是，在我们的例子中这是不切实际的，因为超参数是在训练损失。${\alpha }_t$ 上的一阶梯度将简单地将所有值驱动为 0。因此，在我们的例子中，我们选择将 A 的学习 epoch 减少到 $T_{search}$，远小于 T，然后使用线性插值将其扩展为长度为 T 的序列。$T_{search}$ 的选择可以根据数据集大小动态调整，这将在第 4 节中详细说明。

此外，由于我们有比其他 KD 策略更多的路径，因此不会忽略损失计算成本。
为了减少成本，在步骤 $t$ 中，我们对方程式中的所有${\alpha }_t$ 使用公式2，阈值为 $T=0.5$，当 ̃αt ≤ τ 时丢弃相应的计算。这可以平均节省 60% 的损失计算成本，从而产生与基线相当的 KD 时间。
总之，整体 DistPro 过程在 伪代码流程 中给出1如下图所示，主要包含两个步骤：
该算法的第一阶段是利用部分数据集确定搜索方案 $A$。为此，交替计算${\alpha }_t$ 和 ${\theta }_t$。每个步骤中获得的${\alpha }_t$存储以供将来使用。
该算法的第二阶段是用所有可用的训练数据和搜索的$A$重新训练神经网络。

四、实验结果

在本节中，我们在包括图像分类、语义分割和深度估计在内的几个基准任务上评估所提出的方法。

4.1 图像分类 CIFAR100

1.Implementation details

如前所述，在蒸馏训练之前，我们需要通过如下所述搜索来获得蒸馏过程。
为了搜索过程 A，我们将原始训练集随机分成两个子数据集，80% 的图像用于训练，20% 用于验证。我们运行 40 个 epoch 的搜索，并在 epoch 10、20 和 30 将 θ（学生的参数）的学习率衰减 0.1。按照 [5] 中的设置，我们没有使用所有特征图进行知识蒸馏，但在每个下采样阶段只有那些。变换块用于转换学生特征图。图 4a 显示了块架构。路径的大小为 27，因为我们使用 3 个转换块和 9 个连接。为了公平比较，我们在 [5] 中遵循相同的 HCL 蒸馏损失。一旦获得该过程，为了对齐搜索的蒸馏过程，使用线性插值将 A 的过程从 40 个 epoch（搜索阶段）扩展到 KD 期间的 240 个 epoch（重新训练阶段）。

2.Results

定量结果总结在表1中。
为了证明学习 α 是必不可少的，我们首先将权重相同的 α 分配给所有路径。如“同等加权”行所示，结果比 ReviewKD 差，表明来自 ReviewKD 的选定路径是有用的。对于消融，我们首先在过程结束时采用学习到的${\alpha }_t$。如“使用 ${\alpha }_t$”所示，它在多个设置中优于 ReviewKD。最后一行显示了采用学习过程 A 的结果，根据我们的推导的变化，它明显优于 ReviewKD。

表1:CIFAR100的结果。结果是5次运行的平均值。方差在括号里。“+”表示我们的复现结果。在“同等权重”中，我们在再训练阶段不使用搜索的α。取而代之的是，将α的每个元素统一设为1/L，其中L为α的长度。在“使用${\alpha }_t$”中，最后收敛的α被用于再训练阶段的每次迭代。

4.2 图像分类 ImageNet1K

1.Implementation details
我们遵循 CIFAR100 中使用的搜索策略和 [5] 中的训练配置，批量大小为 512（使用 4GPU）。选定的学生/教师网络是MobileNet/ResNet50、ResNet18/ResNet34和DeiT-tiny/ViT-B。我们为DeiT采用了不同的转换块架构。在搜索阶段，我们使用Tiny-ImageNet搜索A 20个epoch。我们采用余弦lr调度器来学习学生网络。对于搜索空间，使用了 1 个转换块，而我们考虑网络中的 5 个特征图，并且由于 GPU 内存有限，通过删除之后的一些路径来构建 15 个路径。在 KD 阶段，我们以 0.1 的初始学习率和余弦调度器为各种 epoch 训练网络。搜索需要 4 个 GPU 小时和 80 个 GPU 小时来提取 100 个 epoch。由于空间有限，更多细节（ViT 变换块、选择路径等）可以在补充材料中找到

2.快速蒸馏
在之前的工作[39]中，KD被用于加速网络学习，特别是在大数据集的情况下。这是因为，如果使用大量的epoch[16]训练网络，则仅使用Ground Truth标签就可以提高其准确性。之前的工作reviewkd只关注在100个epoch停止的结果。在这里，我们认为用更少的训练成本来评估结果也很重要，这是评估KD方法的另一个重要指标。这对于需要快速学习网络的应用程序是非常有用的。

在这里，我们展示了DistPro在训练成本和准确性之间达到了更好的平衡。在图5中，具有不同网络架构的DistPro(红色曲线)在所有建议的训练阶段都优于具有相同设置的ReviewKD(蓝色虚线曲线)。当需要较少的训练成本时，性能增益更大，例如，在10个epoch的MobileNet上，性能增益为3.22% (65.25% vs 62.03%)，而在100个epoch时下降到0.14%，这仍然是一个不错的进步。ResNet18也得到了类似的结果。这里注意，对于所有不同epoch的结果，我们采用相同的学习过程A，其中时间成本与KD相比可以忽略不计。
在表2中，我们显示了使用ReviewKD[5]实现相同精度时DistPro节省的epoch数。例如，在训练MobileNet时，DistPro只使用50个epoch就达到了72.54%，这与ReviewKD用100个epoch训练(72.56%)相当，产生了2倍的加速度。ResNet18和DeiT也观察到类似的加速

3.Transfer A
在这里，我们研究学习到的A是否可以跨数据集和类似的架构传输。表3显示了结果。
在设置(a)中，我们将CIFAR100中的图像大小调整为224 × 224，并在CIFAR100上进行进程搜索，其搜索成本仅为2个GPU小时。如上一列所示，它只降低了0.06%的精度，这与使用ImageNet搜索的完整结果相当，表明该过程可以跨数据集传输。
在设置(b)中，由于MobileNet/ResNet50的学生/教师网络与ResNet18/ResNet34的网络在相似的对应层上具有相同的特征路径，因此我们采用二者搜索的过程。如图所示，结果是封闭的，表明当相似的路径存在时，该过程可以转移。
在MobileNet/ResNet50中学到的A在ResNet18/ResNet34师生对中也可使用。

4.3 消融实验

仅仅将知识从较低级别的特征映射转移到较高级别的特征映射总是有益的吗?
(Ablation study Is it always beneficial to transfer knowledge only from lower-level feature maps to higher-level feature maps?)
在[5]的表6中，作者对CIFAR100进行了一组实验，结果表明，只有从教师的低级别特征图到学生的高级别特征图的路径是有益的。然而，我们在cityscape上进行了类似的实验，结果并不支持这一说法。具体来说，在ResNet18上使用从最后一个特征图到前三个特征图的路径导致mIoU为73.9%，而使用从所有低级特征图到高级特征图的路径(如[5])导致mIoU为73.5%。
这表明，需要在更大的空间内寻找最优的教学方案。此外，表5中的结果还表明，我们搜索的流程比手工制作的流程更好。

Intuition of process A
在图6中，我们展示了从WRN-40-2蒸馏到WRN-16-2的搜索阶段，学习样本随时间的变化。在其他环境中也发现了类似的观察结果。由图可知，在训练初期，优化的教学方案侧重于将知识从教师的底层特征图传递给学生。随着训练的进行，优化的教学方案逐渐转移到教师的更高层次特征图。直观地说，高级特征映射编码了输入图像的高度抽象信息，因此与低级特征映射相比更难学习。DistPro能够自动找到一个教学方案来利用这种直觉。

图6:WRN-40-2为老师，WRN-16-2为学生，在CIFAR100上进行的搜索蒸馏过程。x轴为迭代次数。左图包含α[1，:，:]，即教师最低级特征图对应的所有元素。中间的对应于α[2，:，:]，右边的对应于α[3，:，:]。

Normalization of ${\alpha }_t$
如前所述，我们对${\alpha }_t$进行归一化，以获得二次梯度逼近后的数值稳定性。在这项研究中，我们评估了几种标准化策略。设非归一化参数为${\alpha }_t$，归一化参数为${\alpha }_t$。下面，我们将${\alpha }_t$视为长度为l的向量。实验在CIFAR100上进行，WRN-28-4作为老师，WRN-16-4作为学生。结果如表7所示，所提出的偏置softmax归一化优于其他方法。直观地说，由于附加了标量1，可以将${\alpha }_t$的值与1进行比较，常数值对分布产生标签平滑[32]效果。

---

结论

在本文中，我们朝着在给定一对想要的学生网络的情况下找到最优 KD 方案并在视觉任务下学习教师网络的问题迈出了一步。具体来说，我们通过在两个网络之间构建路径并沿每条路径路径路径分配一个随机蒸馏过程来设置搜索空间。我们提出了一个元学习框架 DistPro，可以有效地学习这些过程，并找到有效的具有中间特征的 KD。我们展示了它对图像分类的好处，以及图像分割和深度估计等密集预测。我们希望我们的方法能够激发 KD 领域进一步扩展范围，并与 NAS 和超参数调整等其他技术合作。