DINODINO v2：颠覆自监督视觉特征表示学习

DINO

Emerging Properties in Self-Supervised Vision Transformers

导读

DINO 是一种用于自监督视觉学习的深度学习模型，于 2021 年由 Facebook AI 提出。DINO 是最先探讨基于 Transformer 架构的自监督学习代表作之一，其通过在无标签图像上进行自监督训练来学习视觉特征表示。

相比于监督学习需要人为地提供标签告诉模型这是什么，自监督学习无须任何“显示”标签，只需输入图像通过某种机制便能让网络学会理解图像本身的语义信息。例如，我们可以通过图像旋转、随机裁剪等变换，使模型学习到不同角度、不同尺度下的特征。另外，也可以通过模型自身的预测来构建任务，例如预测图像的局部块、颜色等等。这些任务由于不需要人类标注，因此可以在大规模的无标签数据上进行训练，从侧面上提供了一种从无标签数据中学习特征表示的新方法，它可以帮助我们更好地利用现有数据资源，解决监督学习中的一些瓶颈问题。

回到正题，DINO 的核心思想便是通过在大规模的无标签数据集上进行对比学习，学习出一组具有可传递性的视觉特征表示。在 DINO 中，作者通过引入一个新的对比学习方法，将原始图像的特征与随机裁剪的图像的特征进行对比，从而学习到更好的视觉通用表征，最终也获得了非常出色的效果。

DINO 这个名称可以理解为是由 Distillation 和 NO labels 这两个词组成的缩写，既表达了DINO采用自蒸馏方法的特点，也突出了它是一种基于无监督学习的模型。具体来说，DINO 是使用一种称为“无监督自蒸馏”的方法，该方法通过自监督学习来学习模型的知识表示。在这个方法中，模型使用自身的输出来生成“伪标签”，然后使用这些伪标签来重新训练模型，从而进一步提高模型的性能和泛化能力。

方法

作为开始，我们给出一张动图，其非常生动形象的展示了贯穿 DINO 的整个框架和核心思想：

如上所示，DINO 本质上是一种自监督学习方法，通过无监督的方式学习图像特征表示，可用于计算机视觉的其他下游任务，例如分类和检测等。该方法的核心思想是使用一种叫做自蒸馏的方法，即将一个学生模型的表示与一个动量化的教师模型的表示进行比较，以学习出更好的特征表示。

在正式讲解具体细节前，我们可以先看看 DINO 整个处理流程的伪代码：

下面我们将分别从网络结构、数据增强、损失函数三大部分进行详细的介绍。

A. 网络结构
正如我们上面提到过的，DINO 是采用自蒸馏(self-distillation)的方法学习的，其整体框架包含两个相同的架构，分别为教师网络和学生网络，具体的架构可以是 ViT 等 vision transformer 或者诸如 ResNet 等 CNNs 特征提取器，非常灵活方便。当然，通过下述消融实验也知道还是 ViT 的潜力更大。

然而，这种学生和教师网络均输出相同 embeddings 的情况容易出现模式崩塌(mode collapse)的现象。在《Momentum Contrast for Unsupervised Visual Representation Learning》一文中提出了一种解决方案，即应用“动量教师”(momentum tearcher)模型，可以简单地理解为就是教师的模型不是基于反向传播更新的，而是再对学生模型进行梯度回传后，在通过指数移动平均(Exponentially Weighted Average, EWA)，直接将学生网络学习到的模型参数更新给教师网络，换句话就是教师网络的权重更新自学生网络。

DINO 中便是沿用这种方式。具体地，我们可以简单看下教师权重的更新公式：

这里下标 ${\theta }_t$ 和 ${\theta }_s$ 分别指代教师和学生网络对应的模型参数; 而 $\lambda$ 则跟随余弦学习率衰减策略在训练过程中从 0.996 到 1 之间进行变化。

B. 数据增强
DINO 中最核心的数据采样策略便是图像裁剪，这也是自监督学习领域应用非常广泛的主策略之一。一般来说，我们可以将裁剪后的图像分为两种：

• Local views: 即局部视角，也称为 small crops，指的是抠图面积小于原始图像的 50%；
• Global views: 即全局视角，也称为 large crops，指的是抠图面积大于原始图像的 50%；

在 DINO 中，学生模型接收所有预处理过的 crops 图，而教师模型仅接收来自 global views 的裁剪图。据作者称，这是为了鼓励从局部到全局的响应，从而训练学生模型从一个小的裁剪画面中推断出更广泛的上下文信息。

简单来说，就是把局部特征和全局特征分别交给不同的模型来学习，以便在处理整个图像时，能够更好地对局部细节和上下文进行综合判断。

