ICLR2021自蒸馏方法SEED：显著提升小模型性能-技术圈

点蓝色字关注“机器学习算法工程师”

设为星标，干货直达！

注：公众号后台回复：SEED。即可获取本文下载链接。

该文是亚利桑那州立大学&微软联合提出的一种自监督蒸馏表达学习方案，已被ICLR2021接收为Poster。针对现有对比自监督学习方案在小模型上表现性能差的问题，提出将知识蒸馏的思想嵌入到自监督学习框架得到了本文的自监督蒸馏学习。相比常规自监督学习方案，所提方案可以显著的提升轻量型模型的性能，同时具有极好的泛化性能。

Abstract

本文主要聚焦于小模型(即轻量型模型)的自监督学习问题，作者通过实证发现：对比自监督学习方法在大模型训练方面表现出了很大进展，然这些方法在小模型上的表现并不好。

为解决上述问题，本文提出了一种新的学习框架：自监督蒸馏(SElf-SupErvised Distillation, SEED)，它通过自监督方式(SSL)将老师模型的知识表达能力迁移给学生模型。不同于直接在无监督数据上的直接学习，我们训练学生模型去模拟老师模型在一组示例上的相似度得分分布。

所提SEED的简洁性与灵活性不言而喻，包含这样三点：(1) 无需任何聚类/元计算步骤生成伪标签/隐类；(2) 老师模型可以通过优秀的自监督学习（比如MoCo-V2、SimCLR、SWAV等）方法进行预训练；(3)老师模型的知识表达能力可以蒸馏到任意小模型中(比如更浅、更细，甚至可以是完全不同的架构)。

实验表明：SEED可以提升小模型在下游任务上的性能表现。相比自监督基准MoCo-V2方案，在ImageNet数据集上，SEED可以将EfficientNet-B0的精度从42.2%提升到67.6%，将MobileNetV3-Large的精度从36.3%提升到68.2%，见下图对比。

本文的主要贡献包含以下几点：

首次解决了小模型的自监督视觉表达学习问题；
提出一种自监督蒸馏技术用于将大模型的知识迁移到小模型且无需任何带标签数据；
基于所提SEED方案，我们可以显著的提升现有SSL技术在小模型上的性能表现;
系统性的比较了不同蒸馏策略，验证了SEED在多种配置的有效性。

Method

Knowledge Distillation

知识蒸馏最早是Hinton等人于2015年提出用于模型压缩的技术，它旨在将更大的老师模型的知识表达能力迁移给更小的学生模型。该任务可以表示成如下形式：

其中，分别表示输入图像及对应的标签，分别表示学生模型与老师模型的参数。表示监督损失，调控网络预测与标签之间的差异性，对于图像分类而言，它往往为交叉熵损失；表示蒸馏损失，用于调控学生模型预测与老师模型预测之间的差异性，它往往为KL-散度损失。知识蒸馏的有效性已在多个任务、多个数据集上得到了验证；但是在无监督任务上仍是未知的，而这是本文的聚焦所在。

Self-Supervised Distillation

不同于有监督蒸馏，SEED希望在无标签数据上将大模型的知识表达能力迁移给小模型，以促使小模型所学习到的表达能力更好的作用于下游任务。受启发于对比自监督学习，我们设计了一种简单的方案：在实例相似性分布的基础上进行知识蒸馏。类似MoCo，我们维持一个实例队列以保存老师模型的编码输出，对于新样本，我们计算它与队列中所有样本的相似性得分。我们希望：学生模型与老师模型的相似性得分分布尽可能相似，这就构成了老师模型与老师模型的相似性得分分布的交叉熵，见下图。

具体来讲，对于输入及其随机增广，首先进行编码&特征规范化的得到特征向量表达:

其中,分别表示老师模型与学生模型。假设表示实例队列，K表示队列长度，表示老师模型得到的特征向量。类似于对比学习方案，D采用"先进先出"的策略逐步更新并作为蒸馏所得，也就是说：队列中的样本大多是随机的、与当前实例不相关。最小化老师模型与学生模型在D与之间的交叉熵，而不会直接与老师模型进行直接对齐。为缓解不对齐问题，作者添加了老师模型的嵌入特征到上述队列构成了新的队列。

假设表示老师模型提取特征与之间的相似性得分，其定义如下：

注：分别表示老师模型提取的特征、两个特征的内积。类似地，我们定义。

本文所提SEED可以表示为：在所有实例上，最小化老师模型相似性得分与学生模型相似性得分之间的交叉熵。

