ECCV 2022 | CMU提出首个快速学问蒸馏的视觉框架：ResNet50 80.1%精度，训练加速30%

今天介绍一篇来自卡耐基梅隆大学等单位 ECCV 2022 的一篇关于快速学问蒸馏的文章，用基本的训练参数配置就可以把 ResNet-50 在 ImageNet-1K 从头开始 (from scratch) 训练到 80.1% (不运用 mixup，cutmix 等数据增强)，训练速度（尤其是数据读取开销）比拟传统分类框架节省 16% 以上，比之前 SOTA 算法快 30% 以上，是目前精度和速度双双最优的学问蒸馏政策之一，代码和模型已全部开源！

论文和项目网址：http://zhiqiangshen.com/projects/FKD/index.html代码：https://github.com/szq0214/FKD学问蒸馏（KD）自从 2015 年由 Geoffrey Hinton 等人提出之后，在模型紧缩，视觉分类检测等领域产生了巨大影响，后续产生了无数相关变种和扩展版本，但是大体上可以分为以下几类：vanilla KD，online KD，teacher-free KD 等。最近不少研究表明，一个最简单、朴素的学问蒸馏政策就可以获得巨大的性能提升，精度甚至高于很多复杂的 KD 算法。但是 vanilla KD 有一个不可避免的缺点：每次 iteration 都需要把训练样本输入 teacher 前向传播产生软标签 (soft label)，这样就导致很大一部分计算开销花费在了遍历 teacher 模型上面，然而 teacher 的规模通常会比 student 大很多，同时 teacher 的权重在训练过程中都是固定的，这样就导致整个学问蒸馏框架进修效率很低。针对这个问题，本文首先分析了为何没法直接为每张输入图片产生单个软标签向量然后在不同 iterations 训练过程中复用这个标签，其根本原因在于视觉领域模型训练过程数据增强的运用，尤其是 random-resize-cropping 这个图像增强政策，导致不同 iteration 产生的输入样本即使来源于同一张图片也可能来自不同区域的采样，导致该样本跟单个软标签向量在不同 iterations 没法很好的匹配。本文基于此，提出了一个快速学问蒸馏的设计，通过特定的编码办法来处理需要的参数，继而进一步保存复用软标签（soft label），与此同时，运用分配区域坐标的政策来训练目标网络。通过这种政策，整个训练过程可以做到显式的 teacher-free，该方法的特点是既快（16%/30% 以上训练加速，对于集群上数据读取缓慢的缺点尤其友好），又好（运用 ResNet-50 在 ImageNet-1K 上不运用额外数据增强可以达到 80.1% 的精度）。首先我们来回顾一下普通的学问蒸馏结构是如何工作的，以下图所示：

学问蒸馏框架包含了一个预训练好的 teacher 模型（蒸馏过程权重固定），和一个待进修的 student 模型, teacher 用来产生 soft 的 label 用于监视 student 的进修。可以看到，这个框架存在一个比较明显的缺点：当 teacher 结构大于 student 的时候，训练图像前馈产生的计算开销已经超过 student，然而 teacher 权重并不是我们进修的目标，导致这种计算开销本质上是 “无用的”。本文的动机正是在研究如何在学问蒸馏训练过程中避免或者说重复利用这种额外的计算结果，该文章的解决政策是提前保存每张图片不同区域的软监视信号（regional soft label）在硬盘上，训练 student 过程同时读取训练图片和标签文件，从而达到复用标签的效果。所以问题就变成了：soft label 怎么来组织和保存最为有效？下面具体来看该文章提出的政策。1. FKD 算法框架介绍FKD 框架的核心部分包含了两个阶段，以下图：（1）软标签（soft label）的天生和保存；（2）运用软标签（soft label）进行模型训练。