此外，为了使网络更加鲁邦，DINO 中也采用一些其它的随机增强，包括：颜色扰动(color jittering) 、 高斯模糊(Gaussian blur) 和 曝光增强(solarization)

C. 损失函数
在 DINO 中，教师和学生网络分别预测一个一维的嵌入。为了训练学生模型，我们需要选取一个损失函数，不断地让学生的输出向教师的输出靠近。softmax 结合交叉熵损失函数是一种常用的做法，来让学生模型的输出与教师模型的输出匹配。具体地，通过 softmax 函数把教师和学生的嵌入向量尺度压缩到 0 到 1 之间，并计算两个向量的交叉熵损失。这样，在训练过程中，学生模型可以通过模仿教师模型的输出来学习更好的特征表示，从而提高模型的性能和泛化能力。

当然，这也可以看作是一个分类问题，以便网络可以从局部视图中学习更有意义的全局表示。

Centering and Sharpening
在 DINO 论文中，还有两个不得不提的点便是 Centering 和 Sharpening，这是用于防止模式崩塌的两种有效方式。

在自监督学习中，mode collapse是指网络的学习过程中出现了多样性减少的现象。具体来说，当网络学习到一组特征表示时，往往会出现多个输入数据映射到相同的特征表示的情况，这就是所谓的模式崩塌。这种现象通常是由于网络在优化过程中陷入了局部最优解，只能考虑到一部分数据的特征表示，而忽略了其它数据样本的模式和特征，从而导致了多样性缺失的现象，因此会对模型的鲁棒性产生很大的负面影响。

先来看下 Centering。首先，教师模型的输出经过一个 EMA 的操作，从原始激活值中减去得到一个新的结果。简单来说，可以表述为下列公式：
$$Logits=Logits-Logit{s}_{mean}$$

这个操作的目的是使得激活值有时候是正的（当它们高于平均值时），有时候是负的（当它们低于平均值时）。由于 softmax 函数在处理负数时会给出较小的概率值，而在处理正数时会给出较大的概率值，因此这种操作能够防止任何一个特征占据统治地位，因为平均值会在值的范围中间。

最后，再看看 Sharpening。这种技巧通过在 softmax 函数中加入一个 temperature 参数，来强制让模型将概率分布更加尖锐化。由于小差异会被夸大，这会防止所有激活值都是相同的，因为小的差异也会被放大。这个技巧和中心化操作搭配使用，可以使得激活值不断变化，从而引导学生模型更好地了解哪些特征应该变得更加强大。

实验

首先，看下这张效果图：

可以看出，DINO 是能够自动学习特定于类别(class-specific)的特征，从而实现准确的无监督对象分割。

其次，我们将此模型应用于未受过训练的场景，例如用于识别重复图像：

可以看出，DINO 的表现也优于现有的最先进模型，尽管它起初并不是为这一目的设计的！

通过以上可视化结果不难看出，相比于监督学习，DINO 的潜在空间也具有很好的分离类别，这意味着它的特征足够丰富，可以分离物体中的微小差异，这使得它非常适合下游任务和迁移学习。

最后，我们通过 t-SNE 可视化一起看看 DINO 的整个学习表征过程：

总结

本文主要向大家介绍 DINO，这是第一篇探索 ViT 模型在自监督学习领域的经典代表作。DINO 整体架构基于自蒸馏的范式进行构建，包含一个教师网络和学生网络。其中，学生网络学习从局部补丁预测图像中的全局特征，该补丁受动量教师网络嵌入的交叉熵损失的监督，同时进行居中和锐化以防止模式崩溃。

DINO v2

DINOv2:Learning Robust Visual Featureswithout Supervision

导读

时隔两年半，Meta AI 团队再度为大家奉上了 DINO 的进阶版——DINOv2！这篇文章主要是借鉴最近在自然语言处理方面的一些技术和进展，如语言模型的预训练，可以使计算机可以更好地理解语言。这些技术相当重要，因为它们为计算机视觉领域的类似技术提供了新的思路和方法。

作者认为，计算机视觉任务中也可以通过类似的方法来训练模型。这样的模型可以生成一些通用的视觉特征，也就是说，这些特征适用于不同的图像分布和不同的任务，无需进行进一步的微调即可使用。而最重要的是，我们仅需使用已有的自监督学习方法便能够生成这样的特征，当然前提是要有足够多的不同来源的数据进行训练。