老师模型是预训练好的，并且在蒸馏阶段进行参数冻结，序列中的特征在训练阶段具有一致性，更高的意味着上更大的权值。由于归一化的存在，与之间的相似性得分在softmax归一化前为常数1，这也是所有中的最大值。因此，的权值最大且可以通过调节进行单独调节。通过最小化上述损失，将与对齐同时与其他不相关特征形成反差。

当时，的softmax平滑接近与one-hot向量，其中，其他向量为0。在这种极限情况下，损失就退化为：

这种退化后的损失与广泛采用的Info-NCE损失相似。

Experiments

Pre-Training

Self-Supervised Pre-training of Teacher Network 我们默认采用MoCo-V2进行老师模型的预训练，采用不同深度/宽度的ResNet作为骨干网络并添加MLP头，最后特征维度为128。考虑到有限算力问题，所有的老师模型预训练200epoch。由于所提方案与老师模型的预训练方案无关，作者同时还给出了其他自监督方案(比如SWAV、SimCLR)的老师模型的结果对比。

Self-Supervised Distillation on Student Network 我们选择了多个小模型作为学生模型：MobileNetV3-Large、EfficientNet-B0、ResNet18、ResNet34等。类似于老师模型，在学生模型的后端添加了MLP头。蒸馏过程采用标准SGD进行训练，momentum=0.9，weight_decay=1e-4，训练了200epoch，初始学习率为0.03，cosine衰减，batch=256，，队列长度K=65536。

Fine-tuning and Evaluation

为验证自监督蒸馏的有效性，我们选择在几个下游任务上验证学生模型的表达能力。首先，我们先来看一下ImageNet上的性能，结果见下表。

从上表结果我们可以看到：

当仅仅采用对比自监督学习(MoCo-V2)时，越小的模型性能越差。比如MobileNetV3-Large的top1精度仅为36.3%，这与MoCoV2中的结论(越大的模型从对比自监督学习中受益越多)相一致。
大模型的蒸馏有助于提升小模型的性能且提升非常明显。比如以MoCoV2预训练的ResNet152作为老师模型，MobileNetV3-Large的top1精度可以提升到61.4%；当采用ResNet-50x2作为老师模型时，其精度还可以进一步提升达到68.2%。
越小的模型从蒸馏得到的受益越多。

上图给出了半监督训练下的性能对比。可以看到：SEED同样有助于提升小模型的性能，而且越强的老师模型会导致越好的学生模型。

上图给出了所提方案下学生模型在分类任务上的迁移能力对比。可以看到：在所有基准数据集上，所提方案均超越了对比自监督预训练方案，验证了SEED的有效证；这同时也证实了通过蒸馏学习到的表达能力的泛化能力。

上表给出了所提方案下学生模型在检测&分割任务上的迁移能力对比。可以看到：

在VOC数据集上，蒸馏得到的预训练模型取得更大的性能提升；
在COCO数据集上，提升相对较小。原因可能是：COCO训练集包含118k数据，而VOC仅有16.5K数据。更大的训练集+更多的微调迭代会降低初始权重的重要性。

Ablation Study

接下来，我们再对所提方案进行一些消融实验分析，包含老师模型大小的影响、对比自监督学习方法的影响、不同蒸馏技术的硬性等等。

首先，我们先看一下老师模型的影响，见下图中Figure5。从中可以看到：随着老师模型的深度/宽度提升，学生模型的性能可以进一步提升。

然后，我们再看一下不同对比自监督学习方法的影响，见上图Table3.从中可以看到：SEED可以轻易将任意自监督模型纳入到自监督蒸馏体系中并得到不同程度的性能提升，此外还可以看到：更多的训练可以得到进一步的性能增益。

其次，我们看一下不同蒸馏策略对于所提方案的影响，见下图Table4。从中可以看到：(1) 简简单单的距离最小即可取得相当好的精度，这证实了将蒸馏纳入自监督学习的有效性；(2) MoCo-V2的监督作为辅助损失不会产生额外的增益，这也就意味着：SEED的损失可以很大程度上覆盖了原始SSL损失，没有必要在蒸馏阶段引入SSL损失；(3)相比其他蒸馏策略，所提SEED取得了最高精度，这也表明学生模型向老师模型对齐+不相关样本上差异化原则的优越性。

最后，我们再来看一下不同超参对于所提方案的影响，见上图Table5。可以看到：当时，所提方案在ImageNet、CIFAR10上取得了更高的性能；而对于CIFAR100而言，最佳的超参为。当比较大时，与队列中的实例的相似性得分也会变大，这就意味着：在某种程度上，相比其他图像的特征，学生模型提取的图像特征判别能力降低了；当为0时，老师模型将生成one-hot向量，仅将视作正例，队列中的其他实例均为负例。因此，最好的参数设置依赖于数据分布。