当前位置：首页 > news >正文

(TMLR-2024)DINOv2：在无监督条件下学习鲁棒的视觉特征

news 2025/9/9 15:18:32

DINOv2：在无监督条件下学习鲁棒的视觉特征

paper title：DINOv2: Learning Robust Visual Features
without Supervision

paper是Meta发表在TMLR 2024的工作

Code:链接

Abstract

自然语言处理领域在大规模数据上进行模型预训练的最新突破，为计算机视觉中的类似基础模型铺平了道路。这些模型可以通过生成通用的视觉特征——即无需微调即可跨图像分布和任务发挥作用的特征——极大简化图像在任何系统中的使用。本研究表明，现有的预训练方法，特别是自监督方法，如果在来自多样化来源的足够精心整理的数据上进行训练，就可以产生这样的特征。我们重新审视了现有方法，并结合不同技术，在数据和模型规模上扩展了我们的预训练。大部分技术贡献旨在加速和稳定大规模训练。在数据方面，我们提出了一条自动化管道来构建专门的、多样化的、经过精心整理的图像数据集，而不是像自监督文献中通常做的那样使用未经筛选的数据。在模型方面，我们训练了一个具有10亿参数的ViT模型（Dosovitskiy et al., 2021），并将其蒸馏为一系列更小的模型，这些模型在图像和像素层面的多数基准测试上超越了最佳的通用特征——OpenCLIP（Ilharco et al., 2021）。

1 Introduction

在自然语言处理（NLP）中，学习与任务无关的预训练表征已经成为标准（Radford 等, 2019; Raffel 等, 2020; Chowdhery 等, 2022; Hoffmann 等, 2022; Touvron 等, 2023）。人们可以直接使用这些特征，即无需微调，就能在下游任务中取得远优于任务特定模型的性能（Brown 等, 2020）。这种成功主要得益于在大规模原始文本上使用自监督的预训练目标（如语言建模（Radford 等, 2017）或词向量（Devlin 等, 2019）），而无需人工监督。

遵循 NLP 的这一范式转变，我们预计在计算机视觉中也会出现类似的“基础”模型（Bommasani 等, 2021）。这些模型应当能够生成开箱即用的视觉特征，并在图像级别（如图像分类）和像素级别（如分割）任务中发挥作用。目前最有前景的工作主要集中在文本引导的预训练上，即利用某种形式的文本监督来指导特征学习（Joulin 等, 2016; Mahajan 等, 2018; Radford 等, 2021）。然而，这种方法在保留图像信息方面存在局限性，因为图像的描述性字幕只能近似图像中的丰富信息，复杂的像素级信息可能无法通过这种监督方式学习到。此外，这类图像编码器需要对齐的文本-图像语料，因此缺乏 NLP 模型的灵活性，即无法仅从原始数据中学习。

图 1：前几个 PCA 分量的可视化。我们对同一列（a、b、c 和 d）图像的 patch 进行 PCA，并展示前三个分量。每个分量对应一个不同的颜色通道。即使在姿态、风格甚至物体发生变化的情况下，相同部分在相关图像之间仍能匹配。背景通过对第一个 PCA 分量进行阈值化被移除。

另一种替代方法是自监督学习（Caron 等, 2018; Chen 等, 2020; He 等, 2022），其特征完全从图像中学习。这些方法在概念上更接近语言建模等预训练任务，并且能够在图像和像素层面捕获信息（Caron 等, 2021）。此外，已有研究表明自监督模型输出的特征具有多种有用性质，并支持了广泛的应用（Amir 等, 2022; Tumanyan 等, 2022; Ofri-Amar 等, 2023; Hamilton 等, 2022）。然而，尽管自监督学习有潜力学习通用特征，但其大多数进展仍主要依赖在较小的精心整理的数据集（如 ImageNet-1k（Russakovsky 等, 2015））上进行预训练。一些工作尝试将这种方法扩展到更大规模数据（Caron 等, 2019; Goyal 等, 2021; 2022a），但这些方法依赖于未经筛选的数据，通常会显著降低特征质量。这一问题主要是因为缺乏对数据质量和多样性的控制，而这些因素对生成良好的特征至关重要。

