DINO v2

目录

1、引言

2、相关工作

[2.1 图像内自监督训练](#2.1 图像内自监督训练)

[2.2 判别式自监督学习](#2.2 判别式自监督学习)

[2.3 自监督预训练的扩展](#2.3 自监督预训练的扩展)

[2.4 自动数据管理](#2.4 自动数据管理)

3、数据处理

[3.1 数据源](#3.1 数据源)

[3.2 去重](#3.2 去重)

[3.3 自监督图像检索](#3.3 自监督图像检索)

[3.4 实施细节](#3.4 实施细节)

4、判别式自监督预训练

[4.1 图像层面的目标](#4.1 图像层面的目标)

[4.2 图像块层面的目标](#4.2 图像块层面的目标)

[4.3 Untying head weights between both objectives](#4.3 Untying head weights between both objectives)

[4.4 Sinkhorn-Knopp centering](#4.4 Sinkhorn-Knopp centering)

[4.5 KoLeo regularizer](#4.5 KoLeo regularizer)

[4.6 Adapting the resolution](#4.6 Adapting the resolution)

[5、Efffcient implementation](#5、Efffcient implementation)

[5.1 Fast and memory-efficient attention](#5.1 Fast and memory-efficient attention)

[5.2 Sequence packing](#5.2 Sequence packing)

[5.3 Efffcient stochastic depth(高效的随机深度)](#5.3 Efffcient stochastic depth(高效的随机深度))

[5.4 Fully-Sharded Data Parallel (FSDP)](#5.4 Fully-Sharded Data Parallel (FSDP))

[5.5 Model distillation](#5.5 Model distillation)

6、结果


DINOv2是Meta公司于2023年4月推出的基于Vision Transformer架构的自监督视觉模型,采用10亿参数的ViT-g模型在1.42亿张无标注图像数据集(LVD-142M)上进行预训练。该模型通过自监督学习方法构建通用视觉特征,无需任务微调即可应用于图像分类、深度估计、语义分割等下游任务,在COCO等基准测试中性能与OpenCLIP等基线模型相当 。

核心技术融合了DINO损失与iBOT损失的多目标优化策略,结合自动化数据管道优化使训练速度提升2倍。2023年的实验数据显示,其图像相似度计算得分达93%(COCO验证集),在零样本语义分割任务中通过特征融合方案达到SOTA性能。


因为在大量数据上对模型进行预训练,NLP方面取得了突破,这为计算机视觉领域的类似基础模型的开发开辟了道路。这些模型能够极大地简化图像在任何系统中的使用方式,它们能够生成通用的视觉特征,即这些特征能够在不同的图像分布和任务中通用,无需进行任何调整。这项研究表明,现有的预训练方法(尤其是自监督方法)如果基于来自不同来源的充足精心整理的数据进行训练,是能够生成此类特征的。我们重新审视了现有的方法,并将不同的技术结合起来,在预训练时扩大数据量和模型规模。大多数技术聚焦于加快和稳定大规模训练。在数据方面,我们提出了一种自动化的流程,用来构建一个专门、多样且经过精心筛选的图像数据集,而不是像经典自监督那样使用未经筛选的数据。在模型方面,我们训练了一个拥有 10 亿参数的 ViT 模型,并将其蒸馏为一系列规模更小的模型,这些模型在大多数图像和像素级别的基准测试中都超越了现有的最佳通用特征模型---OpenCLIP。

1、引言

与学习任务无关的预训练已成为自然语言处理(NLP)领域的标准做法。人们可以"直接"利用这些特征,无需进行微调,就能在后续任务中取得比特定任务模型更出色的性能。这一成功得益于使用大量原始文本进行预训练,所采用的方法是利用诸如语言建模或词向量等无需监督的预设目标。

随着NLP领域的这一范式转变,我们认为在计算机视觉领域也会出现类似的"基础"模型。这些模型生成在任何任务中都能直接使用的视觉特征,既包括图像层面(例如图像分类),也包括像素层面(例如分割)。针对这些基础模型的最具前景的尝试主要集中在文本引导式预训练上,即利用某种形式的文本监督来指导特征的训练过程(参考SUPIR、InstructBLIP、BLIP3)。这种文本引导式预训练形式会限制对图像所保留信息的范围,因为标题只能大致反映图像中的丰富信息,而像素级的复杂信息可能无法通过这种监督机制得以展现。此外,这些图像编码器需要对齐的文本-图像语料数据集,因此它们不具备与文本编码器相同的灵活性,即无法仅通过原始数据进行学习。

一种替代文本引导式预训练的方法是自监督学习,在这种方法中,特征是仅从图像中学习而得的。这些方法在概念上更接近于诸如语言建模之类的预设任务,并且能够从图像和像素层面获取信息。此外,由自监督模型生成的特征已被证明具有多种有用特性,并且可用于各种各样的应用场景。然而,尽管它们具备学习通用目标特征的能力,但大多数自监督学习方面的进展都是在基于一个小型的精心整理的数据集(即 ImageNet-1k)进行预训练的背景下取得的。已经有人尝试将这些方法推广到 ImageNet-1k 之外的领域,但这些尝试主要针对的是未经筛选的数据集,这通常会导致特征的质量出现显著下降。

在这个研究中,我们在大量精心整理的数据上进行预训练,探索自监督学习是否有可能学习到通用的视觉特征。我们重新审视了现有的基于判别式方法的自监督学习方式,这些方法能够在图像和图像块级别同时学习特征,例如 iBOT,并且我们从更大的数据集的角度重新审视了它们的一些设计选择。我们大部分的技术贡献是扩大模型和数据规模时,稳定并加快判别式自监督学习的过程。这些改进使得我们的方法速度提高了约 2 倍,并且所需内存仅为类似判别式自监督方法的 3 分之一,这使我们能够采用更大的批次大小进行更长时间的训练。

关于预训练数据,我们构建了一个自动化的流程,用于从大量未经筛选的图像集合中过滤并重新平衡数据集。这个流程的设计灵感来源于NLP领域中使用的流程,在该流程中,数据的相似性被用作替代外部元数据,并且无需人工标注。在处理自然图像时,一个主要的难题在于要重新平衡各种概念,并避免过度依赖少数几种主要模式。在这项工作中,一种简单的聚类方法在解决此问题时表现得相当出色。我们收集了一个规模较小但内容多样(共计 1.42 亿张)的图像集合,以验证我们的方法。

最后,我们提供了多种预训练的视觉模型,称为 DINOv2,这些模型是基于不同的视觉转换器(ViT)架构在我们的数据上训练而成的。我们公开了所有模型及代码,以便能够基于任何数据对 DINOv2 进行重新训练。我们在图像和像素层面上对 DINOv2 的质量进行了验证,验证范围涵盖了各种计算机视觉基准测试,相关结果如图 2 所示。
图2

2、相关工作

2.1 图像内自监督训练

第一类自监督方法主要关注基于图像构建的预设任务,即从图像中提取一个信号用于预测图像中的其他部分。这一观点随着Doersch 等人(2015 年)的研究而变得更为普遍,在他们的研究中,他们通过预测特定区域的上下文来进行训练。此外,还引入了许多其他辅助任务,例如图像重新着色、预测变换、图像修复或图像块重新排序等。近来,基于patch的架构(如 ViT)的出现,促使人们重新审视用于图像修复的预训练的方法,这有可能在特征空间中发挥作用。特别值得一提的是,何等人(2022 年)的研究表明,掩码自编码器(MAE)所学习到的特征在针对下游任务进行微调后能带来显著的效果提升。MAEs 的这一特性已在视频(唐等人,2022 年)、音频(许等人,2022 年)以及其他多种形式的模态中得到了进一步验证(吉尔达尔等人,2023 年)。然而,他们的特征需要进行有监督的微调,而我们的特征则无需这样做就能达到良好的效果。

2.2 判别式自监督学习

第二项工作内容,与我们的工作更为接近,利用图像之间或图像组之间的差异信号来学习特征。这一系列方法的起源可追溯至早期的深度学习研究,但随着实例分类方法的出现而变得流行起来。我们对在实例层面的目标或者聚类进行了多项改进。这些方法在诸如 ImageNet 等标准测试集上能提供高效的冻结特征效果,但它们难以扩展到更大的模型规模上。在我们的工作中,我们重新审视了这些方法在大规模预训练数据集和模型背景下的训练情况。特别地,我们在此基础上延续了周等人(2022a)的研究成果,发现这种方法特别适用于进行规模扩展。

2.3 自监督预训练的扩展

越来越多的研究都聚焦于自监督学习在数据量和模型规模方面的扩展能力。这些研究中的大多数都使用了大量未经筛选的数据来训练模型,且在训练过程中没有进行任何监督。他们提供了相关证据表明,基于差异性的方法与数据量呈正比关系,但由于预训练数据的质量不佳,大多数结果都是通过微调特征获得的。特别值得一提的是,戈亚尔等人(2021 年)还表明,只要拥有足够的预训练数据,这些方法就能通过模型规模的扩大而获得优势。这一研究方向质疑了自监督方法在处理任何数据时的适用性,而我们的目标是开发出最优质的预训练编码器。

2.4 自动数据管理

我们的数据集构建借鉴了图像检索领域的做法。特别是,使用检索技术来扩充训练集这一方法已在半监督学习中有了一定的研究。同样地,还可以利用标签或其他元数据,或预训练视觉编码器来过滤未经审核的数据集。与这些工作不同,我们并未使用预训练的编码器、元数据或监督信息来过滤图像,而且也不依赖图像之间的视觉相似性来进行操作。我们的方法受到文本筛选流程的启发,在该流程中,语言模型在维基百科上进行训练,以对从未经筛选的来源提取的文本进行评分。

3、数据处理

我们通过从大量未经筛选的数据中检索出与多个已筛选数据集中的图像较为相似的图像,来构建我们精心整理的 LVD-142M 数据集。下面我们将详细介绍我们数据处理流程中的主要组成部分,包括经过整理/未经整理的数据来源、图像去重的步骤以及检索系统。我们的流程不需要任何元数据或文本,而是直接处理图像,如图 3 所示。
图3

3.1 数据源

我们所选取的精心整理的数据集详情见附录(表 15),其中包括 ImageNet-22k、ImageNet-1k 的训练集拆分部分、Google 地标数据集以及若干细粒度数据集。对于未经筛选的数据源,我们从一个公开可用的爬取网页数据的资源库中收集了一组未经处理的原始图像数据集。对于资源库中的每个网页,我们从 <img> 标签中提取出图像的 URL 链接。我们会丢弃那些不安全或受域名限制的URLs,并对下载的图片进行后处理(如使用主成分分析哈希去重、过滤不适宜内容以及对可识别的人脸进行模糊处理)。这便产生了 12 亿张独一无二的图片。

3.2 去重

我们采用了Pizzi 等人(2022 年)提出的复制检测流程,对未经筛选的数据进行处理,并删除了近似重复的图像。这减少了图像之间的重复部分,同时又增加了图像的多样性。此外,我们还会删除在本研究中所使用的任何基准测试的测试集或验证集中所包含的近似重复的图像。

3.3 自监督图像检索

我们通过从未经过筛选的数据源中检索出与已筛选数据源中的图像较为相似的图像,来构建我们的定制预训练数据集。为了实现这一目标,我们首先使用 ImageNet-22k 数据集预训练一个自监督 ViT-H/16 网络来计算图像的嵌入,并使用余弦相似度来度量图像之间的距离。然后,我们对未经筛选的数据进行 k 均值聚类分析。对于用于检索的查询数据集,如果其规模足够大,我们将为每个查询图像检索出 N 个(通常为 4 个)最近邻图像。如果规模较小,我们就从与每个查询图像对应的集群中采样出M 张图像。尽管视觉检查似乎表明当 N 大于 4 时检索质量良好,但这却导致了更多的冲突(有些图像会被检索给多个查询)。我们选择 N = 4 这个值,因为从这个角度来看,它在平衡性方面表现得较为理想。

3.4 实施细节

我们流程中的去重和检索阶段依赖于 Faiss 库(约翰斯等,2019 年)来高效地对嵌入进行索引以及进行批量近似嵌入搜索的计算。特别是,我们充分利用了其对 GPU 加速索引的支持,采用了带有产品量化编码的倒序文件索引。整个处理过程是在一个由 20 个节点组成的计算集群中完成的,该集群配备了 8 块 32GB 的 V100 图形处理器,生成 LVD-142M 数据集仅需不到两天的时间。

4、判别式自监督预训练

我们采用一种判别式自监督方法来学习特征,该方法可以看做是 DINO 和 iBOT 损失的组合,并结合了 SwAV 的中心化处理。我们还加入了正则化项以使特征分布更加均匀,以及使用高分辨率进行了一个较短的训练阶段。

4.1 图像层面的目标

我们对从学生网络和教师网络中提取的特征计算交叉熵损失。这两个特征都来自一个ViT的类别标记,该标记是通过对同一张图片的不同部分进行裁剪而获得的。我们将学生类别令牌通过学生的DINO head传递出去。这个head是一个 MLP 模型,输出一个分数向量,我们将其叫做"原型分数"。然后我们应用 softmax 函数来得到 ps 。同样地,我们将教师 DINO head 应用到教师类别标记上,以获得教师原型分数。然后,我们应用一个 softmax 操作,接着进行中心化处理(采用移动平均法,或者按照以下内容所述采用 Sinkhorn-Knopp 中心化方法),以得到 pt 。对学生的参数进行更新,并利用过去的几个迭代值使用指数移动平均的方法来构建教师模型。DINO 损失项对应于:

4.2 图像块层面的目标

我们对学生的输入块进行了随机遮挡处理,但教师的数据未进行处理。然后,我们将学生的遮挡token输入到学生iBOT head。同样地,我们将未遮挡的教师块token 送入教师iBOT head 。然后,我们应用 softmax 和归一化操作,从而得到 iBOT 损失项:

其中 i 表示被遮蔽标记的块索引。与上述情况类似,我们学习学生的参数,并通过指数移动平均法构建教师head 。

4.3 Untying head weights between both objectives

DINO 和 iBOT 两者均使用可学习的多层感知机投影头。它应用于输出标记,并且损失是在此之上计算得出的。在周等人(2022a)的研究中的一项消融实验表明,在 DINO 和 iBOT 头部之间共享参数能够带来更好的性能。从大规模的数据来看,我们发现情况恰恰相反,因此在我们所有的实验中都使用了两个独立的heads。

4.4 Sinkhorn-Knopp centering

Ruan等人(2023 年)建议将 DINO 和 iBot 中的教师 softmax 均衡步骤替换为 SwAV 的Sinkhorn-Knopp(SK)批量归一化方法。我们对教师中的Sinkhorn-Knopp算法进行了 3 次迭代运算。对于学生,我们采用 softmax 正则化方法。

4.5 KoLeo regularizer

KoLeo正则化器源自Kozachenko-Leonenko的微分熵估计器,并且会促使同一批次中的特征具有均匀的分布范围。给定一组 n 个向量 (x1, ..., xn),其定义如下:

其中 表示 xi 与批次内任何其他点之间的最小距离。在计算此正则化项之前,我们还会对特征进行 ℓ2 标准化处理。

4.6 Adapting the resolution

提高图像分辨率对于像素级别的后续任务(如分割或检测)至关重要,因为在低分辨率下,小物体会变得难以辨认。然而,高分辨率的训练需要耗费大量时间和内存资源。因此,我们在预训练的最后阶段将图像分辨率提高到 518×518,且在高分辨率输入时进行一段较短的时间的训练。这也与 UniViT和FlexiViT 的训练类似。

5、Efffcient implementation

我们考虑了一些改进措施,以便在更大规模上训练模型。我们使用 PyTorch 2.0 在 A100 GPU 上训练模型。使用相同的硬件,与 iBOT 实现方案相比,DINOv2 代码的运行速度提高了约 2 倍,同时仅占用 1/3 的内存。

5.1 Fast and memory-efficient attention

我们开发了自己版本的FlashAttention(源自 Dao 等人,2022 年)来提高自注意力层的内存使用效率和运算速度。在我们所考虑的所有情况下,我们的版本都与原版相当甚至更胜一筹,同时还能涵盖更多的应用场景和硬件类型。由于 GPU 的硬件特性,当每个头的嵌入维度为 64 的倍数时效率最高;而当整个嵌入维度为 256 的倍数时,矩阵运算的效果会更好。因此,我们的 ViT-g 架构与Zhai等人提出的架构略有不同,目的是为了最大限度地提高计算效率。我们使用的是 1536 的嵌入维度和 24 个头(每个头 64 维),而不是 1408 的嵌入维度和 16 个头(每个头 88 维)。我们的实验显示在准确率方面没有明显的差异,我们的 ViT-g 主干网络包含 11 亿个参数。

5.2 Sequence packing

DINO 算法需要同时前向传播大尺寸裁剪(分辨率为 224)和小尺寸裁剪(分辨率为 98)。当这些数据被分割成多个块时,这两组数据将以不同长度的标记序列形式呈现,因此无法一起进行前向传播。为了加快训练速度,我们采用了名为"序列打包"的技巧,该技巧源自自然语言处理领域(Krell 等人,2022 年)。这个思路很简单:我们将必须通过transformer前向传播的序列连接成一个长序列。我们像往常一样将此序列通过transformer块。然而,在注意力层中,对自注意力矩阵应用了一个块对角掩码,防止不同序列之间的注意力。这样,这个前向传播就严格等效于对每个序列分别进行前向传播。与之前分别进行前向传播和后向传播的实现方式相比,这种技巧为我们带来了显著的计算效率提升。我们的设置的底层组件可在 xFormers 库中找到(勒福杜瓦等,2022 年)。

5.3 Efffcient stochastic depth(高效的随机深度)

我们采用了随机深度算法的改进版,该算法不再对被舍弃的残差进行计算,而是直接忽略其结果。得益于特定的融合核技术,这在一定程度上节省了内存和计算资源,其节省比例与丢弃率大致相当。由于具有较高的丢弃率(在本研究中d= 40%),这使得计算效率和内存使用量得到了显著的提升。该实现过程包括在批次维度上随机打乱 B 个样本,并且从这组样本中选取前 (1 - d) × B 个样本用于该块中的计算。

5.4 Fully-Sharded Data Parallel (FSDP)

使用 AdamW 优化器来最小化我们的目标值,需要 4 个 32 位浮点数精度的模型副本:学生模型、教师模型、优化器的第一阶矩、优化器的第二阶矩。这意味着对于像我们ViT-g这种拥有十亿参数的模型而言,其内存需求总计为 16GB。为了减少每个 GPU 的内存占用量,我们将模型副本分布在多个 GPU 上,即使用 PyTorch 实现的 FSDP 技术将 16GB 的内存分布在各个 GPU 上。因此,模型的大小不受单个 GPU 内存的限制,而是受计算节点上 GPU 内存总和的限制。PyTorch 实现的 FSDP 还带来了第二个优势,那就是能够降低跨 GPU 通信的成本:权重碎片按照优化器的要求以 float 32 位精度进行存储,但权重的广播操作和梯度的缩减在骨干网络中则是以 float16 精度进行,对于骨干网络(为了避免训练过程中的不稳定情况,MLP 头的梯度则是在float 32 位精度下进行缩减)。这使得通信成本相较于其他自监督预训练方法中使用的 DistributedDataParallel(DDP)的 float32 梯度全归约操作减少了约 50%。因此,在增加 GPU 节点数量时,训练过程的扩展效率比使用 float16 自动类型转换的 DDP 更高。总体而言,在我们遇到的几乎所有情况下,PyTorch-FSDP 混合精度都优于带有自动类型转换的 DDP。

5.5 Model distillation

我们在训练流程中的大多数技术改进旨在提高使用大数据进行大型模型训练的效果。对于较小的模型,我们是从我们最大的模型------ViT-g 中对其进行蒸馏,而不是从头开始训练它们。知识蒸馏(Hinton 等人,2014 年)旨在通过最小化给定输入下两个模型输出之间的某些距离,用较小的模型来重现大型模型的输出。由于我们的目标函数是将教师网络的特征提炼至学生网络,因此我们采用了与教师网络相同的训练流程,但有一些例外:我们使用一个较大的模型作为冻结的教师模型,保留一个备用的 EMA(指数移动平均)版本的学生模型作为最终模型,取消掩码和随机深度,并在两个全局裁剪上应用 iBOT 损失。在我们的消融实验中,我们发现这种方法的性能要优于从头开始训练的方法,即便是在使用 ViT-L 的情况下也是如此。我们的蒸馏方法与Duval等人(2023 年)所描述的方法非常接近,只是我们没有对蒸馏的损失项进行修改,并且评估的是学生的指数移动平均值。

6、结果

在本节中,我们将对我们的模型在众多图像理解任务中的表现进行实证评估。我们对全局和局部图像表示进行了评估,涵盖了类别和实例级别的识别、语义分割、单目深度预测以及动作识别等方面。评估的目标有两个方面。第一,我们的自监督特征的表现远远优于当前的SOTA。其次,它们在大量任务中与弱监督特征的表现相当,甚至更胜一筹。

相关推荐
九.九6 小时前
ops-transformer:AI 处理器上的高性能 Transformer 算子库
人工智能·深度学习·transformer
春日见6 小时前
拉取与合并:如何让个人分支既包含你昨天的修改,也包含 develop 最新更新
大数据·人工智能·深度学习·elasticsearch·搜索引擎
恋猫de小郭6 小时前
AI 在提高你工作效率的同时,也一直在增加你的疲惫和焦虑
前端·人工智能·ai编程
deephub7 小时前
Agent Lightning:微软开源的框架无关 Agent 训练方案,LangChain/AutoGen 都能用
人工智能·microsoft·langchain·大语言模型·agent·强化学习
偷吃的耗子7 小时前
【CNN算法理解】:三、AlexNet 训练模块(附代码)
深度学习·算法·cnn
大模型RAG和Agent技术实践7 小时前
从零构建本地AI合同审查系统:架构设计与流式交互实战(完整源代码)
人工智能·交互·智能合同审核
老邋遢7 小时前
第三章-AI知识扫盲看这一篇就够了
人工智能
互联网江湖7 小时前
Seedance2.0炸场:长短视频们“修坝”十年,不如AI放水一天?
人工智能
PythonPioneer7 小时前
在AI技术迅猛发展的今天,传统职业该如何“踏浪前行”?
人工智能
冬奇Lab8 小时前
一天一个开源项目(第20篇):NanoBot - 轻量级AI Agent框架,极简高效的智能体构建工具
人工智能·开源·agent