自监督学习，就看这篇！详解DINO：视觉 Transformer 的自监督学习

极市导读

自监督训练得到的 ViT 包含关于图像语义分割的显式信息，这在以往的有监督训练和卷积网络里面都是不具备的。这些特征也是优秀的 k-NN 分类器，ViT-Small 在 ImageNet 上达到了 78.3% 的 top-1 精度。DINO 还研究了自监督训练中的 momentum encoder，multi-crop training 以及在 ViT 中使用小 Patch 的技巧。 >>加入极市CV技术交流群，走在计算机视觉的最前沿

本文目录

1 DINO：视觉 Transformer 的自监督学习
(来自 Facebook AI Research)
1 DINO 论文解读
1.1 背景和动机
1.2 DINO 算法介绍
1.3 DINO 网络架构
1.4 DINO 训练策略
1.5 ImageNet 实验结果
1.6 消融实验结果

太长不看版

DINO 是视觉 Transformer 做自监督学习的非常经典的工作。DINO 所要探究的问题是：自监督学习算法是否能够为视觉 Transformer 带来新的特性。本文给出了以下的观察：首先，自监督训练得到的 ViT 包含关于图像语义分割的显式信息，这在以往的有监督训练和卷积网络里面都是不具备的。其次，这些特征也是优秀的 k-NN 分类器，ViT-Small 在 ImageNet 上达到了 78.3% 的 top-1 精度。DINO 还研究了自监督训练中的 momentum encoder，multi-crop training 以及在 ViT 中使用小 Patch 的技巧。

1 DINO：视觉 Transformer 的自监督学习

论文名称：Emerging Properties in Self-Supervised Vision Transformers (ICCV 2021)

论文地址：

https//arxiv.org/pdf/2104.14294.pdf

代码地址：

https//github.com/facebookresearch/dino

• 1 DINO 论文解读：

1.1 背景和动机

在 DINO 诞生的时期，视觉 Transformer 与 ConvNet 相比才刚刚具备竞争力，但是那是还没有明显的优势。而且，当时视觉 Transformer 对于计算量要求更高，需要更多的训练数据，且特征没有表现出独特的属性。

因此，本文想要探索一下 Transformer 成功的关键是不是来自 Self-supervised Learning。其主要的动机是 Transformer 在 NLP 中成功的主要原因就是 Self-supervised Learning，比如 BERT 的 Masked Language Modeling 或者 GPT 的 Language Modeling。自监督训练会根据上下文创建一个任务，这些任务相比于有监督训练的预测标签来讲，可以提供更加丰富的学习信号。同样的道理，图像的有监督训练是把图像中包含的丰富的视觉信息压缩到只有类别的标签信息。

因此，本文研究自监督预训练对 ViT 特征的影响。本文给出了几个有趣的观察，这些观察在以往的有监督训练以及卷积网络里面都没有：

• 自监督训练 ViT 得到的特征显式地包含场景布局，尤其是对象的边界。这个信息一般在自监督训练的视觉 Transformer 的最后一个 Block 得到，如图1所示。这个发现作者认为几乎是所有的自监督学习方法的共同特性。
• 自监督训练的视觉 Transformer 的特征可以在不经过任何微调，线性分类器，以及数据增强的前提下，仅仅使用最简单的 k-NN 分类器，就能实现 78.3% 的 top-1 ImageNet 精度。这个发现作者认为并不是所有自监督方法的共性，而是仅仅当包含一些特殊的组件比如 momentum encoder，或者使用了特殊的数据增强比如 multi-crop augmentation 的时候才会出现。

图1: 8×8 Patch 自监督训练的视觉 Transformer 的最后一层的 [CLS] token 的注意力图。

本文提出了一种简单的自监督方法，可以理解为一种没有标签的知识蒸馏方法。得到的结果称为 DINO，它通过使用标准 Cross-Entropy Loss 直接预测教师网络的输出来简化自监督训练。

1.2 DINO 算法介绍