在本研究中，我们探讨了如果在大规模精心整理的数据上进行预训练，自监督学习是否有可能学到通用视觉特征。我们重新审视了现有的判别式自监督方法，这些方法能在图像和 patch 层面学习特征，例如 iBOT（Zhou 等, 2022a），并在更大规模数据集下重新考察其设计选择。我们的多数技术贡献旨在在模型和数据规模扩展时稳定和加速判别式自监督学习。这些改进使我们的方法速度提高约 2 倍，内存需求减少 3 倍，相较于类似的判别式自监督方法，从而能够使用更大的批量进行更长时间的训练。

在预训练数据方面，我们建立了一条自动化管道，用于从大规模未经筛选的图像集合中过滤和重新平衡数据。这一管道借鉴了 NLP 中的方法（Wenzek 等, 2020），利用数据相似性而非外部元数据，无需人工标注。处理真实世界图像的一个主要困难是如何平衡概念分布，避免模型在少数主导模式上过拟合。在本研究中，一个简单的聚类方法已经能够很好地解决这个问题。我们最终收集了一个小而多样的数据集，包含 1.42 亿张图像，用于验证我们的方法。

最后，我们提供了一系列预训练的视觉模型，称为 DINOv2，这些模型使用不同的 Vision Transformer (ViT)（Dosovitskiy 等, 2016）架构在我们的数据上训练得到。我们开源了所有模型以及在任意数据上重新训练 DINOv2 的代码。我们在多个计算机视觉基准上验证了 DINOv2 的质量，包括图像级和像素级任务，结果如图 2 总结。我们得出的结论是，仅凭自监督预训练就足以学习可迁移的冻结特征，这些特征与当前最好的公开弱监督模型相比具有竞争力。

图 2：随参数规模扩展的性能演化。我们展示了在八类视觉任务上的性能（见第 7 节），并计算了每类任务的平均指标。特征由我们的自监督编码器 DINOv2（深蓝色）提取，并与自监督方法（浅橙色）以及弱监督方法（深粉色）进行比较。我们将表现最好的弱监督模型的性能用虚线水平线表示。我们的模型系列在自监督学习上相较于之前的最新方法有了显著提升，并达到了可与弱监督特征媲美的性能。详细分析见第 7 节。

2 Related Work

图像内自监督训练。第一类自监督方法专注于从图像构建的预训练任务，即从图像中提取一个信号，并要求模型从图像的其他部分来预测这个信号。这一思想在 Doersch 等人 (2015) 的工作中变得流行，他们通过预测给定图像块的上下文来进行训练。随后提出了许多其他预训练任务，例如图像重新上色 (Zhang et al., 2016)、预测变换 (Gidaris et al., 2018)、修复缺失区域 (Pathak et al., 2016) 或图像块重排序 (Noroozi & Favaro, 2016; Misra & Maaten, 2020)。最近，随着基于图像块的架构（如 ViTs）的出现，人们重新研究了用于预训练的修复任务 (He et al., 2022; Bao et al., 2021; El-Nouby et al., 2021)，这种修复甚至可以在特征空间中进行 (Assran et al., 2023; Baevski et al., 2022)。值得注意的是，He et al. (2022) 展示了掩码自动编码器 (MAE) 学习到的特征在下游任务上微调时能显著提升性能。MAE 的这一特性也在视频 (Tong et al., 2022)、音频 (Xu et al., 2022) 以及其他模态 (Girdhar et al., 2023) 上得到了验证。然而，它们的特征需要监督微调，而我们的特征则在不经任何调整时就能有良好表现。

判别式自监督学习。第二类方法更接近我们的工作，它通过图像或图像组之间的判别信号来学习特征。这类方法的根源可以追溯到早期深度学习工作 (Hadsell et al., 2006)，但在实例分类方法 (Dosovitskiy et al., 2016; Bojanowski & Joulin, 2017; Wu et al., 2018) 出现后变得流行。随后基于实例级目标 (Hénaff et al., 2019; He et al., 2020; Chen & He, 2021; Chen et al., 2020; Grill et al., 2020; Caron et al., 2021) 或聚类 (Caron et al., 2018; Asano et al., 2020; Caron et al., 2020) 的方法不断改进。这些方法在 ImageNet (Russakovsky et al., 2015) 等标准基准上提供了优异的冻结特征，但在扩展到更大模型规模时较为困难 (Chen et al., 2021)。在本工作中，我们在大规模预训练数据和模型的背景下重新审视这些方法的训练。特别是，我们基于 Zhou et al. (2022a) 的方法进行改进，发现其特别适合扩展。

