【ICLR2021自监督蒸馏】SEED: Self-supervised Distillation For Visual Representation论文精读

论文地址：SEED: Self-supervised Distillation For Visual Representation
论文代码：无

摘要

  本文研究小模型的自监督学习。作者的实证研究发现，虽然广泛使用的对比自监督学习方法在大模型训练方面取得了很大的进步，但它在小模型训练中效果不佳。为了解决这个问题，作者提出了一种新的学习范式，称为自监督蒸馏（SEED），利用一个更大的网络（作为教师）以自监督的方式将其代表性知识转移到一个更小的架构（作为学生）中。不是直接从未标记的数据中学习，而是训练学生编码器来模拟教师在一组实例上推断的相似性分数分布。目的是促进在预训练阶段使用蒸馏技术对小模型进行视觉表征学习。SEED显著提高了小型网络在下游任务上的性能，与自监督基线相比，SEED在ImageNet-1k数据集上的top-1准确性从EfficientNet-B0的42.2%提高到67.6%，在MobileNetV3-Large上从36.3%提高到68.2%。

背景

  自监督学习（SSL）的成功主要体现在其从大规模未标记数据中学习的非凡能力。SSL的另一个显著优点是，它可以获得与任务无关的表示，同时允许训练免除昂贵的数据标记。视觉SSL的主要分支包括pretext task、对比表征学习、在线/离线聚类等。其中，最近的一些工作在转移到下游任务（如半监督分类、目标检测）时，取得了与监督预训练相当甚至更好的精度。然而上述性能最好的SSL算法都涉及大型网络（例如ResNet-50或更大），却很少关注小型网络。作者发现现有的对比学习等技术在小型网络上效果不佳。如下图，结合MoCo-V2在ImageNet上使用MobileNet-V3-Large的top-1精度仅为36.3%，远远低于其有监督训练的准确率75.2%；对于EfficientNet-B0，准确率为42.2%，而其对应的有监督训练的准确率为77.1%。作者推测这是因为参数较少的较小模型无法有效地学习具有大量数据的实例级区分表示。

  为了解决这一问题，作者将知识蒸馏（KD）注入到自监督学习中，并提出自监督蒸馏（SEED）作为一种新的学习范式。也就是说，以自监督的方式训练较大的模型，并蒸馏到较小的模型。SEED没有直接在较小的模型上进行自监督训练，而是首先以自监督的方式训练一个较大的模型（作为老师），然后将知识蒸馏到较小的模型（作为学生）。请注意，传统的蒸馏是有监督学习的方式，而这里的蒸馏是在自监督设置下进行的，没有任何标记数据。有监督蒸馏可以表述为训练学生去模仿教师模型预测的类上概率质量函数。然而，在无监督的知识蒸馏环境中，类上的分布不是直接可以实现的。因此，作者提出了一种简单而有效的自监督蒸馏方法。与何凯明大大的MoCo论文的方法类似，作者也维护了一个数据样本队列。给出一个实例，首先使用教师网络获得其与队列中所有数据样本以及实例本身的相似度得分。然后训练学生编码器模拟教师在这些数据样本上推断的相似性分数分布。

贡献

• 它不需要任何聚类/典型的计算过程来检索伪标签或潜在类；
• 教师模型可以使用任何高级SSL方法进行预培训，例如MoCo-V2、SimCLR、SWAV等；
• 这些知识可以蒸馏到任何目标小型网络（较浅、较薄或完全不同的体系结构）。

方法

  自监督蒸馏方法如下图。教师编码器通过SSL进行预训练，并在蒸馏过程中冻结。学生编码器是在动态维护的队列上计算一个图像的增强视图，通过最小化教师和学生的概率交叉熵来训练的。即该实例队列用于存储来自教师编码器关于数据样本的编码输出。给定一个新样本，使用教师和学生模型计算其与队列中所有样本的相似度得分，并要求学生计算的相似性分数分布与教师计算的相似性分数分布相匹配，其公式为最小化学生和教师的相似性分数分布之间的交叉熵。

  具体而言，对于图像的随机增强视图x_i，首先将其映射并归一化为特征向量表示:

  其中z_Ti，z_Si∈ R^D、f^T_θ和f^S_θ分别表示教师和学生编码器。设D=[d₁，…，d_K]表示实例队列，其中K是队列长度，d_j是从教师编码器获得的特征向量。与对比学习框架类似，D在“先进先出”策略下随着蒸馏的进行而逐步更新。也就是说，将教师编码的当前batch的视觉特征进入队列，并在迭代结束时将最早的样本出队。另外，队列D中维护的样本大多是随机的，与目标实例x_i无关。最小化学生和教师基于D计算的相似度得分分布之间的交叉熵，将x_i与随机选择的样本进行软对比，而不直接与教师编码器对齐。为了解决这个问题，我们将教师嵌入（z_Ti）添加到队列中，形成D⁺=[d₁,…,d_K，d_K+1]，其中d_K+1=z_Ti。
  p^T(x_i；θ^T；D⁺)表示提取的教师特征z_Ti和教师模型计算的d_j(j=1，…，K+1)之间的相似性得分：

