本文证明,在大型、多样化和精选的数据集上预训练计算机视觉模型可以产生通用的视觉特征,这些特征在各种任务中表现良好,无需微调。该研究结合了现有方法,并侧重于缩放数据和模型大小。提出了一种自动管道来创建专用的、多样化的和精选的图像数据集。研究人员训练了一个 10 亿参数的 ViT 模型并将其提炼成更小的模型,这些模型在图像和像素级别的大多数基准测试中优于以前最好的通用功能 OpenCLIP。
数据处理
研究人员通过从未经整理的来源检索图像来组装一个名为 LVD-142M 的精选数据集,这些来源与几个精选数据集(ImageNet-22k、ImageNet-1k 的训练拆分、Google Landmarks 和几个细粒度数据集)中的图像相似。他们使用公开可用的网络数据创建一个包含 12 亿张独特图像的未经过滤的原始数据集,并应用复制检测管道删除近似重复项。为了构建精选的预训练数据集,他们使用自监督的 ViT-H/16 网络计算图像嵌入,并对未精选的数据执行 k 均值聚类。根据查询数据集检索来自相应集群的最近邻居或图像。重复数据删除和检索阶段使用 Faiss 库进行高效索引和批量搜索。整个过程分布在一个由 20 个节点组成的计算集群上,每个节点有 8 个 V100,生成 LVD-142M 数据集的时间不到两天。
判别式自监督预训练
研究人员采用一种有区别的自我监督方法来学习特征,结合了 DINO、iBOT 和 SwAV 方法的元素。他们的方法包括图像级目标、补丁级目标、未绑定的头部权重、Sinkhorn-Knopp 居中、KoLeo 正则化器和自适应分辨率。
- 图像级和补丁级目标涉及学生和教师网络之间的交叉熵损失。
- 解开头部重量可提高两种秤的性能。
- Sinkhorn-Knopp 居中取代了 teacher softmax 居中步骤。
- KoLeo 正则化器鼓励在一个批次中使用统一的特征跨度。
- 最后,他们调整分辨率,在预训练结束时将分辨率短时间增加到 518x518,以更好地处理像素级下游任务。
高效实施
研究人员实施了多项改进以更大规模地训练模型:
- 他们开发了一个更快、内存效率更高的 FlashAttention 版本,优化了 ViT-g 架构以提高计算效率。由于 GPU 硬件特性,当每个头的嵌入维度为 64 的倍数时效率最佳,而当完整嵌入维度为 256 的倍数时矩阵运算甚至更好。
- 自注意力中的嵌套张量用于全局和局部作物,从而显着提高计算效率。
- 实现随机深度的改进版本(在批次维度上打乱样本并将一些第一个样本切片用于块中的计算)以节省内存和计算(比例大约等于丢弃率(本文中为 40%),这要归功于特定的融合内核)。
- 采用完全分片数据并行 (FSDP) 技术来减少每个 GPU 的内存占用并节省跨 GPU 通信成本。与 DistributedDataParallel 相比,这可以更有效地扩展 GPU 节点的数量。
- 对于较小的模型,他们使用知识蒸馏,从最大的模型 ViT-g 中提取它们,而不是从头开始训练。这种方法获得了更好的性能,并利用相同的训练循环进行了一些修改(使用更大的模型作为冻结教师,保留学生的备用 EMA 作为最终模型,删除掩蔽和随机深度,并将 iBOT 损失应用于两种全球作物)。
消融研究
在本节中,研究人员通过训练多个模型并观察它们的表现来评估改进后的训练方法中每个组成部分的重要性。通常,每个组件都会提高 k-NN 或线性探测的性能,并且在大多数情况下,两者都会提高。
预训练数据源。特征的质量直接关系到预训练数据的质量。他们比较了在 LVD-142M、ImageNet-22k 和未经整理的数据集上训练的模型的性能。结果表明,在大多数基准测试中,对一组经过整理的图像进行训练比对未经整理的数据进行训练效果更好,这证实了整理数据对自我监督预训练的好处。与在 ImageNet-22k 上训练的模型相比,在 LVD-142M 上的训练在除 ImageNet-1k 之外的所有基准测试中均优于在 ImageNet-22k 上训练的模型,这表明更多样化的图像集可以提高特征质量。总的来说,LVD-142M 数据集提供了不同类型图像的良好平衡,从而获得最佳性能。
型号尺寸和数据。随着模型大小的增长,在 LVD-142M 上的训练变得比在 ImageNet-22k 上的训练更有益。在 LVD-142M 上训练的 ViT-g 在 ImageNet-1k 上与在 ImageNet-22k 上训练的模型的性能相匹配,同时在其他基准测试中优于它。
损失组件。KoLeo 损失将实例检索性能提高了 8% 以上,而不会对其他指标产生负面影响。来自 iBOT 的蒙版图像建模术语对于密集预测任务至关重要,可将性能提高近 3%。
知识蒸馏的影响。对于小型架构,提取较大的模型比从头开始训练它们更有效。在 12 个基准测试中的 10 个上,经过蒸馏的 ViT-L/14 优于从头开始训练的 ViT-L/14。
分辨率的影响。高分辨率训练可在不同分辨率下获得更好的性能,但计算量更大。仅在一小部分训练期间(例如,10k 次迭代)以高分辨率进行训练会产生几乎同样好的结果,同时需要的计算量要少得多。因此,作者在训练结束时包括了这个步骤,而不是从头开始以高分辨率进行训练。
结果
作者通过在冻结的主干上训练一个简单的线性分类器来测试模型生成的图像表示的质量。他们的方法在线性评估方面比之前的技术水平(在 ImageNet-22k 上训练的 iBOT ViT-L/16)有显着改进(+4.2%),并且在替代测试集上具有更强的泛化能力。
作者还将它们的功能与最先进的弱监督模型进行了比较,发现它们的主干超过了具有 ViT-G/14 架构的 OpenCLIP 和具有 ViT-g/14 的 EVA-CLIP 的性能。此外,它们在 ImageNet-V2 测试集上的性能明显更好,表明泛化能力更好。
执行健全性检查以确定他们的模型是否可以通过对特定数据集的监督进行微调。当对主干进行微调时,ImageNet-1k 验证集上的 Top-1 精度提高了 +2% 以上。他们最好的微调性能仅略低于绝对最先进的水平。
最后,对域泛化基准进行稳健性分析。结果表明,与 iBOT 相比,他们的模型具有明显更好的鲁棒性,并且改进了 ImageNet-A 上最好的弱监督模型,而落后于 ImageNet-R 和 Sketch。
公平与偏见分析
作者在 Dollar Street 数据集上评估了地理公平性,该数据集比较了不同国家和收入水平的表现。他们的模型在地区和收入方面比 SEERv2 模型更公平,并且明显优于监督基线。然而,他们的模型仍然明显偏向西方国家,尤其是在非洲,与欧洲相比,非洲的表现下降了 25.7%。该模型对高收入家庭的表现也优于低收入家庭,相差 31.7%。尽管有所改进,该模型仍然表现出对西方国家富裕家庭的偏见。
作者使用 Goyal 等人的协议评估了他们对不同性别、肤色和年龄的人的图像进行分类的模型。(2022b)。他们在 ImageNet-22k 的 619 个类的子集上训练多类分类器,并将它们分组到更广泛的类别中。他们使用 Casual Conversations 数据集将他们的模型与 SEERv2 进行比较。他们的模型通常将所有群体的图像分类为人类,肤色之间没有大的偏差,并且不会预测来自非人类或犯罪元类别的有害标签。然而,该模型经常触发 Possibly-Human 类,其中包含通常与人类相关的对象。在这项研究中,没有明确的模式表明对特定群体的偏见,但作者承认,更彻底的评估可能会揭示他们模型中的缺陷。