自监督预训练的扩展。越来越多的工作集中于自监督学习在数据和模型规模上的扩展能力 (Caron et al., 2019; Goyal et al., 2019; Tian et al., 2021; Goyal et al., 2022a)。大多数这些工作利用大量未经筛选的数据来进行无监督训练。他们展示了判别方法确实能随数据扩展而提升，但由于预训练数据质量较差，大部分结果需要通过微调才能获得。值得注意的是，Goyal et al. (2021) 还展示了在预训练数据充足的情况下，这些方法可以从更大模型规模的扩展中获益。这一研究方向质疑了自监督方法是否能在任意数据上都有效，而我们的重点则在于产生最好的预训练编码器。

自动化数据筛选。我们的数据集构建借鉴了图像检索社区的工作 (Weinzaepfel et al., 2021; Radenović et al., 2018b; Berman et al., 2019; Douze et al., 2009; Tolias et al., 2016; Revaud et al., 2019)。特别是，利用检索来增强训练集的方法已经在半监督学习中被研究过 (Yalniz et al., 2019)。类似地，其他研究也使用了标签或其他元数据 (Mahajan et al., 2018; Radford et al., 2021)，或预训练的视觉编码器 (Schuhmann et al., 2021; 2022) 来过滤未经筛选的数据集。与这些工作不同，我们没有使用预训练编码器、元数据或监督来过滤图像，而是利用图像之间的视觉相似性。我们的方法受到文本筛选管道 (Wenzek et al., 2020) 的启发，在那里语言模型在 Wikipedia 上进行训练以对从未经筛选的来源中提取的文本进行打分。

3 Data Processing

我们通过从一个大规模的未筛选数据池中检索与若干已筛选数据集中的图像接近的图像，组装了我们的精炼数据集 LVD-142M。下面我们描述数据管道的主要组成部分，包括已筛选/未筛选数据源、图像去重步骤以及检索系统。我们的管道不需要任何元数据或文本，直接处理图像，如图 3 所示。更多细节请参见附录 A。

图 3：我们的数据处理管道概览。来自已筛选和未筛选数据源的图像首先被映射为嵌入。未筛选图像在与已筛选图像匹配之前会先进行去重。最终的组合通过自监督检索系统扩充了初始数据集。

数据源。我们选择的已筛选数据集在附录（表 15）中详细列出，包括 ImageNet-22k、ImageNet-1k 的训练集、Google Landmarks 以及若干细粒度数据集。对于未筛选的数据源，我们从一个公开可用的网页抓取数据仓库中收集原始的未过滤图像数据。从仓库中的每个网页，我们提取 <img> 标签中的图像 URL 链接。我们丢弃不安全或受域限制的 URL，并对下载的图像进行后处理（PCA 哈希去重、NSFW 过滤以及模糊化可识别的人脸）。最终得到 12 亿张唯一图像。

去重。我们对未筛选数据应用 Pizzi et al. (2022) 的拷贝检测管道，去除近似重复图像。这减少了冗余并提高了图像的多样性。我们还移除了与本工作中使用的任何基准测试的测试集或验证集中的图像近似重复的图像。

自监督图像检索。我们通过从未筛选数据源中检索与已筛选数据源中的图像接近的图像来构建精炼的预训练数据集。具体做法是，首先使用一个在 ImageNet-22k 上预训练的自监督 ViT-H/16 网络计算图像嵌入，并采用余弦相似度作为图像间的距离度量。然后，我们对未筛选数据进行 k-means 聚类。对于检索任务，如果查询数据集足够大，我们为每个查询图像检索 N（通常为 4）个最近邻。如果查询数据集较小，我们从与每个查询图像对应的簇中采样 M 张图像。尽管人工检查表明 N 远大于 4 时也能获得不错的检索质量，但这会带来更多冲突（即同一图像被多个查询检索到）。我们选择 N = 4，因其在检索质量和冲突率之间提供了较好的折中。