本文主要介绍了一些技术上的创新，以充分挖掘自监督学习的潜力。其中，包括自动化数据管道用于获取更好的数据集，以及训练包含 10 亿个可调参数的 ViT 模型，并通过无监督蒸馏方法，将其压缩成一系列能够应用于不同任务的小模型。这些技术可以让计算机视觉领域变得更加高效、精准，并且更好理解。

上图计算来自同一列的图像块之间的 PCA 可视化结果，并显示它们的前 3 个组件，每个组件都匹配到不同的颜色通道。可以看出，尽管姿势、风格甚至物体发生变化，但相关图像之间的相同部分是匹配的。（这里背景主要是通过阈值处理来屏蔽掉的）

数据处理

本文的主要贡献之一便是创建了一个大规模的数据集——LVD-142M。同以往的人工采集、标注和清洗流程不同，此数据集是通过从大量未标注的数据中检索出与几个经过精心整理过的数据集中存在相似度很高的那部分样本所组成的。下面我们一起遵循以下图例看看这个pipeline是如何构建起来的。

数据源
首先，LVD-142M 数据集的来源共包含两部分，即公开数据集和网络数据集。

1. 公开数据集:
   

2. 网络数据集:
   网络数据集简单理解就是通过爬虫将网页的图片download下来。相信这也是绝大部分算法工程师的日常工作，想必很多人都会遇到同样的问题，那便是直接扒下来的数据很“脏”！DINOv2 的做法是分三步清洗：首先，针对每个感兴趣的网页，从标签中提取出图像的 URL 链接；其次，排除掉具有安全问题或受限于域名的 URL 链接；最后，对下载的图像进行后处理，包括 PCA 哈希去重、NSFW 过滤和模糊可识别的人脸；整合起来，我们便生成了 1.2 亿张独一无二的图像。

   （1）去重
   经过第一个步骤，我们将一些质量太差或者格式破损的图像给剔除掉，然而这当中必然会存在许多冗余图像。本文采用《A self-supervised descriptor for image copy detection》这篇论文提出的copy detection pipeline进行图像查重，同样是 Meta 的工作，发表于 CVPR 2022，本质上也是通过深度学习的方法计算相似度，有兴趣的同学可以自行尝试下。
   
   经过上述操作，可以有效减少冗余并增加了图像之间的多样性。此外，作者还删除了这项工作中使用的任何基准的测试或验证集中包含的重复图像。

   （2）自监督图像检索

   这一步便是我们上面提到的，从大量未标注的数据中检索出与几个经过精心整理过的数据集中存在相似度很高的那部分样本，下面我简单总结下。

   首先，当我们要从一个大型的未筛选数据集中挑选图像来用于预训练时，我们需要先对这些图像进行聚类，以便在检索时能快速找到与查询图像相似的图像。聚类算法会对所有的未筛选图像进行分组，使得同一组内的图像在视觉上非常相似。

   那么，为了聚类过程能够顺利进行，我们需要先计算每个图像的嵌入。一个嵌入是指将一个图像转换成一个向量，使得这个向量包含了这个图像的信息。本文使用了一个在 ImageNet-22k 上进行预训练过的 ViT-H/16 的自监督神经网络来计算每个图像的嵌入。

   一旦我们计算出了每个图像的嵌入向量，我们就可以使用聚类算法来将这些向量分组。原文直接采用了 k-means 聚类算法，它会将嵌入向量相似的图像放到同一个聚类中。

   紧接着，给定一个查询图像，我们从与查询图像所在的聚类中检索N（通常为4）个最相似的图像。如果查询图像所在的聚类太小，我们会从聚类中抽样M张图像（M是由视觉检查结果后决定的）。最后，我们将这些相似的图像和查询图像一起用于预训练，那么便能够获得一个更加优质、精心筛选过的大规模预训练数据集。

需要注意的是，在这里， “视觉检查” 是指通过人工检查了一些从聚类中检索到的图像，以确保它们与查询图像在视觉上的相似度。例如，如果我们发现从聚类中检索到的图像与查询图像相似度较高，那么我们就可以增加 N 的值，或者减少 M 的值。反之，如果从聚类中检索到的图像与查询图像相似度较低，我们就可以减少 N 的值，或者增加 M 的值。这个过程是由人工进行的，因为只有人类才能够准确地判断图像之间的相似度。如此一来，我们就可以更好地确定哪些图像应该作为预训练数据集的一部分。

（3）实现细节
总结成以下四点：

