NeurIPS 2021: All Tokens Matter: Token Labeling for Training Better Vision Transformers论文精读与解析

摘要

本文提出了一种新的训练目标——token labeling,用于训练高性能视觉Transformer(ViT)。

与ViT的标准训练目标(计算附加的class token的分类损失)不同,我们提出的目标利用所有的patch token以密集方式计算训练损失。
具体来说,token labeling将图像分类问题重新定义为多个token级别的识别问题,并为每个patch token分配一个由机器注释器生成的特定于位置的监督。

实验表明,token labeling可以在宽光谱范围内清晰一致地提高各种ViT模型的性能。以具有26M可学习参数的ViT为例,通过token labeling,该模型在ImageNet上可以达到84.4%的Top-1精度。通过将模型尺寸稍微扩大到150M,结果可以进一步提高到86.4%,使之前的模型(250M+)中最小尺寸的模型达到86%。

我们还表明,token labeling可以明显提高预处理模型对具有密集预测的下游任务(如语义分割)的泛化能力。

介绍

在过去的几年里,Transformer[39]在几乎所有的自然语言处理(NLP)任务中都取得了优异的性能[4、14、24]。由于这种成功,最近,许多研究人员尝试为视觉任务构建变压器模型,他们令人鼓舞的结果表明,基于变压器的模型在图像分类方面具有巨大潜力[6、15、25、36、40、46],特别是自我注意机制在构建成对输入标记之间的长期依赖关系方面的强大优势。

尽管收集长期依赖关系很重要,但最近关于本地数据增强的工作[57]表明,良好的建模和利用本地信息进行图像分类将避免模型偏向于倾斜和非通用模式,并大大提高模型性能。然而,最近的视觉变换器通常使用聚合全局信息的类标记来预测输出类,而忽略了其他在其各自的局部图像patch上编码丰富信息的token labeling的作用。

在本文中,我们提出了一种新的视觉变换器训练目标,称为token labeling,它同时利用了patch tokens和class tokens。我们的方法采用机器注释器生成的K维得分图作为监督,以密集方式监督所有标记,其中K是目标数据集的类别数。通过这种方式,每个token labeling都明确地与一个单独的特定位置监督相关联,指示相应图像patch内目标对象的存在,从而提高视觉变压器的对象识别能力,而计算开销可以忽略不计。据我们所知,这是首次证明密集监督对图像分类中的视觉变换器有益的工作。

根据我们的实验,利用提出的token labeling可以明显提高视觉变换器的性能。如图1所示,我们的模型名为LV ViT,具有56M参数,在ImageNet[13]上获得85.4%的前1精度,比所有其他不超过100M参数的基于变压器的模型表现更好。当模型规模扩大到150M时,结果可以进一步提高到86.4%。此外,我们经验发现,带有token labeling的预处理模型也有利于密集预测的下游任务,例如语义分割。
在这里插入图片描述

相关工作

Transformer[39]指的是完全依赖于自我注意机制来构建全局依赖的模型,这些模型最初是为自然语言处理任务设计的。由于其捕捉空间信息的强大能力,Transformer也已成功应用于各种视觉问题,包括图像增强等低水平视觉任务[7,45],以及图像分类[9,15]、目标检测[5,11,55,61]、分割[7,33,41]和图像生成[28]等更具挑战性的任务。一些工作还扩展了视频和3D点云处理的变压器[50、53、60]。

Vision Transformer(ViT)是较早在ImageNet分类上实现最先进性能的尝试之一,它使用纯变压器作为基本构建块。然而,ViT需要对非常大的数据集(如ImageNet-22k和JFT-300M)进行预处理,并需要巨大的计算资源,以实现与ResNet[18]相当的性能,并在ImageNet上训练类似的模型大小。后来,DeiT[36]通过简单地调整网络架构并添加额外的token和知识提取class token[21,47]来解决数据低效问题,以提高模型性能。

最近的一些工作[6,16,43,46]还试图通过修改patch embedding block或transformer block或两者,将局部依赖性引入视觉变压器,从而显著提高性能。此外,还有一些作品[20、25、40]采用了金字塔结构,以减少总体计算,同时保持模型捕获低级特征的能力。
与前面大多数设计新变压器块或变压器架构的工作不同,我们试图通过研究在patch中嵌入了本地信息的tokenlabeling的作用来改进视觉变压器。我们表明,通过稍微调整视觉变换器的结构并采用所提出的token labeling目标,我们可以在不同的模型大小级别上为变换器模型实现强大的基线。

方法

在本节中,我们首先简要回顾了ViT的结构[15],然后描述了token labeling。

ViT