实现细节。我们的管道在去重和检索阶段依赖 Faiss 库 (Johnson et al., 2019)，用于高效索引和批量计算最近邻嵌入搜索。特别是，我们大量利用其对 GPU 加速索引的支持，使用带有产品量化码 (Jegou et al., 2010) 的倒排文件索引。整个处理过程分布在一个由 20 个节点组成的计算集群上，每个节点配备 8 张 V100-32GB GPU，总耗时不到两天就生成了 LVD-142M 数据集。

4 Discriminative Self-supervised Pre-training

我们使用一种判别式自监督方法来学习特征，这种方法可以看作是 DINO 和 iBOT 损失结合 SwAV (Caron et al., 2020) 的中心化。我们还添加了一个正则项来扩展特征，并引入了一个短暂的高分辨率训练阶段。我们在这里简要介绍这些方法，但更多细节可以在相关论文或我们的开源代码中找到。

• 图像级目标 (Caron et al., 2021)。我们考虑从学生网络和教师网络提取的特征之间的交叉熵损失。这些特征来自ViT的class token，由同一图像的不同裁剪获得。我们将学生class token输入到学生DINO head中。该head是一个MLP模型，输出一组分数，我们称之为“prototype scores”。然后我们应用softmax得到 $p_s$ 。类似地，我们将教师DINO head应用到教师class token上以获得教师prototype分数。接着我们应用softmax并使用滑动平均（或之后详细介绍的Sinkhorn-Knopp居中）得到 $p_t$ 。DINO损失项对应如下：

$LDINO=−∑ptlog⁡ps\mathcal{L}_{DINO} = - \sum p_t \log p_s$

我们学习学生的参数，并通过过去迭代的指数滑动平均来构建教师head (He et al., 2020)。

• Patch级目标 (Zhou et al., 2022a)。我们随机遮挡一些输入patch给学生，但不遮挡教师。然后我们将学生iBOT head应用于学生的mask token。同样地，我们将教师iBOT head应用到对应的（可见的）教师patch token。接着我们应用softmax和居中步骤，得到iBOT损失项：

$LiBOT=−∑iptilog⁡psi\mathcal{L}_{iBOT} = - \sum_i p_{ti} \log p_{si}$

其中 $i$ 是被mask的patch索引。类似于上述过程，我们学习学生的参数，并通过指数滑动平均构建教师head。

• 解除两个目标之间的head权重绑定。DINO和iBOT损失都使用一个可学习的MLP投影head。它被应用到输出token上并在其之上计算损失。在(Zhou et al., 2022a)中，一项消融研究表明在DINO和iBOT head之间共享参数可以带来更好的性能。但在大规模实验中，我们观察到相反的结果，因此我们在所有实验中使用两个独立的head。

• Sinkhorn-Knopp居中 (Caron et al., 2020)。Ruan et al. (2023)推荐用Sinkhorn-Knopp (SK)批量归一化替换DINO和iBOT的教师softmax居中步骤，这也是SwAV (Caron et al., 2020)中的做法。我们运行Sinkhorn-Knopp算法迭代3次。对于学生，我们应用softmax归一化。

• KoLeo正则项 (Sabléayrolles et al., 2019)。KoLeo正则项来源于Kozachenko-Leonenko微分熵估计器 (Beirlant et al., 1997; Delattre & Fournier, 2017)，鼓励在一个batch内特征的均匀分布。给定一组向量 $(x1,…,xn)(x_1, \dots, x_n)$ ，其定义如下：

$Lkoleo=−1n∑i=1nlog⁡(dn,i),\mathcal{L}_{koleo} = -\frac{1}{n}\sum_{i=1}^n \log(d_{n,i}),$

其中 $dn,i=min⁡j≠i∥xi−xj∥d_{n,i} = \min_{j \ne i} \|x_i - x_j\|$ 是 $x_i$ 与batch内任意其他点之间的最小距离。我们还对特征进行 $ℓ2\ell_2$ 归一化后再计算该正则项。

• 分辨率自适应 (Touvron et al., 2019)。提高图像分辨率对像素级下游任务（如分割或检测）至关重要，因为小物体在低分辨率下会消失。然而，高分辨率训练在时间和显存上开销很大，因此我们在预训练结束的短时间内，将图像分辨率提高到 $518 \times 518$ 。这与UniViT (Likhomanenko et al., 2021)训练和FlexiViT (Beyer et al., 2023)训练的方法类似。

