Skip to content

DINO

摘要

在本文中,我们探讨了自监督学习是否为视觉变换器(ViT)[19] 提供了与卷积网络(convnets)相比有显著区别的新属性。除了将自监督方法适配到这种架构上特别有效之外,我们还做出以下观察:首先,自监督 ViT 特征包含了关于图像语义分割的显式信息,这在有监督的 ViTs 以及 convnets 中并不明显。其次,这些特征还是优秀的 k-NN 分类器,在 ImageNet 上使用小型 ViT 达到了 78.3% 的 top-1 准确率。我们的研究还强调了动量编码器 [33]、多裁剪训练 [10] 以及在 ViTs 中使用小块的重要性。我们将我们的发现实现到一个简单的自监督方法中,称为 DINO,我们将其解释为一种无标签的自我蒸馏形式。我们展示了 DINO 和 ViTs 之间的协同作用,通过使用 ViT-Base 在 ImageNet 的线性评估中达到了 80.1% 的 top-1 准确率。

引言

最近,变换器(Transformers)[70] 作为一种替代卷积神经网络(convnets)用于视觉识别的方法出现了 [19, 69, 83]。它们的采用与一种受自然语言处理(NLP)启发的训练策略结合在一起,即在大量数据上进行预训练,并在目标数据集上进行微调 [18, 55]。由此产生的视觉变换器(ViT)[19] 与 convnets 具有竞争力,但它们尚未提供超越后者的明显好处:它们在计算上要求更高,需要更多的训练数据,且其特征没有展示出独特的属性。

在本文中,我们探讨了 Transformers 在视觉领域中受到限制的成功是否可以通过它们在预训练中使用监督学习来解释。我们的动机是,Transformers 在 NLP 中成功的主要因素之一是使用了自监督预训练,以 BERT[18] 中的闭环程序或 GPT[55] 中的语言建模的形式。这些自监督预训练目标利用句子中的单词创建前置任务,提供比预测单个标签的监督目标更丰富的学习信号。类似地,在图像中,图像级别的监督经常将包含在图像中的丰富视觉信息简化为从预定义的几千个对象类别中选出的单一概念 [60]。

虽然 NLP 中使用的自监督前置任务是特定于文本的,但许多现有的自监督方法已经展示了它们在使用 convnets 的图像上的潜力 [10, 12, 30, 33]。它们通常具有类似的结构,但设计了不同的组件来避免平凡解(崩溃)或改善性能 [16]。在这项工作中,受到这些方法的启发,我们研究了自监督预训练对 ViT 特征的影响。特别值得关注的是,我们已经发现了几个有趣的属性,这些属性既不在有监督的 ViTs 中出现,也不在 convnets 中出现:

• 自监督 ViT 特征显式包含场景布局,特别是对象边界,如图 1 所示。这些信息可以直接在最后一个块的自注意力模块中访问。

• 在没有任何微调、线性分类器或数据增强的情况下,自监督 ViT 特征通过一个基本的最近邻分类器(k-NN)表现特别好,达到了 ImageNet 上 78.3% 的 top-1 准确率。

分割掩码的出现似乎是自监督方法共有的属性。然而,与 k-NN 的良好性能只有在结合某些组件,如动量编码器 [33] 和多重裁剪增强 [10] 时才会显现。我们研究中的另一个发现是使用更小的 ViT 块来提高结果特征质量的重要性。

总的来说,我们关于这些组件重要性的发现导致我们设计了一个简单的自监督方法,可以被解释为一种无标签的知识蒸馏 [35] 形式。由此产生的框架,DINO,通过直接预测一个建立在动量编码器上的教师网络的输出,使用标准的交叉熵损失来简化自监督训练。

有趣的是,我们的方法只需对教师输出进行中心化和锐化处理以避免崩溃,而其他流行的组件,如预测器 [30]、高级归一化 [10] 或对比损失 [33] 在稳定性或性能方面的好处很少。特别重要的是,我们的框架灵活,可以在不需要修改架构或调整内部归一化 [58] 的情况下,同时适用于卷积网络和 ViT。

我们进一步通过在 ImageNet 线性分类基准测试上超越以前的自监督特征,使用小块的 ViT-Base 达到 80.1% 的 top-1 准确率,验证了 DINO 与 ViT 之间的协同作用。我们还确认 DINO 适用于卷积网络,通过与 ResNet-50 架构匹配达到了最新水平。最后,我们讨论了在计算和内存容量有限的情况下使用 DINO 与 ViT 的不同场景。特别是,使用 ViT 训练 DINO 只需两个 8-GPU 服务器 3 天就能在 ImageNet 线性基准测试上达到 76.1%,这超过了基于大小相当的卷积网络的自监督系统,显著减少了计算需求 [10, 30]。