DINO 算法如下图2所示。知识蒸馏是一种学生网络 匹配给定教师网络 的学习范式，其参数分别是 和 。给定输入图片 ，两个网络都输出 维上的概率分布，用 和 表示。概率 是通过使用 Softmax 函数对网络 的输出进行归一化来获得的:

给定一个固定的教师网络 ，通过最小化学生网络 和教师网络 的参数的交叉熵损失来学习匹配这些分布:

式中， 。

首先，作者使用了多种数据增强策略得到不同的试图，从给定的图像中生成一组不同视图 。这个集合包含两个全局视图 和 以及几个分辨率较小的局部视图。所有局部视图输入给学生模型，全局视图输入给教师模型。然后使用下面的损失函数优化:

这种损失是通用的，可以用于任意数量的视图 (公式中是 2)。分辨率为 224×224 的2个全局视图覆盖原始图像的大 (比如大于 50%) 部分区域，分辨率为 96×96 的局部视图仅仅可以覆盖原始图像的小区域 (比如小于 50%)。

教师网络和学生网络共享相同的架构 ，其参数分别是 和 。作者通过使用随机梯度下降最小化方程来学习参数 :

如图2所示是 DINO 的算法流程，一张图片 的两个 View 分别输入给学生模型 和教师模型 。两个网络具有相同的架构但参数不同。教师网络的输出通过一个 Batch 的平均值进行 centering 操作，每个网络输出一个 维的特征，并使用 Softmax 进行归一化。然后使用交叉熵损失作为目标函数计算学生模型和教师模型之间的相似度。在教师模型上使用停止梯度 (stopgradient, sg) 算子来阻断梯度的传播，只把梯度传给学生模型使其更新参数。教师模型使用学生模型的 EMA 进行更新。

图2：DINO 算法流程

DINO 算法伪代码如下。

# gs, gt: student and teacher networks
# C: center (K)
# tps, tpt: student and teacher temperatures
# l, m: network and center momentum rates
gt.params = gs.params
for x in loader: # load a minibatch x with n samples
    x1, x2 = augment(x), augment(x) # random views

    s1, s2 = gs(x1), gs(x2) # student output n-by-K
    t1, t2 = gt(x1), gt(x2) # teacher output n-by-K

    loss = H(t1, s2)/2 + H(t2, s1)/2
    loss.backward() # back-propagate

    # student, teacher and center updates
    update(gs) # SGD
    gt.params = l*gt.params + (1-l)*gs.params
    C = m*C + (1-m)*cat([t1, t2]).mean(dim=0)

def H(t, s):
    t = t.detach() # stop gradient
    s = softmax(s / tps, dim=1)
    t = softmax((t – C) / tpt, dim=1) # center + sharpen
    return – (t * log(s)).sum(dim=1).mean()

DINO 对于教师模型，不像知识蒸馏方法那样直接给定参数，而是从学生网络中进行迭代构建。而且作者观察到，在整个训练过程中，教师模型比学生模型有更好的表现，因此，通过提供更高质量的目标特征来指导学生的训练。

1.3 DINO 网络架构

DINO 的模型 是由一个 Backbone (ResNet 或者 ViT) 和一个Projection head 组成。Projection head 就是一个3 层的 MLP。DINO 不使用 BYOL 中的 predictor，使得教师和学生的架构完全一致。而且，作者注意到与标准 ConvNet 不同，ViT 架构默认不使用批量归一化 。因此，当将 DINO 应用于 ViT 时，我们也在投影头中不使用任何 ，使系统完全达到 -free 的架构。

通过下面这行代码：

t = softmax((t - C) / tpt, dim=1) # center + sharpen


可以发现，对于教师模型的操作只有 center + sharpen。center 可以稳定 momentum 带来的影响，但是会使得教师模型的输出趋近于均匀分布，而 sharpen 正好可以解决这个问题。center 的超参数 可以通过下式来计算:

其中， 是比例参数。sharpen 是通过把知识蒸馏中的温度系数 设置得比较低来实现。

