论文阅读笔记: DINOv2: Learning Robust Visual Features without Supervision

DINOv2: Learning Robust Visual Features without Supervision

论文地址: https://arxiv.org/abs/2304.07193
代码地址: https://github.com/facebookresearch/dinov2

摘要

• 大量数据上的预训练模型在NLP方面取得突破，为计算机视觉中的类似基础模型开辟了道路。这些模型可以通过生成通用视觉特征(即无需微调即可跨图像分布和任务工作的特征)来极大地简化任何系统中图像的使用。
• 现有的预训练方法，尤其是自监督方法，如果经过来自不同来源的足够的精选数据的训练，可以产生这样的特征。
• 本文重新审视现有方法结合不同的技术来扩展预训练数据和模型大小。大多数技术贡献旨在加速和稳定大规模训练。
• 在数据方面，本文提出了一个自动管道来构建专用的、多样化的和精选的图像数据集，而不是像自我监督文献中通常所用的那种未经精选的数据。
• 在模型方面，训练了一个具有1B参数的ViT模型，并将其提炼乘一系列较小的模型。

引言

• 学习与任务无关的预训练表示已成为NLP的标准。
• 随着NLP范式的转变，类似的"基础"模型将出现在计算机视觉。这些模型应该生成在任何任务上即用的视觉特征，无论是图像级别还是像素级别。
• 文本引导预训练的替代方案是自监督学习，其中特征仅从图像中学习。这些方法在概念上更接近语言建模等替代任务，可以捕获图像和像素级别的信息。此外，自监督模型输出的特征已被证明表现出各种有用的属性，并且用了多种应用。然而，尽管它们具有学习通用特征的潜力，但自监督学习的大部分进步都是在小型精选数据集ImageNet-1K上进行预训练的背景下取得的。但他们专注于未经整理的数据集，这通常会导致在功能质量上显著下降。这是因为缺乏对数据质量和多样性的控制，而数据质量和多样性对于产生良好的特征至关重要。
• 本文探讨了如果大量精选数据进行预训练，自监督学习是否有潜力学习通用视觉特征。重新审视了现有的区分性自监督方法，这些方法在图像和block级别上学习特征。大部分技术贡献都是为了在模型和数据规模扩展时稳定和加速有区别的自监督学习。这些改进使本文方法比四类的判别式自监督方法快两倍左右，所需内存减少三倍。
• 关于预训练数据，构建了一个自动管道，用于从大量未经整理的图像中过滤和重新平衡数据集。在处理图像时的一个主要困难是重新平衡概念并避免对一些主导模式的过度你和。在这项工作中，朴素的聚类方法可以很好地解决这个问题。
• 最后，提供了各种预训练的视觉模型DINOv2。在数据集上使用不同的ViT架构进行训练。自监督预训练本身就是学习可转移冻结特征的良好候选者。

相关工作

图像内自监督训练

• 该方法侧重于从图像构建的借口任务，即从图像中提取信号以根据图像的其余部分进行预测。它们通过预测给定补丁的上下文来进行训练，引入了许多其他替代任务，例如着色图像、预测变换、修复或补丁重排序等。
• 最近，基于补丁的架构的出现，引发了对预训练修复的重新审视，可能在特征空间中。
• He等人表明mask自动编码器(MAE)学习的功能在对下游任务进行微调时可 提供实质性改进。然而，他们的功能需要监督微调。

规模自监督学习

• 这个方法与本文工作更近，是使用图像或图像组之间的判别信号来学习特征。
• 早起来源于Hadsell等人，随着实力分类方法的出现而变得流行。基于实例级目标进行多项改进或聚类。这些方法在ImageNet等标准基准上提供了高性能的冻结特征，但很难扩展到更大的模型尺寸。

扩展自监督预训练

• 越来越多的工作集中在数据和模型大小方面的自监督学习的扩展能力。
• 这些工作大多数都使用大量未经整理的数据来训练模型，而无需监督。
• 他们证明了判别方法会随着数据的变化而扩展，但由于预训练数据的质量较差，大多数结果都是通过对特征进行微调来获得的。特别地，在给定足够的预训练数据的情况下，这些方法受益于模型大小的缩放，
• 这一系列工作质疑自监督方法处理任何数据的能力，而本文专注于生产最好的预训练编码器。

自动数据管理

• 数据集构建借鉴了图像检索社区。
• 在半监督学习的背景下研究了使用检索来增强训练集。
• 其他人也使用主题标签或其他元数据或预训练视觉编码器来过滤未经整理的数据集。
• 与之前工作不同的是，本文不适用预训练的编码器、元数据或监督来过滤图像并利用图像之间的视觉相似度。

方法

有区别的自监督预训练

本文通过一种有区别的自监督方法来学习特征。该方法可以被视为以SwAV为中心的DINO核iBOT损失的结合。还添加了一个正则化器来扩展特征和一个简短的高分辨率训练阶段。

图像级目标