image.png

图 2:无标签的自我蒸馏。为了简单起见,我们在单个视图对 \(\left(x_1, x_2\right)\) 的情况下示例化 DINO。模型将输入图像的两个不同随机变换传递给学生和教师网络。两个网络具有相同的架构但参数不同。教师网络的输出通过批处理上计算的平均值进行中心化。每个网络输出一个维度为 \(K\) 的特征,该特征通过特征维度上的温度 softmax 进行归一化。然后通过交叉熵损失测量它们的相似性。我们在教师上应用了停止梯度(sg)操作符,以仅通过学生传播梯度。教师参数通过学生参数的指数移动平均(ema)进行更新。

相关工作

自监督学习。大量关于自监督学习的工作集中于被称为实例分类 [12, 20, 33, 73] 的判别方法,该方法将每张图像视为不同的类别,并通过区分它们以及数据增强来训练模型。然而,显式学习一个分类器来区分所有图像 [20] 随着图像数量的增加并不易于扩展。Wu 等人 [73] 提出使用噪声对比估计器(NCE)[32] 来比较实例而不是对它们进行分类。这种方法的一个缺点是它需要同时比较大量图像的特征。实践中,这需要大批量 [12] 或内存库 [33, 73]。几种变体允许以聚类形式 [2, 8, 9, 36, 42, 74, 80, 85] 自动分组实例。

近期的工作表明,我们可以在不区分图像的情况下学习无监督特征。特别值得注意的是,Grill 等人 [30] 提出了一种称为 BYOL 的度量学习公式,其中特征通过与动量编码器获得的表示相匹配来训练。像 BYOL 这样的方法即使没有动量编码器也能工作,但代价是性能下降 [16, 30]。

其他几项工作也反映了这一方向,显示出可以匹配更精细的表示 [26, 27],通过将特征匹配到均匀分布 [6] 或使用白化 [23, 81] 来训练特征。我们的方法从 BYOL 中汲取灵感,但采用不同的相似性匹配损失,<u>并且学生和教师使用完全相同的架构。通过这种方式,我们的工作补充了 BYOL 对自监督学习作为无标签的 Mean Teacher 自我蒸馏 [65] 的解释。</u>

自我训练和知识蒸馏。自我训练旨在通过将少量初始标注传播到大量未标注实例来提高特征质量。这种传播可以通过硬标签分配 [41, 78, 79] 或软分配 [76] 来完成。当使用软标签时,这种方法常被称为知识蒸馏[7, 35],主要设计用于训练小网络模仿大网络的输出以压缩模型。Xie 等人 [76] 已经展示了蒸馏可以用来在自我训练流程中将软伪标签传播到未标记数据,画出了自我训练和知识蒸馏之间的重要联系。我们的工作基于这种关系,并将知识蒸馏扩展到没有标签可用的情况。以前的工作也结合了自监督学习和知识蒸馏 [25, 63, 13, 47],实现了自监督模型压缩和性能提升。然而,这些工作依赖于预训练的固定教师,而我们的教师在训练过程中动态构建。这样,知识蒸馏不是作为自监督预训练的后处理步骤,而是直接作为一种自监督目标。最后,我们的工作也与共蒸馏 [1] 相关,其中学生和教师具有相同的架构并在训练期间使用蒸馏。然而,在共蒸馏中,教师也从学生那里获得蒸馏,而在我们的工作中,教师是通过学生的平均值更新的

方法

带有知识蒸馏的自监督学习

我们在这项工作中使用的框架,DINO,与近期的自监督方法 [10, 16, 12, 30, 33] 共享相同的总体结构。然而,我们的方法也与知识蒸馏 [35] 有相似之处,并且我们从这个角度来介绍它。我们在图 2 中展示了 DINO,并在算法 1 中提供了伪代码实现。

知识蒸馏是一种学习范式,其中我们训练一个学生网络 \($g_{\theta_s}$\) 来匹配给定教师网络 \($g_{\theta_t}$\) 的输出,分别由 \($\theta_s$\)\($\theta_t$\) 参数化。给定输入图像 \(x\),两个网络都输出在 \($K$\) 维上的概率分布,分别表示为 \($P_s$\)\($P_t$\)。概率 \($P$\) 通过使用 softmax 函数归一化网络 \(g\) 的输出来获得。更具体地说,