如图所示，上半部分展示了软标签的天生过程，作家通过输入多个 crops 进入预训练好的 teacher 来产生需要的软标签向量，同时作家还保存了：（1）每个 crop 对应的坐标和（2）是否翻转的 Boolean 值。下半部分展示了 student 训练过程，作家在随机采样图片的时候同时也会读取它们对应的软标签文件，从中选取 N 个 crops 用于训练，额外数据增强比如 mixup，cutmix 会放在这个阶段，从而节省了由于引入更多数据增强参数带来的额外保存开销。2. 采样政策本文还提出了一个 multi-crop sampling 的政策，即在一个 mini-batch 里面每张图片采样多个样本 crops。当总的训练 epochs 不变的前提下，该采样办法可以大大减少数据读取的次数，对于一些数据读取不是非常高效或者产生严重瓶颈的集群设备，这种政策的加速效果非常明显（以下表格所示）。同时在一张图片采样多个 crops 可以减少训练样本间的方差，帮助稳定训练，作家发现如果 crops 的数目不是太大的情况下可以明显提升模型精度，但是一张图片里面采样太多 crops 数目会造成每个 mini-batch 里面训练样本的信息差异不足（过于相似），因此过度采样会影响性能，所以需要设置一个合理的数值。3. 加速比作家在实验部分跟标准的训练办法以及 ReLabel 训练进行了速度的比较，结果以下表格所示：可以看到，比拟正常的分类框架，FKD 会快 16% 左右，而比拟 ReLabel 则快了 30%，因为 ReLabel 比拟正常训练需要读取双倍的文件数目。需要注意的是这个速度对比实验中，FKD crop 数目为 4，如果选取更大的 crop 数目可以得到更高的加速比。

加速原因分析：除了上述介绍的采用多个 crops 来进行加速外，作家还分析了其他一些加速的因素，以下图所示，ReLabel 在训练模型阶段需要天生采样数据的坐标，同时需要运用 RoI-Align 和 Softmax 来天生所需的软标签，比拟而言，FKD 直接保存了坐标信息和最终软标签格式，因此读取标签文件之后不需要做任何额外的后处理就可以直接训练，速度比拟 ReLabel 也会更快。

4. 标签质量分析软标签质量是保证模型训练精度的一项最重要的指标，作家通过可视化标签分布以及计算不同模型预测之间的交叉熵（cross-entropy）来证明了所提出的办法拥有更好的软标签质量。

上图展示了 FKD 和 ReLabel 软标签分布的情况对比，得到以下结论：（第一行）FKD 比拟 ReLabel 置信度更加平均也与输入样本内容更加一致，作家分析原因是 ReLabel 将全局图像输入到模型中，而不是局部区域，这使得天生的全局标签映射编码了更多全局类别信息同时忽略了背景信息，使得天生的软标签过于接近单个语义标签。（第二行）虽然存在一些样本 ReLabel 和 FKD 之间的最大预测概率相似，但 FKD 包含更多标签分布中的从属类别概率，而 ReLabel 的分布中并没有捕获这些从属类别的信息。（第三行）对于某些异常情况，FKD 比 ReLabel 更加健壮，例如目标框含有松散边界，或者只定位部分目标等。（第四行）在有些情况下，ReLabel 的标签分布意外的崩溃了（均匀分布），没有产生一个主要的预测，而 FKD 仍然可以预测得很好。5. 标签紧缩、量化政策1）硬化 (Hardening)。在该政策中，样本标签 Y_H 运用 teacher 预测的最大 logits 的索引。标签硬化政策产生的依然是 one-hot 的标签，以下公式所示：

2）光滑 (Smoothing)。光滑量化政策是将上述硬化后的标签 Y_H 替换为软标签和均匀分布的分段函数组合，以下所示：

3）边际光滑 (Marginal Smoothing with Top-K)。边际光滑量化政策比拟单一预测值保留了更多的边际信息（Top-K）来光滑标签 Y_S：

4）边际光滑归一化 (Marginal Re-Norm with Top-K)。边际光滑归一化政策会将 Top-K 预测值重新归一化到和为 1，并保持其他元素值为零（FKD 运用归一化来校准 Top-K 预测值的和为 1，因为 FKD 保存的软标签是 softmax 处理之后的值）：

具体对应上述各种量化政策的图示以下图所示：