• 考虑从学生和教师网络提取的特征之间的交叉熵损失。这两个特征都来自ViT的token类，是从同一个图像的不同裁剪中获得的。通过学生DINO头来传递学生token类。这个头是一个输出分数向量的MLP模型，称之为"原型分数"。然后应用softmax来得到$p_s$。
• 将教师DINO头应用到教师类token中，得到教师原生分数。然后应用softmax并使用移动平均进行居中以获得$p_t$。DINO损失项为:
  $$L_{DINO}=-\sum p_t\log p_s$$

patch级目标

• 随机屏蔽一些提供给学生的输入patch，但不屏蔽给老师的输入patch。
• 将学生的iBOT头应用到学生的mask token中。
• 同样的，将教师iBoT头应用于与学生中屏蔽的token相对应的教师patch token。
• 应用上面的softmax和中心化步骤，并获得iBOT损失项
  $$L_{iBOT}=-\sum _t{p}_{ti}\log p_{si}$$
  其中$i$是屏蔽token的patch索引。

解除两个目标之间的head权重

• DINO和iBOT 损失都使用可学习的MLP任务头。它应用于输出Token, 并在顶部计算损失。
• Zhou的消融表明DINO和iBOT之间共享参数可以带来更好的性能。在规模上，事实恰恰相反，因此在所有的实验中使用两个单独的头。

KeLeo正则化器

• KoLeo正则化器院子Kozachenko-Leonenko微分熵估计器，并鼓励bacth内特征的统一跨度。
• 给定一组n个向量$(x_1,\dots ,x_n)$, 它被定义为
  $$L_{koleo}=-\frac{1}{n}\sum _{i=1}^n\log (d_{n,i})$$
  其中$d_{n,i}=mi{n}_{j\ne i}||x_i-x_j||$是$x_i$与批次内任何其他点之间的最小距离。作者还在计算此正则化之前对特征进行了$l_2$规范化。

调整分辨率

• 提高图像分辨率是分割或检测等像素级下游任务的关键，其中小物体在低分辨率下会消失
• 高分辨训练需要时间和内存，因此在预训练结束后的短时间内将图像分辨率提高到$518\times 518$。

高效实施

考虑到对更大规模的训练模型进行一些改进，使用Pytorch2.0在A100 GPU上训练模型。使用相同的硬件，与iBOT实现相比，DINOv2代码仅使用1/3内存，运行速度大约2x左右。

快速且高效记忆的注意力

• 实现了自己的FlashAttention版本，以提高自注意力层的内存使用率和速度。
• 在考虑所有的情况下与原始版本相当或更好，同时涵盖了更多的用例和硬件。
• 由于GPU硬件的具体情况，当每个头的embedding维度为64的倍数时，效率最佳，当完整的embedding维度为256的倍数时，矩阵运算甚至更好。
• ViT-g架构与Zhai等人提出的架构略有不同，为了最大限度地提高计算效率，使用1536的embedding维度和24个头，而不是1408的embedding和16个头。

序列打包

• DINO算法需要前向计算大crop(分辨率224)和小crop(分辨率98)。当分割成patch时，这两组由不同长度的词符序列表示，并且不能一起前向计算。
• 为了加速训练，使用了一种名为"sequence packing"的技巧。这个方法为: 通过Transformer序列连接成一个长序列进行前向。
• 然而，对注意力层中的自注意力矩阵应用了block对角mask，从而防止了不同序列之间的注意力。这样，相当于分辨对每个序列进行前向。
• 与之前的实现中使用单独的前向和后向传播相比，这个技巧带来了显著的效率提升。

有效的随机深度

• 实现了随机深度的改进版本，它跳过了被丢弃的残差计算，而不是屏蔽结果。
• 得益于特定的融合内核，这可以节省内存和计算，其比例大约等于丢弃率。
• 由于丢弃率过高，因此可以大幅提高计算效率和内存使用率。
• 该实现包括在batch维度上随机混合B样本，并切片第一个$(1-d)\times B$样本以进行block中的计算。

全分片数据并行(FSDP)

• 利用AdamW优化器最小化我们的目标需要4个float32精度的模型副本——学生、教师、优化器第一时刻、优化器第二时刻。
• 对于想ViT-G这样的十亿参数模型来说，总计为16GB内存。
• 为了减少每个GPU的内存占用，在GPU之间分割模型副本，即使用FSDP的Pytorch实现在GPU之间进行分片。
• 因此，模型大小不受单个GPU内存的限制，而是受计算节点伤的GPU内存综合限制，
• FSDP还可以节省跨GPU通信成本。权重分片根据优化器要求以float32精度存储，但骨干网络的广播权重和减少梯度以float16精度完成。

模型蒸馏

• 对训练训练的大多数技术改进旨在改进大型模型在大量数据上的训练。
• 对于较小的模型，是从最大的模型ViT-g中提取它们，而不是从头开始训练它们。
• 知识蒸馏旨在通过最小化一组给定输入的两个输出之间的距离，用较小的模型再现大模型的输出。