• Faiss库：Faiss 是一个专门用于高维向量相似度搜索的库。在这个预训练数据集的去重和检索阶段中，我们可以采用 Faiss 库来高效地索引和搜索最近嵌入向量。
• GPU加速：本文采用了 Faiss 库的GPU-accelerated indices，利用倒排文件索引（inverted file indices）和乘积量化编码（product quantization codes）进行处理。
• 集群计算：整个预处理过程是在一个由 20 个节点组成的计算集群上完成的。每个节点都配备了 8 个 V100-32 GB 的 GPU。这样的配置能够更快地进行嵌入向量的计算和搜索，只不过一般人谁玩得起？
• 通过以上操作，最终生成的 LVD-142M 数据集是一个包含 142 百万张图像的数据集。整个预处理过程只需要不到两天的时间就能够完成。呵呵，笑笑不说话~~~

方法

本文提到了一种用于学习特征的判别式自监督方法，它是 DINO 和 iBOT 损失的结合，并引入了 SwAV 的居中方法。此外，我们还添加了一个正则项来扩展特征，并进行了一个短的高分辨率训练阶段。

简单点理解就是，这种自监督方法是由多个损失函数组成的，包括 DINO（Transformers之间的局部信息最大化），iBOT（特征之间的相似度最小化）和 SwAV（样本中心化）。同时，添加了一个正则项，以使特征在特征空间中更加均匀地分布。此外，DINOv2 中还进行了一个短暂的高分辨率训练阶段，以进一步提高特征的鲁棒性。这些方法的详细实现可以在相关的论文或我们的开源代码中找到。

1. Image-level objective
这个方法被称为图像级目标，是一种用于学习特征的判别式自监督方法。其基本思想是将来自同一图像不同裁剪的视图作为正样本，将来自不同图像的视图作为负样本，使用交叉熵损失函数来衡量这些视图之间的相似性和差异性，从而训练一个学生网络。另外，我们使用指数移动平均方法构建一个教师网络，其参数是过去迭代的加权平均值，以减少训练中的波动。最终，我们使用这两个网络的类令牌特征作为特征表示。

2. Patch-level objective
这个方法是另一种用于学习特征的自监督方法（具体细节见 iBOT），称为 Patch 级目标。在这种方法中，作者将输入的一些 Patch 随机地遮盖掉，只将未被遮盖的 Patch 提供给学生网络，然后使用交叉熵损失函数来衡量学生和教师网络在每个被遮盖的 Patch 上的特征表示的相似性和差异性，从而训练学生网络。同时，我们可以将 Patch 级别的损失与图像级别的损失相结合，以便在训练过程中兼顾整体和局部特征。

3. Untying head weights between both objectives
此方法是针对前两个方法的实验发现进行的改进。在前两个方法中，图像级别和Patch级别的损失函数都共享了一个网络的参数（权重）。但是经过实验观察发现，当两个级别的损失函数共享同样的参数时，模型在Patch级别会欠拟合，在图像级别会过拟合。因此，我们可以考虑将这些参数（权重）解绑，使得模型在两个级别都能够更好地学习特征表示。这个方法的优化目标是在两个级别都得到最佳的结果。

4. Sinkhorn-Knopp centering
这个方法是对 DINO 和 iBot 两种方法中的一些步骤进行改进。在原来的方法中，教师模型中的 softmax-centering 步骤在某些情况下可能导致不稳定性，因此本文采用了 Sinkhorn-Knopp（SK）批量归一化方法来代替。这个方法的核心思想是通过正则化来使学生和教师网络在特征表示上更加接近。在这个方法中，作者使用了 3 次 Sinkhorn-Knopp 算法迭代来实现归一化。对于学生网络，则仍然使用 softmax 归一化。通过这个方法，我们可以更好地训练学生模型，并获得更好的性能。

5. KoLeo regularizer
KoLeo regularizer 是一种正则化方法，它通过计算特征向量之间的差异来确保它们在批次内均匀分布。具体来说，它使用了一种名为 Kozachenko-Leonenko 差分熵估计的技术，这是一种估计随机样本密度的方法。在计算这个正则化器之前，特征向量会被进行“2-范数归一化”（将每个向量的所有元素平方和开根号并除以该和），以确保它们具有相同的长度。这个正则化器的作用是减少特征向量之间的差异，从而使它们在整个批次内均匀分布。

6. Adapting the resolution
这一步主要是涉及在预训练的最后一段时间内，将图像的分辨率提高到 518×518 ，便在下游任务中更好地处理像素级别的信息，例如分割或检测任务。高分辨率的图像通常需要更多的计算资源和存储空间，因此只在预训练的最后阶段使用这种方法，以减少时间和资源成本。

实现

下面介绍下几个训练模型的优化方法。本文方法采用的是最新发版的 PyTorch 2.0，同时在 A100 GPU 上进行训练，作者开放了包括预训练模型在内的代码。