6. 不同标签量化 / 紧缩政策的保存大小比较不同标签紧缩方法需要的保存空间以下表格所示，所运用的数据集为 ImageNet-1K，其中 M 是软标签天生阶段每张图像被采样的数目，这里作家选取了 200 作为示例。Nim 是图像数量， ImageNet-1K 数据集为 1.2M，SLM 是 ReLabel 标签矩阵的大小，Cclass 是类的数量，DDA 是需要保存的数据增强的参数维度。

从表格中可以看到，在不做任何紧缩的情况下 FKD 软标签需要的保存空间为 0.9T，这在实际运用中显然是不现实的，标签数据的大小已经远远超过训练数据本身了。通过标签紧缩可以极大减少保存大小，同时后面实验也证明了合适的紧缩办法并不会损害模型精度。7. 自监视进修任务上的应用FKD 的训练办法也可以应用于自监视进修任务。作家运用自监视算法比如 MoCo，SwAV 等来预训练 teacher 模型，然后按照上述办法天生用于自监视的软标签（unsupervised soft label），这个步骤跟监视进修得到的 teacher 很相似。天生标签过程会保留原始自监视模型中 projection head 并运用之后的最终输出向量，然后将这个向量作为软标签保存下来。得到该软标签后，可以运用同样的监视式的训练办法来进修对应的 student 模型。8. 实验结果1）首先是在 ResNet-50 和 ResNet-101 上的结果，以下表所示，FKD 取得了 80.1%/ResNet-50 和 81.9%/ResNet-101 的精度。同时训练时间比拟普通训练和 ReLabel 都快了很多。

2）作家还测试了 FKD 在 MEAL V2 上的结果，同样得到了 80.91% 的结果。

3）Vision Transformer 上的结果：接下来作家展示了在 vision transformer 上的结果，在不运用额外数据增强的情况下，FKD 就可以比之前学问蒸馏方法得到将近一个点的提升，同时训练速度快了 5 倍以上。

4）Tiny CNNs 上的结果：

5）消融实验：

首先是不同紧缩政策，综合考虑保存需求和训练精度，边际光滑政策是最佳的。接下来是训练阶段不同 crop 数目的对比，MEAL V2 由于运用了 pre-trained 的参数作为初始化权重，因此不同 crop 数目下性能都比较稳定和接近。而 vanilla 和 FKD 在 crop=4 的时候表现得最好。尤其 vanilla，比拟 crop=1 精度提升了一个点，crop 大于 8 之后精度下降明显。

6）自监视任务上的结果：以下表所示，在自监视进修任务上 FKD 办法还是可以很好的进修目标模型，同时比拟双子结构自监视网络训练和蒸馏训练，可以加速三到四倍。

9. 下游任务下表给出了 FKD 模型在 ImageNet ReaL 和 ImageNetV2 两个数据集上的结果，可以看到，FKD 在这些数据集上取得了稳定的提升。

下表是 FKD 预训练模型在 COCO 目标检测任务上的结果，提升同样明显。

10. 可视化分析以下两张可视化图所示，作家通过可视化中间特征层（attention map）的办法探索 FKD 这种 region-based 训练办法对模型产生的影响，作家对比了三种不同训练办法得到的模型：正常 one-hot label，ReLabel 和本文提出的 FKD。（i） FKD 的预测的概率值比拟 ReLabel 更加小（soft），因为 FKD 训练过程引入的上下文以及背景信息更多。在 FKD 随机 crop 的训练政策中，许多样本采样于背景（上下文）区域，来自 teacher 模型的软预测标签更能真实的反映出实际输入内容，并且这些软标签可能与 one-hot 标签完全不同，FKD 的训练机制可以更好的利用上下文中的额外信息。（ii） FKD 的特征可视化图在物体区域上具有更大的高响应值区域，这表明 FKD 训练的模型利用了更多区域的线索进行预测，进而捕获更多差异性和细粒度的信息。（iii）ReLabel 的注意力可视化图与 PyTorch 预训练模型更加接近，而 FKD 的结果跟他们比拟具有交大差异性。这说明 FKD 办法进修到的注意力机制跟之前模型有着显著的差别，从这点出发后续可以进一步研究其有效的原因和工作机理。