其中T^T，T^S都表示温度超参数，使用()^T、()^S分别表示教师、学生网络中的特征，并使用（·）表示两个特征之间的内积。学生模型同理：

自监督蒸馏可以表示为在所有实例x_i上最小化教师的相似性分数p^T(x_i；θ^T；D⁺)和学生的相似性分数p^S(x_i；θ^S；D⁺)之间的交叉熵，即把上述（2）和（3）式代入得：

以此来更新学生模型中的参数。
  由于教师网络是预先训练和冻结的，因此在训练过程中队列中的特征是一致的。p^T_j值越高，施加在p^S_j上的权重越大。由于l₂归一化，在softmax归一化之前，z_Ti和d_K+1之间的相似度得分保持1不变，这在p^T_j中是最大的。因此，p^S_K+1的权重是最大的，可以通过调整τ^T的值来单独调整。通过最小化损失，z_Si的特征可以与z_Ti对齐，同时与D中其他不相关的图像特征进行对比。
与Info-NCE loss的关系。当τ^T→ 0时，p^T的softmax函数平滑地接近一个one-hot向量，其中p^T_K+1等于1，p^T中其他所有值均为0。在这种极端情况下，损失变成：

即上述（4）式中前k项均为0（因为0乘任何数均为0），仅计算最后一项，此时p^T(x_i；θ^T；D⁺)=[p₁^T,…,p_k^T]=[0,…0,1]。

实验

  表中数据显示了各种教师网络和学生网络的结果。在第一行为每个学生架构列出了使用MoCo-V2进行对比自监督预训练的基线。可以清楚地看到，较小的网络性能相当差。例如MobileNet-V3只能达到36.3%。这与先前的结论很好地吻合，即需要更大的模型才能在基于对比的自监督预训练中表现更好。作者推测，这主要是由于较小的网络无法区分大规模数据集中的实例。结果还清楚地表明，从较大网络中蒸馏有助于提高小网络的性能，并显示出明显的改善。
  另外。图中展示了ImageNet-1k上Top-1半监督评估准确度使用1%（红线）、10%（蓝线）标记进行线性微调，与SEED的完全有监督（绿线）线性评估基线相比。对于教师模型参数为0的点，作者给出了MoCo-V2在没有任何蒸馏的情况下的半监督线性评价结果。在从较大的教师那里提炼出来后，学生模型往往在半监督任务上表现更好。

其它实验如迁移到下游任务的实验结果图可见原论文，这里就不一一赘述了。

消融实验（丰富，推荐阅读原论文）

• 不同的教师网络结构；
• 不同的教师模型预训练算法；
• 不同的蒸馏策略；
• 超参数设置。

  在表3中，作者展示使用不同的预训练算法从ResNet-50中提取的ResNet-18的Top-1精度，即MoCo-V1、MoCo-V2、SimCLR和SWAV。值得注意的是，上述方法都一致采用了基于对比的预训练方法，这是一种基于在线聚类的方法。作者发现SEED对预训练方法是不可知的，因此很容易在自监督蒸馏中使用任何自监督模型（包括基于聚类的方法，如SWAV）。此外，观察到教师SSL和蒸馏阶段的更多训练阶段都可以带来有益的收益。图5总结了从更广泛和更深入的ResNet体系结构中提取的ResNet-18和EfficientNet-B0的准确性。随着教师网络深度和宽度的增加，可以看到明显的性能改善：与ResNet-50相比，更深（ResNet-101）和更宽（ResNet-50×2）显著提高了准确性。然而，进一步的架构扩展的影响相对有限，作者认为在这种情况下，准确性可能会受到学生网络容量的限制。

  探讨了几种替代蒸馏策略。l₂距离：教师和学生嵌入的l₂距离最小化。K-Means：利用K-Means聚类来根据教师网络的表示分配一个伪标签。在线聚类：在蒸馏过程中不断更新聚类中心以生成伪标签。二元对比损失：采用信息相似损失进行对比蒸馏。作者在原论文A.4中提供了其他策略的详细信息。表4显示了从ResNet-50中提取的ResNet-18（学生）上每种方法的结果。从结果来看，简单的l₂距离最小化方法可以达到较高的精度，这证明了将蒸馏思想应用于自监督学习的有效性。除此之外，作者还研究了原始SSL（MoCo-V2）监管对SEED的补充损失的影响，发现它不会给蒸馏带来额外的好处。发现这两种策略（Top-1线性Acc）的结果相近，SEED达到57.9%，而SEED+MoCo-V2达到57.6%。这意味着SEED的损失在很大程度上可以覆盖原始SSL的损失，并且在蒸馏过程中无需进一步进行SSL。同时，作者提出的SEED以最高的准确性优于这些备选方案，这表明了将学生与老师对齐并与无关样本进行对比的优势。