5 Efficient implementation

我们考虑了若干改进来在更大规模上训练模型。我们在A100 GPU上使用PyTorch 2.0训练模型。代码和预训练模型在Apache 2.0许可证下发布。模型的详细信息见附录表17。在相同硬件下，与iBOT实现相比，DINOv2代码运行速度快约2倍，且只使用1/3的显存。

快速且内存高效的注意力机制。我们实现了自己的FlashAttention版本 (Dao et al., 2022)，以改进自注意力层的内存使用和速度。我们的版本在所有考虑的情况下与原版相当或更好，同时覆盖更多用例和硬件。由于GPU硬件的特性，当每个head的embedding维度是64的倍数时效率最佳，当整个embedding维度是256的倍数时矩阵运算效果更佳。因此，我们的ViT-g架构与Zhai et al. (2022)提出的架构略有不同，以最大化计算效率。我们使用1536维embedding，配合24个head（64维/head），而不是1408维embedding和16个head（88维/head）。实验表明最终精度没有显著差异，我们的ViT-g骨干网络参数量为11亿。

序列打包。DINO算法需要同时处理大裁剪（分辨率224）和小裁剪（分辨率98）。当切分为patch后，这两组对应的token序列长度不同，无法同时前向计算。为了加速训练，我们采用了一种来自NLP的“序列打包”技巧 (Krell et al., 2022)。其核心思想是：我们将必须前向的不同序列拼接成一个长序列，通过Transformer块正常传递。但在注意力层中，对自注意力矩阵应用块对角掩码，阻止不同序列之间的注意力。这种方式在计算上严格等价于分别前向每个序列，但相比之前的分别前向和反向传播实现，大幅提升了效率。我们实现的底层组件在xFormers库 (Lefaudeux et al., 2022) 中可用。

高效的随机深度。我们实现了一个改进版本的随机深度 (Huang et al., 2016)，它跳过被丢弃残差的计算，而不是对结果进行mask。这样在内存和计算上可节省与drop rate成比例的开销，多亏了特定的融合内核。在高drop rate (d = 40%)的情况下，这极大地提高了计算效率和内存使用率。实现方式是随机打乱批次维度上的B个样本，只对前 $(1−d)×B(1-d)\times B$ 个样本进行块内计算。

完全分片数据并行 (FSDP)。使用AdamW优化器最小化目标函数需要4个模型副本（float32精度）：学生、教师、优化器一阶矩、优化器二阶矩。这对一个十亿参数的模型（如ViT-g）来说共需16 GB显存。为了减少单GPU的显存开销，我们使用PyTorch的FSDP实现，将模型副本分片到多个GPU上。因此，模型规模不再受单GPU显存限制，而是受计算节点总显存限制。PyTorch-FSDP还有第二个优势，即节省跨GPU通信开销：权重分片存储在float32中（优化器需要），但权重广播和梯度规约在backbone部分使用float16（MLP head梯度规约仍用float32以避免不稳定）。这使通信成本比DDP中的float32梯度all-reduce操作降低约50%。因此，在扩展GPU节点数量时，该训练过程比带autocast的DDP更高效。总体上，PyTorch-FSDP混合精度在几乎所有情况下都优于DDP+autocast。

模型蒸馏。我们的大多数训练环节技术改进旨在提高大模型在大规模数据上的训练效率。对于小模型，我们不是从零开始训练，而是从最大模型ViT-g进行蒸馏。知识蒸馏 (Hinton et al., 2014) 旨在通过最小化大模型与小模型输出之间的某种距离，使小模型复现大模型的输出。由于我们的目标函数本质上就是从教师网络到学生网络的一种蒸馏，因此我们使用相同的训练流程，但有一些例外：使用大模型作为冻结的教师，保留学生的EMA作为最终模型，去掉mask和随机深度，并将iBOT损失应用于两个全局裁剪。在消融实验中，我们观察到这种方法即便对ViT-L也优于从零开始训练。我们的蒸馏方法与Duval et al. (2023)描述的方法接近，但我们不修改蒸馏的损失项，并评估学生的EMA。