\[ P_s(x)^{(i)}=\frac{\exp \left(g_{\theta_s}(x)^{(i)} / \tau_s\right)}{\sum_{k=1}^K \exp \left(g_{\theta_s}(x)^{(k)} / \tau_s\right)}, \]

其中 \($\tau_s>0$\) 是一个温度参数,控制输出分布的锐度,对于 \($P_t$\) 有类似的公式,但温度为 \($\tau_t$\)。给定一个固定的教师网络 \($g_{\theta_t}$\),我们通过最小化学生网络参数 \($\theta_s$\) 的交叉熵损失来学习匹配这些分布:

\[ \min _{\theta_s} H\left(P_t(x), P_s(x)\right), \]

其中 \(H(a, b)=-a \log b\)

接下来,我们详细介绍如何将等式 (2) 中的问题适应到自监督学习。首先,我们使用多裁剪策略 [10] 构造图像的不同扭曲视图或裁剪。更具体地说,从给定图像中,我们生成一组不同的视图 \(V\)。这个集合包含两个全局视图 \($x_1^g$\)\($x_2^g$\) 以及几个分辨率较小的局部视图。所有裁剪都通过学生网络,而只有全局视图通过教师网络,因此鼓励“局部到全局”的对应关系。我们最小化损失:

\[ \min _{\theta_s} \sum_{x \in\left\{x_1^g, x_2^g\right\}} \sum_{\substack{x^{\prime} \in V \\ x^{\prime} \neq x}} H\left(P_t(x), P_s\left(x^{\prime}\right)\right) . \]

这种损失是通用的,可以用于任意数量的视图,甚至仅 2 个。然而,我们遵循多裁剪标准设置,使用 2 个分辨率为 \(224^2\) 的全局视图,覆盖原始图像的大面积(例如大于 \(50\%\)),以及几个分辨率为 \(96^2\) 的局部视图,仅覆盖原始图像的小面积(例如小于 \(50\%\))。除非另有说明,否则我们将这种设置称为 DINO 的基本参数化。

两个网络共享相同的架构 \(g\),但有不同的参数集 \(\theta_s\)\(\theta_t\)。我们通过最小化方程(3)以随机梯度下降的方式学习参数 \(\theta_s\)

教师网络 与知识蒸馏不同,我们没有预先给定的教师 \(g_{\theta_t}\),因此,我们从学生网络的过去迭代中构建它。我们在第 5.2 节研究教师的不同更新规则,并显示在我们的框架中,冻结一个时期的教师网络效果出奇的好,而复制学生权重给教师则无法收敛。特别值得注意的是,使用学生权重的指数移动平均(EMA),即动量编码器 [33],特别适合我们的框架。更新规则是 \(\theta_t \leftarrow \lambda \theta_t+(1-\lambda) \theta_s\),其中 \(\lambda\) 在训练期间 [30] 按余弦计划从 0.996 变化到 1。最初,动量编码器被引入作为对比学习中队列的替代品 [33]。然而,在我们的框架中,它的角色不同,因为我们既没有队列也没有对比损失,其角色可能更接近自训练中使用的平均教师 [65]。事实上,我们观察到这种教师执行了一种类似于 Polyak-Ruppert 平均的模型集成方法,具有指数衰减 [51, 59]。使用 PolyakRuppert 平均进行模型集成是改进模型性能的标准做法 [38]。我们观察到,这种教师在整个训练过程中的性能都优于学生,因此,通过提供更高质量的目标特征来指导学生的训练。这种动态在以前的工作中没有被观察到 [30,58]。

网络架构 神经网络 \(g\) 由一个主干网络 \(f\)(ViT [19] 或 ResNet [34])和一个投影头 \(h: g=h \circ f\) 组成。在下游任务中使用的特征是主干网络 \(f\) 的输出。投影头由一个 3 层的多层感知机(MLP)组成,隐藏维度为 2048,后接 \(\ell_2\) 归一化和一个权重归一化的全连接层 [61],其维度为 \(K\),这与 SwAV [10] 的设计类似。我们测试了其他投影头,这个特定的设计对于 DINO 来说似乎效果最好(附录 C)。我们不使用预测器 \([30,16]\),结果在学生和教师网络中采用完全相同的架构。特别值得注意的是,与标准的卷积网络不同,ViT 架构默认不使用批量归一化(BN)。因此,在将 DINO 应用到 ViT 时,我们也不在投影头中使用任何 BN,使得系统完全没有 \(BN\)

