（2024|TMLR|Meta，DINOv2，ViT，自蒸馏，iBOT，SwAV 中心化，判别式自监督预训练，分类/分割，分辨率调整）无监督稳健的视觉特征学习

DINOv2: Learning Robust Visual Features without Supervision论文地址https://arxiv.org/abs/2304.07193项目页面https://github.com/facebookresearch/dinov2进 Q 学术交流群922230617 或加 CV_EDPJ 进 W 交流群目录1. 引言2. iBOT3. 数据处理4. 判别式自监督预训练5. 高效实现6. 消融实验7. 结果7.1 ImageNet 分类7.2 其他视频和图像分类任务7.3 实例识别7.4 密集预测7.5 定性分析1. 引言在自然语言处理领域模型在大规模数据上进行预训练为计算机视觉领域类似的 “基础模型” 开辟了道路。这类模型通过生成通用视觉特征——即无需微调即可适用于不同图像分布和任务的特征——可以极大地简化图像在任何系统中的使用。本工作表明现有的预训练方法尤其是自监督方法如果在足够多样化的精选数据上进行训练就能产生这样的通用视觉特征。本文重新审视了现有的判别式自监督方法如 iBOT并在更大数据集的背景下重新考虑其设计选择。本文的大部分技术贡献旨在在扩展模型和数据规模时稳定和加速判别式自监督学习。这些改进使得比类似的判别式自监督方法快约 2 倍内存消耗减少 3 倍从而能够利用更大的批量进行更长时间的训练。在预训练数据方面本文构建了一个自动化的流程从一个大规模未精选图像集合中筛选和重新平衡数据集。该流程的灵感来自 NLP 中的方法利用数据相似性而非外部元数据且不需要人工标注。处理野外wild图像的一个主要困难是重新平衡概念避免对少数主导模式的过拟合。在这项工作中一个简单的聚类方法就能很好地解决这个问题。本文收集了一个包含 142M 图像的多样化小型语料库来验证本文的方法。本文提供了多种预训练的视觉模型称为DINOv2它们使用不同的 Vision Transformer 架构在本文的数据上训练。本文在各种图像级和像素级的计算机视觉基准测试上验证了 DINOv2 的质量并展示了其随规模扩展的性能。本文的结论是自监督预训练本身是学习可迁移的、可直接使用的冻结特征的优秀候选方法其性能可与最佳的开源弱监督模型例如OpenCLIP相媲美。图 2参数扩展时的性能演变。本文展示了在八种视觉任务上的性能并给出了每种类型的相关平均指标。特征是从本文的自监督编码器 DINOv2深蓝色中提取的并将其与自监督方法浅橙色以及弱监督方法深粉色进行了比较。本文报告了弱监督模型中表现最佳的性能作为虚线水平线。本文的模型系列在自监督学习方面大幅超越了之前的最先进水平并达到了与弱监督方法相当的性能。2. iBOTiBOT全称 image BERT pre-training with Online Tokenizer是一种自监督学习框架。其核心思路是借鉴自然语言处理中 BERT 的成功经验为视觉模型尤其是 Vision Transformer, ViT训练一个更好的 “在线标记器”。核心机制在线标记器 (Online Tokenizer)。在自然语言处理中文本可以被清晰地分词但图像的语义是连续且模糊的。iBOT 的核心创新在于引入了在线标记器。它不像以往方法如 BEiT那样需要提前训练好一个固定的分词器而是通过自蒸馏的方式让教师网络动态地作为学生网络的标记器两者可以同步优化。训练方式掩码图像建模 (MIM) 自蒸馏。iBOT 结合了两种策略使得模型既能像 BERT 一样理解局部结构又能把握图像的全局语义。掩码图像建模 (MIM)模型被要求去预测图像中被遮挡部分的特征迫使模型关注图像的内部结构。自蒸馏通过学生网络和教师网络互相学习提升模型对全局语义的理解。iBOT 和 DINO 的核心关系是iBOT DINO 掩码图像建模 (MIM)。它们同属自蒸馏框架共享相同的骨干架构和避免模型崩塌的技术。iBOT 在 DINO 仅关注图像全局语义的基础上增加了一个局部细节的学习任务。衍生与演进iBOT 是 DINO 的直接扩展可以直接理解为 iBOT 是在 DINO 代码库上为投影头增加了处理 “图像块” 的能力并增加了一个损失函数项。DINOv2 的集大成者Meta 后续推出的DINOv2进一步融合了这两者的优点它同时使用了 DINO 的图像级损失和 iBOT 的图像块级损失成为了目前性能最强的视觉基础模型之一。3. 数据处理给定多个精选数据集中的图像本文通过从一个大规模未精选数据池中检索与其相似的图像来组装本文的精选数据集 LVD-142M。数据源精选数据集包括 ImageNet-22k、ImageNet-1k 的训练集、Google Landmarks 以及多个细粒度数据集。未精选数据源包含约 1.2B 张图像。图像去重本文首先对未精选数据源进行自去重移除几乎相同的图像。然后本文进一步进行相对去重移除与评估数据集的训练集和测试集过于相似的图像以保证评估的公正性。检索系统本文使用基于样本和基于聚类两种方法来增强数据集。对于大型数据集本文为每个样本检索其最近邻图像对于小型数据集本文先对未精选数据源进行聚类然后从与目标数据集相关的聚类中挑选图像。4. 判别式自监督预训练2021|ICCVDINOViT自监督学习知识蒸馏自监督视觉 Transformer 的新特性本文采用一种结合了 DINO 和 iBOT 损失函数并融合 SwAV 中心化方法的判别式自监督方法来学习特征。此外本文添加了正则化项以分散特征并在训练末期加入了短的高分辨率训练阶段。图像级目标基于学生网络与教师网络从同一图像的不同裁剪区域提取的特征ViT 的类别 token计算交叉熵损失。学生 token 通过学生 DINO 头MLP得到原型分数经 softmax 得到 p_s​教师 token 经教师 DINO 头后通过 softmax 和中心化移动平均或 Sinkhorn-Knopp得到 p_t​。损失函数为学生参数通过梯度更新学习教师参数则通过学生参数的指数移动平均EMA进行更新。补丁patch级目标对学生网络的输入图像块进行随机掩码而教师网络则输入完整图像。将学生网络的掩码 token 通过学生 iBOT 头教师网络中对应位置的可见 token 通过教师 iBOT 头经 softmax 和中心化后得到 p_ti​ 与 p_si​。损失函数为其中 i 为被掩码图像块的索引教师网络同样通过 EMA 更新。解耦头权重在原始 iBOT 中DINO 与 iBOT 的投影头共享参数被认为是有益的。但在大规模训练中本文发现相反解耦两者效果更好。因此本文在所有实验中都使用两个独立的投影头。Sinkhorn-Knopp 中心化本文采纳了 Ruan 等人的建议将 DINO 和 iBOT 中教师网络的 softmax 中心化步骤替换为 SwAV 中使用的 Sinkhorn-KnoppSK批归一化方法。本文运行 SK 算法 3 次迭代而学生网络仍使用标准的 softmax 归一化。KoLeo 正则化KoLeo 正则化源自 Kozachenko-Leonenko 微分熵估计器旨在使一个批次内的特征在流形上均匀分布。对于一组 n 个向量 (x_1​, …, x_n​)其定义为其中d 是 x_i 到批次内其他点的最小距离。计算此正则化前特征会先进行 ℓ2 归一化。分辨率调整提高图像分辨率对分割、检测等像素级下游任务至关重要。然而全程使用高分辨率训练成本过高。因此本文仅在预训练末期将图像分辨率短暂提升至 518×518以实现效率与性能的平衡。5. 高效实现本文采用了几项改进来高效地训练大规模模型。相较于 iBOT 的实现DINOv2 在相同硬件上运行速度提升约 2 倍内存占用仅为其 1/3。快速且内存高效的注意力机制本文实现了定制版的 FlashAttention在自注意力层上提升了速度并降低了内存占用。该版本在各项测试中性能不逊于原始实现且覆盖了更多使用场景和硬件。受 GPU 硬件特性影响当每个头的嵌入维度为 64 的倍数、总嵌入维度为 256 的倍数时效率最高。因此为最大化计算效率本文的 ViT-g 架构与原始设计略有不同嵌入维度设为 1536包含 24 个头每个头 64 维而非 1408 维、16 个头每个头 88 维。实验表明这一调整对最终精度无显著影响调整后的 ViT-g 骨干网络参数量为 1.1B。序列打包Sequence packingDINO 算法需同时处理大尺寸分辨率 224和小尺寸分辨率 98的裁剪图像。两者分割为图像块后产生的 token 序列长度不同无法在同一批次中并行处理。为此本文借鉴 NLP 领域的 “序列打包” 技巧将需前向传播的多个序列拼接成一个长序列传入 Transformer 模块同时在自注意力矩阵上应用分块对角掩码阻止不同序列间的相互注意力。这使得前向传播在数学上等价于分别处理各序列但显著提升了计算效率。高效随机深度本文对随机深度进行了改进不再对丢弃的残差结果进行掩码而是直接跳过其计算。通过专用的融合内核节省的内存和计算量约等于丢弃率。本文使用 d 40% 的丢弃率该改进带来了计算效率和内存使用的显著提升。具体实现为在批次维度上随机打乱 B 个样本仅取前 (1−d)×B 个样本参与当前块的计算。全分片数据并行FSDP使用 AdamW 优化器进行训练需要维护 4 份 float32 精度的模型副本学生、教师、一阶矩、二阶矩。对于 ViT-g 这样的 1.1B 参数模型这总计需要 16 GB 内存。为降低单卡内存占用本文利用 PyTorch 的 FSDP 实现将这 16 GB 的模型副本分片到多张 GPU 上使模型规模不再受限于单卡显存而是由计算节点总显存决定。此外FSDP 还降低了跨 GPU 通信开销权重分片以 float32 存储以满足优化器要求但主干网络的广播权重与梯度规约操作在 float16 精度下进行MLP 头的梯度仍在 float32 下规约以防止训练不稳定。相比其他自监督方法使用的 DDP 和 float32 梯度全局规约FSDP 的通信成本降低了约 50%。因此在扩展 GPU 节点数量时使用 FSDP 混合精度的训练效率远超 DDP 搭配 float16 自动混合精度。模型蒸馏本文对训练流程的改进主要面向大数据集上的大模型训练。对于较小的模型本文选择从最大的 ViT-g 模型进行蒸馏而非从头训练。知识蒸馏的目标是让小型模型模仿大型模型的输出。由于本文的训练目标本身已包含从教师网络到学生网络的蒸馏形式因此本文沿用相同的训练流程仅作少数调整使用一个冻结的 ViT-g 作为教师保留学生的指数移动平均EMA作为最终模型移除掩码机制和随机深度并在两个全局裁剪上应用 iBOT 损失。消融实验表明即使对于 ViT-L 模型此蒸馏方法的效果也优于从头训练。6. 消融实验本文通过一系列实验验证了各组件的有效性。改进的训练方法逐步添加改进组件如 KoLeo、Sinkhorn-Knopp 等表明这些组件通常会持续提升 k-NN 和线性探针的性能尽管有些组件如 LayerScale虽然会暂时降低线性探针性能但能显著提升训练稳定性。预训练数据源对比在 ImageNet-22k、未精选数据和 LVD-142M 上训练的模型结果表明在精选数据上训练比在未精选数据上训练效果更好。使用更多样化的 LVD-142M 数据集在 ImageNet-1k 之外的基准测试上表现更优。模型规模与数据规模随着模型规模增大在更大、更多样化的 LVD-142M 上训练带来的收益也越大。损失函数组件KoLeo 正则化显著提升了实例检索任务的性能超过 8%且对其他任务无害。iBOT 掩码图像建模损失对密集预测任务如分割至关重要。知识蒸馏从 ViT-g 蒸馏出的 ViT-L 在所有测试基准上均优于从头训练的 ViT-L验证了蒸馏方法的有效性。分辨率在 224 和 416 分辨率下训练的模型在不同测试分辨率下的性能如上图所示。在训练末期进行短暂的高分辨率微调能以较低的计算成本获得接近全程高分辨率训练的性能。7. 结果本文在图像分类、实例识别、密集预测分割、深度估计和视频分类等多个任务上评估了 DINOv2 特征。7.1 ImageNet 分类DINOv2 的冻结特征在线性探针评估上显著优于现有自监督模型提升 4.2%并超越了 OpenCLIP 和 EVA-CLIP 等弱监督模型。微调对模型进行微调仍能带来性能提升表明其不仅适合做冻结特征也适合微调。泛化DINOv2 在 ImageNet-A、-R 和 Sketch 等分布外数据集上表现出色。7.2 其他视频和图像分类任务在 iNaturalist 等细粒度分类和 Places205 等场景分类任务上DINOv2 大幅超越现有自监督模型并优于 OpenCLIP。在视频分类任务UCF101 Kinetics-400 SSv2上DINOv2 同样表现优异。7.3 实例识别在牛津、巴黎等图像检索基准上DINOv2 的特征在平均精度mAP上显著超越了自监督和弱监督模型OpenCLIP。7.4 密集预测语义分割在 ADE20K、CityScapes 和 Pascal VOC 上使用简单的线性分类器DINOv2 的冻结特征就取得了极具竞争力的结果甚至接近一些需要微调的方法。深度估计在 NYUd、KITTI 和零样本迁移NYUd → SUN RGB-D任务上DINOv2 的特征在 RMSE 指标上均优于所有对比模型。7.5 定性分析语义分割与深度估计DINOv2 产生的分割图和深度图更平滑、更完整伪影更少。分布外泛化模型在绘画、动物等未见过的图像上仍能产生合理的深度和分割结果。PCA 分析对图像块特征进行 PCA第一主成分能有效分离前景和背景后续主成分能对应物体的不同部分并在同类物体间保持语义对应关系。跨图像块匹配模型能准确匹配不同图像中功能相似的部分如鸟的翅膀和飞机的机翼展现出对物体部件的理解能力。