图1:在没有监督的情况下训练的具有8 × 8patch的视觉transformer的self-attention。我们来看看最后一层顶层的[CLS]token的self-attention。此token不附加任何标签或监督。这些图表明,该模型自动学习特定于类的特征,从而实现无监督的对象分割。
Abstract
在本文中,我们质疑自监督学习是否为Vision Transformer(ViT)提供了与卷积网络(convnets)相比突出的新特性[19]。
除了将自监督方法适应于这种架构特别有效的事实之外,我们还进行了以下观察:首先,自监督ViT特征包含关于图像语义分割的明确信息,这在监督ViT和convnets中都没有清晰地出现。其次,这些特征也是优秀的k-NN分类器,在较小的ViT下,在ImageNet上达到78.3%的top-1。
我们的研究还强调了momentum encoder[33]、multi-crop learning[10]以及使用small patches of ViTs的重要性。我们将我们的发现应用到一种简单的自我监督方法中,称为DINO,我们将其解释为a form of self-distillation with no lables。通过使用ViT-Base进行线性评估,在ImageNet上获得80.1%的top-1,我们展示了DINO和vit之间的协同作用。
Introduction
Transformers[70]最近作为卷积神经网络(convnets)的替代方案出现在视觉识别中[19,69,83]。它们的采用与受自然语言处理(NLP)启发的训练策略相结合,即对大量数据进行预训练并对目标数据集进行微调[18,55]。由此产生的Vision transformer(Vision transformer, ViT)[19]与convnets具有竞争力,但它们尚未提供明显的优势:它们对计算的要求更高,需要更多的训练数据,并且它们的特征没有表现出独特的属性。
在本文中,我们质疑transformers在视觉上的muted success是否可以通过在其预训练中使用监督来解释。我们的动机是,transformers在NLP中取得成功的主要因素之一是使用自监督预训练,以BERT[18]中的封闭过程或GPT[55]中的语言建模的形式。这些自我监督的预训练目标使用句子中的单词来创建pretext任务,这些任务提供了比预测每个句子的单个标签的监督目标更丰富的学习信号。同样,在图像中,图像级监督通常会将图像中包含的丰富视觉信息简化为从预定义的几千个对象类别中选择的单个概念[60]。
而NLP中使用的自监督pretext任务是针对文本的,许多现有的自监督方法已经在带有convnets的图像上显示了它们的潜力[10,12,30,33]。它们通常具有相似的结构,但具有不同的组件,旨在避免琐碎的解决方案(崩溃)或提高性能[16]。在这项工作中,受这些方法的启发,我们研究了自监督预训练对ViT特征的影响。特别有趣的是,我们已经确定了几个有趣的属性,这些属性不会出现在监督vit中,也不会出现在convnets中:
•自监督ViT功能明确包含场景布局,特别是对象边界,如图1所示。这些信息可以在最后一个块的自关注模块中直接访问。
•自监督ViT特征在基本近邻分类器(k-NN)下表现特别好,无需任何微调,线性分类器或数据增强,在ImageNet上达到78.3%的top-1准确率。
分割掩码的出现似乎是自监督方法共享的属性。然而,只有结合momentum encoder[33]和multi-crop augmentation[10]等组件时,k-NN的良好性能才会显现出来。我们研究的另一个发现是使用更小的vit patch来提高最终特征的质量的重要性。
总的来说,我们关于这些组件的重要性的发现使我们设计了一种简单的自我监督方法,可以被解释为一种没有标签的知识蒸馏[35]。由此产生的框架,DINO,通过使用标准的交叉熵损失,直接预测由动量编码器构建的教师网络的输出,从而简化了自监督训练。有趣的是,我们的方法只能通过对教师输出的centering和sharpening来避免崩溃,而其他流行的组件,如predictor[30]、advanced normalization[10]或contrastive loss[33],在稳定性或性能方面几乎没有好处。特别重要的是,我们的框架是灵活的,可以在convnets和vit上工作,而不需要修改架构,也不需要适应内部规范化[58]。
我们进一步验证了DINO和ViT之间的协同作用,通过使用带有小补丁的ViT- base,在ImageNet线性分类基准上以80.1%的top-1准确率优于先前的自监督特征。我们还确认DINO通过与ResNet-50架构匹配最先进的状态与convnets一起工作。最后,我们讨论了在计算和内存容量有限的情况下使用DINO和vit的不同场景。特别是,使用ViT训练DINO只需要两台8-GPU服务器在3天内就可以在ImageNet线性基准测试上达到76:1%,这优于基于同等规模的convnets的自监督系统,并且显著降低了计算需求[10,30]。
Approach
SSL与知识蒸馏(SSL with Knowledge Distillation)
这项工作使用的框架DINO与最近的自监督方法具有相同的总体结构[10,16,12,30,33]。然而,我们的方法与知识蒸馏也有相似之处[35],我们从这个角度来介绍它。我们在图2中演示了DINO,并在算法1中提出了一个伪代码实现。
知识蒸馏是一种学习范式,我们训练学生网络gθs来匹配给定教师网络gθt的输出,分别用θs和θt参数化。给定输入图像x,两个网络输出K维上的概率分布,表示为Ps和Pt。概率P是通过使用softmax函数对网络g的输出进行归一化得到的。更准确地说,
当τs > 0时,温度参数控制输出分布的清晰度,对于温度为τt的Pt,也有类似的公式。给定一个固定的教师网络gθt,我们学习通过最小化交叉熵损失w.r.t.来匹配这些分布。学生网络的参数θs:
H (a;B) =−a log B。
在下文中,我们将详细介绍如何将Eq.(2)中的问题应用于自监督学习。首先,我们使用multi-crop策略构建图像的不同扭曲视图或作物[10]。更准确地说,从给定的图像中,我们生成不同视图的集合V。这个集合包含两个全局视图,x g1和x g2以及几个较小分辨率的局部视图。所有crops都通过学生传递,而只有全局视图通过教师传递,因此鼓励“局部到全局”的对应。我们minimize the loss:
这种损失是普遍的,可以用于任意数量的视图,甚至只有2个。然而,我们遵循多重裁剪的标准设置,使用2个分辨率为224^2的全局视图,覆盖原始图像的大部分区域(例如大于50%),以及几个分辨率为962的局部视图,仅覆盖原始图像的一小部分区域(例如小于50%)。除非另有说明,否则我们将此设置称为DINO的基本参数化。
两个网络共享相同的体系结构,但有不同的参数集θs和θt。我们通过随机梯度下降最小化Eq.(3)来学习参数θs。
教师网络(Teacher Network)。与知识蒸馏不同,我们没有先验的教师gθt,因此,我们从学生网络的过去迭代(past iterations)中构建它。我们在第5.2节中研究了不同的教师更新规则,并表明在一个epoch内冻结教师网络在我们的框架中工作得非常好,而复制教师的学生权重则无法收敛。特别有趣的是,在学生权重上使用指数移动平均(EMA),即momentum encoder[33],特别适合我们的框架。更新规则为θt⬅λθt +(1−λ)θs;在训练期间,λ遵循从0:996到1的余弦schedule[30]。最初,动量编码器被引入作为对比学习中queue的替代品[33]。然而,在我们的框架中,它的作用有所不同,因为我们没有队列(queue)也没有对比损失(contrastive loss),并且可能更接近于自我训练中使用的平均教师的作用[65]。事实上,我们观察到教师执行了一种类似于指数衰减的Polyak-Ruppert平均的模型集成形式[51,59]。使用PolyakRuppert平均进行模型集成是提高模型性能的标准做法[38]。在整个训练过程中,我们观察到该教师的表现优于学生,从而通过提供更高质量的目标特征来指导学生的培训。在以前的研究中没有观察到这种动态[30,58]。
网络体系结构(Network architecture)。神经网络g由主干f (ViT[19]或ResNet[34])和投影头h: g = h◦f组成。下游任务中使用的特征是输出主干f。投影头由一个隐藏维数为2048的3层多层感知器(MLP)和一个K维的权值归一化全连接层[61]组成,与swv[10]的设计相似。我们已经测试了其他投影头,这种特殊的设计似乎最适合DINO(附录C)。我们没有使用预测器[30,16],导致学生和教师网络中的架构完全相同。特别有趣的是,我们注意到,与标准convnets不同,ViT体系结构默认情况下不使用Batch Normalization(BN)。因此,当将DINO应用于ViT时,我们在projection头中也不使用任何BN,使系统完全无BN。
避免崩溃(Avoiding collapse)。几种自监督方法的不同之处是,它们通过对比损失[73]、聚类约束[8,10]、预测器[30]或批归一化[30,58]来避免崩溃。虽然我们的框架可以通过多次归一化来稳定[10],但它也可以只使用动量教师输出的中心和锐化来避免模型崩溃。如5.3节实验所示,定心防止一维占主导地位,但会促使坍缩至均匀分布,而锐化则相反。运用这两种操作平衡了它们的效果,这足以避免在动量老师面前崩溃。选择这种方法来避免崩溃,以减少对批的依赖来交换稳定性:居中操作仅依赖于一阶批统计,可以解释为向教师添加一个偏置项c: gt(x) gt(x) + c。中心c用指数移动平均线更新,这使得该方法可以很好地跨不同的批大小,如第5.5节所示:
其中m > 0为速率参数,B为批大小。输出锐化是通过在教师softmax归一化中使用温度τt的低值来获得的。
实施和评估(Implementation and evaluation protocols)
在本节中,我们提供了DINO训练的实现细节,并介绍了我们实验中使用的评估方法。
视觉transformer(Vision Transformer)。我们简要描述了视觉transformer(Vision Transformer, ViT)的机制[19,70],关于变形器的详细信息请参考Vaswani等人[70],关于其对图像的适应请参考Dosovitskiy等人[19]。我们遵循DeiT[69]中使用的实现。我们在表1中总结了本文中使用的不同网络的配置。ViT架构将分辨率为N × N的非重叠连续图像块网格作为输入。在本文中,我们通常使用N = 16(“/16”)或N = 8(“/8”)。然后,这些贴片通过一个线性层,形成一组嵌入。我们在序列[18,19]中添加了一个额外的可学习token。这个令牌的作用是聚合来自整个序列的信息,我们将投影头h附加到它的输出。为了和之前工作保持一致,我们将此token称为class token[CLS],尽管在我们的案例中,它没有附加任何标签或监管。Patch token和[CLS]token集通过“预规范”层规范化(pre-norm layer normalization)被馈送到标准Transformer网络[11,39]。Transformer是一个自关注和前馈层的序列,与跳跃连接并行。自注意层通过使用注意机制查看其他token表示来更新token表示[4]。
实现细节(Implementation details)。我们在没有标签的ImageNet数据集[60]上预训练模型。我们使用adamw优化器[44]进行训练,批处理大小为1024,在使用vits /16时分布在16个gpu上。在前10次迭代中,学习率线性增加到其基本值,该基本值由以下线性缩放规则确定[29]:lr = 0:0005∗batchsize=256。在此warm up之后,我们使用余弦调度来衰减学习率[43]。权重衰减也遵循余弦时间表从0:04到0:4。温度τs设为0:1,而我们在前30个epoch的0:04至0:07对τt使用线性预热。我们遵循BYOL[30](颜色抖动、高斯模糊和日晒)和多重裁剪[10]的数据增强,使用双三次插值使位置嵌入适应尺度[19,69]。复制我们的结果的代码和模型是公开的。
评估协议(Evaluation protocols)。自监督学习的标准方案是要么在固定的特征上学习线性分类器[82,33],要么在下游任务上微调特征。对于线性评估,我们在训练期间应用随机调整作物大小和水平翻转增强,并报告中心作物的准确性。对于调优评估,我们使用预训练的权值初始化网络,并在训练期间对其进行调整。然而,这两种评估都对超参数敏感,例如,当改变学习率时,我们观察到两次运行之间的准确性差异很大。因此,我们也使用简单加权最近邻分类器(k-NN)评估特征的质量,如[73]所示。我们冻结预训练模型来计算和存储下游任务训练数据的特征。然后,最近邻分类器将图像的特征与投票给标签的k个最接近的存储特征进行匹配。我们扫描了不同数量的近邻,发现20 NN在大多数情况下都是最有效的。该评估协议不需要任何其他超参数调优,也不需要数据扩展,并且只需通过下游数据集一次即可运行,从而大大简化了特征评估。