避免崩溃 几种自监督方法通过使用不同的操作来避免崩溃,包括对比损失 [73]、聚类约束 [8, 10]、预测器 [30] 或批量归一化 \([30,58]\)。虽然我们的框架可以通过多种归一化 [10] 来稳定,但也可以仅通过对动量教师输出进行中心化和锐化来避免模型崩溃。如第 5.3 节的实验所示,中心化可以防止一个维度占主导地位,但鼓励向均匀分布崩溃,而锐化则有相反的效果。同时应用这两种操作可以平衡它们的效果,这足以在有动量教师的情况下避免崩溃。选择这种方法来避免崩溃,是为了减少对批量的依赖而牺牲稳定性:中心化操作仅依赖于批量的一阶统计量,并可以解释为向教师添加一个偏差项 \(c\)\(g_t(x) \leftarrow g_t(x)+c\)。中心 \(c\) 通过指数移动平均更新,这使得该方法能够在不同批量大小下良好工作,如第 5.5 节所示:

\[ c \leftarrow m c+(1-m) \frac{1}{B} \sum_{i=1}^B g_{\theta_t}\left(x_i\right), \]

其中 \(m>0\) 是一个比率参数,\(B\) 是批量大小。通过在教师 softmax 归一化中使用低温度值 \(\tau_t\) 来获得输出锐化。

实施和评估协议

在本节中,我们提供了使用 DINO 训练的实现细节,并介绍了我们实验中使用的评估协议。

视觉变换器。我们简要描述了视觉变换器(ViT)[19, 70] 的机制,并引用 Vaswani 等人 [70] 关于变换器的详细信息以及 Dosovitskiy 等人 [19] 对其图像适应性的描述。我们遵循在 Dei T [69] 中使用的实现。我们在表 1 中总结了本文中使用的不同网络的配置。ViT 架构接受一个由非重叠连续图像块组成的网格作为输入,分辨率为 \(N \times N\)。在本文中,我们通常使用 \(N=16\)(“/16”)或 \(N=8\)(“/8”)。然后通过线性层传递这些块以形成一组嵌入。我们在序列中额外添加了一个可学习的标记 [18,19]。这个标记的作用是从整个序列中聚合信息,我们在其输出处附加了投影头 \(h\)。我们将这个标记称为类标记 [CLS],以与之前的作品 \([18,19,69]\) 保持一致,尽管在我们的案例中它没有附加任何标签或监督。补丁标记集和 [CLS] 标记被送入一个标准的变换器网络,其中包含一个“预规范化”的层规范化 [11, 39]。变换器是一系列自注意力和前馈层的序列,与跳过连接并行。自注意力层通过查看其他标记表示与注意力机制 [4] 来更新标记表示。

实现细节。我们在没有标签的 ImageNet 数据集 [60] 上预训练模型。我们使用 adamw 优化器 [44] 进行训练,批量大小为 1024,在使用 ViT-S/16 时分布在 16 个 GPU 上。学习率在前 10 个周期内线性增加到其基础值,该基础值通过以下线性缩放规则确定 [29]:\(lr=0.0005 *\) 批量大小 \(/ 256\)。在这个热身之后,我们用一个余弦计划 [43] 衰减学习率。权重衰减也遵循从 0.04 到 0.4 的余弦计划。\(\tau_s\) 的温度设置为 0.1,而我们对 \(\tau_t\) 使用从 0.04 到 0.07 的线性热身,在前 30 个周期内。我们遵循 BYOL[30] 的数据增强方法(颜色抖动、高斯模糊和日光化)和多裁剪 [10],使用双三次插值来适应尺度 \([19,69]\) 的位置嵌入。我们的代码和模型是公开可用的,以复现我们的结果。评估协议。自监督学习的标准协议是要么在冻结特征上学习线性分类器 [82,33],要么在下游任务上微调特征。对于线性评估,我们在训练期间应用随机大小裁剪和水平翻转增强,并在中心裁剪上报告准确性。对于微调评估,我们用预训练的权重初始化网络,并在训练期间调整它们。然而,这两种评估都对超参数敏感,例如,我们观察到在改变学习率时,准确性之间存在很大的方差。因此,我们还使用简单的加权最近邻分类器 \((k-NN)\) 来评估特征的质量,如 [73] 所示。我们冻结预训练模型以计算并存储下游任务训练数据的特征。然后,最近邻分类器将图像的特征与为标签投票的 \(k\) 个最近存储特征匹配。我们在不同数量的最近邻居上进行扫描,发现对于我们的大多数运行,\(20NN\) 始终表现最好。这种评估协议不需要任何其他超参数调整,也不需要数据增强,且只需对下游数据集进行一次传递,极大地简化了特征评估。