1.4 DINO 训练策略

DINO 中的 ViT 架构使用 DeiT 的实现，Patch Size 使用8或者16。DINO 也像 ViT 那样给模型加了一个 [CLS] token，即使它没有附加到任何标签或监督信息。Self-attention 机制更新这个 token。

在评估时，除了 linear evaluation 和 finetuning evaluation 之外，作者还额外尝试了 k-NN evaluation。作者冻结了预训练模型，然后存储下游任务训练数据的特征。

1.5 ImageNet 实验结果

如下图3所示是跨方法层面和跨架构层面的 ImageNet 实验结果对比。ViT-S 的选择的动机是它与 ResNet-50 沿几个轴的相似性：参数数量 (21M vs 23M)，吞吐量 (1237imgs/sec VS 1007 imgs/sec) 和 ImageNet 上的有监督训练的性能 (79.3% VS 79.8%)。首先，作者观察到 DINO 在 ResNet-50 上的表现与最先进的几个方法相当，验证了 DINO 作为标准自监督学习方法的性能。当切换到 ViT 架构时，DINO 在线性分类方面优于 BYOL、MoCov2 和 SwAV +3.5%，在 k-NN 评估中优于 +7.9%。更令人惊讶的是，使用简单的 k-NN 分类器的性能几乎与 Linear Classifier 相当 (74.5% 和 77.0%)。值得一提的是这个属性仅在使用带有 ViT 架构的 DINO 时出现，并且不会出现在其他现有的自监督方法和 ResNet-50 中。

图3：ImageNet 实验结果

1.6 消融实验结果

不同组件的作用

作者在下图4中报告了当添加或删除组件时的不同模型的性能。首先可以观察到，在没有 Momentum 的情况下，DINO 框架不起作用。SK 操作也可以避免坍塌的现象。此外，有了 Momentum 的情况下，使用 SK 几乎没有影响。

比较第 3 行和 9 行：突出了 Momentum 对性能提升的重要性。

比较第 4 行和 5 行：DINO 中的 Multi-crop 训练和 CE Loss 是获得良好特征的重要组成部分。

比较第 6 行：向学生网络添加 Predictor 几乎没有影响。

图4：不同组件的作用

Patch Size 的作用

在下图5中，作者比较了不同 Patch size (16×16, 8×8 和 5×5) 的 ViT 模型的 k-NN 分类性能。所有模型都训练了 300 个 Epoch。可以观察到，当减小 Patch size 的大小时，性能会大大提高。有趣的是，在不增加额外参数的情况下可以大大提高性能。然而，使用较小 Patch size 的性能增益是以牺牲吞吐量为代价的。比如，当使用 5×5 的 Patch size 时，吞吐量下降到 44 im/s，而 8×8 Patch size 的吞吐量下降到 180 im/s。

教师模型训练策略的作用

如下图5所示作者比较了不同的教师模型训练策略，比如：

Previous epoch：使用来自前一个 Epoch 的学生模型作为教师模型。Previous iter：使用来自前一个 Iteration 的学生模型作为教师模型。Student copy：使用学生模型作为教师模型。

图5：教师模型训练策略的作用

可以观察到使用来自前一个 Epoch 的教师模型得到的性能并不会崩溃，且与 MoCo-v2 或 BYOL 等现有框架提供有竞争力的 k-NN 评估性能。

避免坍塌

作者从交叉摘分解的角度来验证实验现象: 交叉熵 可以被分解为熵 和 KL 散度 ，如下是所示:

散度 为0意味着 constant output，如下图6所示是教师输出的熵 和 KL 散度 在训练过程中的演变。可以发现，当只使用 sharpening 或者 centering 策略时，KL 散度都直接为 0 ，代表崩溃。当只使用 sharpening 时，教师输出的熵为 0 ；当只使用 centering 时，教师输出的摘为 。表明这两种操作都会导致不同形式的崩溃。

图6：坍塌现象研究