关于模型细节可以在附录的表格17中查看。与 iBOT 实现相比，DINOv2 的代码在相同硬件下运行速度提高了2倍，内存使用量也只有 iBOT 的 1/3 。

1. Fast and memory-efficient attention
首先，为了提高自注意力层的速度和内存使用效率，作者基于 FlashAttention 实现了一个新的版本。由于 GPU 硬件的特殊性，每个 head 的嵌入维度是 64 的倍数时效率最佳，并且当整个嵌入维度是 256 的倍数时，矩阵运算效率更好。因此，为了最大化计算效率，本文使用嵌入维度为 1536 的 24 个头（64 dim/head），而不是使用嵌入维度为 1408 的 16 个头（88 dim/head）。实验结果表明，最终的准确率没有显著差异，整个 ViT-g 骨干网络包含 1.1B 个参数。

2. Nested tensors in self-attention
为了提高计算效率，本文实现了一个可以在同一个前向传递中运行全局裁剪和局部裁剪（它们具有不同数量的补丁标记）的版本，这种实现与之前的实现相比具有更好的计算效率。此外，还使用了 xFormers 库的一些低级组件。

3. Efficient stochastic depth
在传统的stochastic depth方法中，由于每个残差块在训练过程中可能会随机丢弃，因此需要在前向计算时对每个残差块的输出进行屏蔽（mask）以确保正确性。而在这里，实现了一种改进版本的stochastic depth方法，它跳过了被丢弃的残差的计算，而不是对结果进行屏蔽。这样可以节省内存和计算资源，并且在高丢弃率下具有更好的计算效率和内存使用率。具体地，实现方式是在批次维度上随机重排样本，并将第一个（1 - d）× B个样本用于块的计算。

4. Fully-Sharded Data Paralle
FSDP 技术的作用是使训练过程可以更加高效地扩展至多个 GPU 节点。具体地，在使用 AdamW 优化器进行训练时，需要使用4个模型副本，包括 student, teacher, optimizer 的一阶和二阶动量，对于一个类似于 ViT-g 这样动不动 10 亿参数的模型，这意味着我们至少需要 16GB 的内存。因此，为了降低每个 GPU 的内存占用，作者使用了 FSDP 技术，将模型副本分片放置在多个 GPU 上。

如此一来，模型大小不再受单个 GPU 的内存限制，而是受到整个计算节点的 GPU 内存总和的限制。FSDP 还具有第二个优点，即可以节省跨 GPU 通信成本，因为权重分片存储在 float32 精度中，而广播权重和减小梯度则使用 float16 精度进行计算。与其他自监督预训练方法中使用的 float32 gradient all-reduce operation 相比，这样做可以将通信成本降低约50%。因此，与 DDP 和 float16 自动类型转换相比，Pytorch-FSDP 混合精度在几乎所有情况下都更加优越。

Gradient AllReduce 是在分布式深度学习中，一种将梯度从不同 GPU 节点上的模型参数收集和组合成单个梯度张量的操作。通常，在使用梯度下降算法时，需要将所有 GPU 节点的梯度合并成单个梯度张量，以更新模型参数。此操作便是实现这一目的的一种并行算法，可以提高模型训练的速度和效率。

5. Model distillation
模型蒸馏是一种通过在小模型中复制大模型的输出来训练小模型的方法。这种方法旨在通过最小化大模型和小模型对一组给定输入的输出之间的距离来实现。对于该论文中的训练，研究人员利用了同样的训练循环，使用一个更大的模型作为冻结的教师模型，保留学生模型的 EMA 作为最终模型，并且移除了遮蔽和随机深度。

实验

1. 定量分析

2. 定性分析

总结

本文介绍了第一个基于图像数据的 SSL 工作——DINOv2，这是一种用于在大型图像数据集上预训练图像编码器，以获得具有语义的视觉特征。这些特征可用于广泛的视觉任务，无需微调即可获得与有监督模型相当的性能。该方法的关键是基于精心设计的数据流管道构建大规模的图像数据集，以及一些技术改进，如 Stochastic Weight Averaging、Efficient Stochastic Depth和Fully-Sharded Data Parallel等。此外，本文充分展示了这些模型的一些性质，如对场景几何和对象部分的理解。作者计划继续在模型和数据规模上进行扩展，并希望将这些视觉特征与简单的线性分类器相结合，用于创建一个语言-视觉联合系统，以处理视觉特征并提取所需的信息。

参考文献：

1. 重塑自监督学习：DINO 网络如何颠覆视觉特征表示的常规方法
2. 全网最详细的 DINOv2 论文解读来啦！