典型的视觉变换器[15]首先将固定大小的输入图像分解为一系列小块。每个patch通过线性层投影映射到一个特征向量,或称为token。然后,将所有token与用于分类分数预测的额外可学习cls token相结合,发送到变压器块堆栈中进行特征编码。

在损失计算中,通常从最后一个变压器块的输出token中选择cls token,并将其发送到线性层,以进行分类分数预测。数学上,给定一个图像I,将最后一个变压器块的输出表示为[Xcls,X1,…,XN],其中N是补丁标记的总数,Xcls和X1,…,XN分别对应于类标记和补丁标记。图像I的分类损失可以写为:

在这里插入图片描述
H:softmax交叉熵损失

Token Labeling

上述分类问题只采用图像级别标签作为监督,而忽略了每个图像补丁中嵌入的丰富信息。在本小节中,我们提出了一种新的训练目标标记,它利用了patch token和class token之间的互补信息。

Token Labeling:原始的ViT中,测量单个class token(表示整个输入图像)和相应图像级别标签之间的距离,token labeling强调所有输出标记的重要性,并主张每个输出token都应该与一个特定于位置的标签相关联。因此,在我们的方法中,输入图像的ground truth不仅涉及单个K维的one hot向量ycls,还涉及一个K×N矩阵,称为K维score map,由[y1,…,yN]表示,其中N是输出的patch token的数量。

给定输出patch tokens[X1,…,XN]和相应的标签[y1,…,yN],token labeling目标可以定义为:

在这里插入图片描述

因此,总损失函数可以写成:
在这里插入图片描述
其中β是平衡这两个项的超参数。在我们的实验中,我们根据经验将其设置为0.5。

在这里插入图片描述

除了使用类标记(粉色矩形)之外,我们还利用了所有输出patch标记(橙色圆角矩形),方法是为每个补丁标记分配一个由机器注释器[3]生成的特定于位置的预测作为监督(请参阅红色虚线矩形中的部分)。我们提出的token
labeling方法可以作为辅助目标,为每个贴片标记提供局部细节,以帮助视觉变换器更准确地定位和识别目标对象。请注意,传统的视觉变压器训练不包括红色虚线矩形部分。

我们的token labeling具有以下优点。

● 首先,与知识蒸馏不同,token labeling是一种廉价的操作。稠密分数图可以通过预先处理的模型生成(例如,EfficientNet[34]或NFNet[3])。
● 在训练期间,我们只需要裁剪分数图并执行插值,使其与空间坐标中的裁剪图像对齐。因此,附加计算可以忽略不计。
● 其次,与大多数分类模型和ReLabel策略[49]中使用单个标签向量作为监督不同,我们还利用分数图来密集监督模型,因此每个补丁标记的标签提供特定于位置的信息,这有助于训练模型轻松发现目标对象并提高识别精度。
● 由于在训练中采用了密集监督,我们发现带有token labeling的预处理模型有利于密集预测的下游任务,如语义分割。

Token Labeling with MixToken

在训练ViT时,先前的研究[36、46]表明,像MixUp[52]和CutMix[48]这样的增强方法可以有效地提高模型的性能和鲁棒性。
●Mixup:将随机的两张样本按比例混合,分类的结果按比例分配;
●Cutout:随机的将样本中的部分区域cut掉,并且填充0像素值,分类的结果不变;
●CutMix:就是将一部分区域cut掉,但不填充0像素,而是随机填充训练集中的其他数据的区域像素值,分类结果按一定的比例分配。

在这里插入图片描述
然而,ViT依赖于基于patch的tokenization来将每个输入图像映射到一系列token,我们的token labeling策略也对基于patch的token进行操作。如果我们直接在原始图像上应用CutMix,一些生成的patch可能包含来自两个图像的内容,导致小patch中的混合区域,如图3所示。

在这里插入图片描述

在执行token labeling时,很难为每个输出标记分配一个干净正确的标签。考虑到这种情况,我们重新考虑了CutMix增强方法并提出了MixToken,它可以被视为patch embedding后对token操作的CutMix的修改版本,如图3右侧所示。

具体地,对于图像I1与I2,它们的token labels分别为Y1=[y11,y12,…,y1N],Y2=[y21,y22, …, y2N],我们首先将这两张图片进行patch embedding,生成两个token序列,记为T1=[t11,t12,…,t1N],T2=[t21,t22,…,t2N]。接下来,通过使用二进制掩码M应用MixToken生成一个新的token序列,如下所示:

在这里插入图片描述
对于相应的token labels,我们还使用相同的掩码M混合它们:

在这里插入图片描述
从而,class token的标签为:

在这里插入图片描述

实验

请参考原论文。

猜你喜欢

转载自blog.csdn.net/qq_41112170/article/details/127500366
今日推荐