SCRATCHING VISUAL TRANSFORMER’S BACK WITH UNIFORM ATTENTION

摘要

视觉变压器(ViTs)的良好性能通常归因于多头自注意(MSA)。MSA使ViT模型的每一层都能进行全局交互，这是与卷积神经网络(cnn)形成鲜明对比的特征，卷积神经网络逐渐增加了多层交互的范围。我们研究了注意力密度的作用。我们的初步分析表明，注意图的空间相互作用更接近于密集相互作用，而不是稀疏相互作用。这是一个奇怪的现象，因为密集的注意力图对模型来说很难学习，因为它们周围有更陡峭的softmax梯度。我们将此解释为对包含密集交互的ViT模型的强烈偏好。

因此，我们手动将统一的注意力插入到ViT模型的每一层，以提供急需的密集交互。我们称这种方法为上下文广播，CB。我们观察到，CB的加入降低了原始注意图的密度程度，增加了ViT模型的容量和可泛化性。CB产生的成本可以忽略不计:模型代码中的一行，没有额外的参数，以及最少的额外操作。

1. 介绍

Transformers(Vaswani等人，2017)在语言领域取得成功后，Dosovitskiy等人(2021)将其扩展到视觉变形金刚(ViTs)，其操作与变形金刚几乎相同，但用于计算机视觉任务。最近的研究(Dosovitskiy等，2021;Touvron等人(2021b)的研究表明，ViTs在图像分类任务上取得了优异的性能。

这种良好的表现通常归因于vit中的多头自注意(MSA) (Dosovitskiy等2021;Touvron等，2021b;Wang et al .， 2018;Carion等，2020;Strudel等人，2021;Raghu等人，2021)，这有利于长期依赖具体来说，MSA是为各层空间信息的全局交互而设计的。这是与大量成功的前辈卷积神经网络(cnn)在结构上形成对比的特征，卷积神经网络通过堆叠许多固定和硬编码的局部操作(即卷积层)逐渐增加交互范围。Raghu等人(2021)和Naseer等人(2021)已经证明，与cnn相比，vit中的自注意对空间信息的全局交互具有有效性。

与以往的研究不同，我们研究了空间注意力密度的作用。远程依赖，或全局交互，表示从引用令牌到达远程位置的连接。密度是指所有令牌之间非零交互的比例。我们在图1中说明了两者之间的区别。注意，“全局”并不一定意味着“密集”，反之亦然，因为注意力地图可以是密集局部的，也可以是稀疏全局的。

我们首先检查ViT模型是学习密集的还是稀疏的注意图。我们基于熵测量的初步分析表明，学习到的注意力地图在所有空间位置都趋于密集。这是一个奇怪的现象，因为密集的注意力图更难学习。自关注图是softmax操作的结果。正如我们在分析中显示的那样，梯度在密集的注意力地图周围变得不那么稳定。换句话说，vit正在努力学习密集注意图，尽管通过梯度下降学习它们很困难。虽然通过梯度下降很难学习密集注意，但手动实现它很容易。我们通过所有token的平均池化来插入均匀的注意力，这是最密集的注意力形式。我们将这个模块称为上下文广播(CB)。该模块将平均池令牌添加到每个中间层的每个令牌。

我们发现，当将CB添加到ViT模型中时，它降低了所有层注意图的密度程度。CB还使ViT模型的整体优化更容易，并提高了其泛化能力。CB模块在ImageNet上的图像分类任务中带来一致的收益(Russakovsky et al .， 2015;Recht等人，2019;Beyer等人，2020)和ADE20K上的语义分割任务(Zhou等人，2017;2019)。总体而言，CB模块似乎有助于ViT模型将其资源从学习密集注意图转移到学习其他信息信号上。

这些好处的成本可以忽略不计。只需要在model.py中插入一行代码。没有引入额外的参数;只有极少的操作是。

我们的贡献是：

• 我们观察到ViTs受益于进一步密集的相互作用(第2.1节);
• 一个简单有效的模块，CB和CBS，用于注入密集的相互作用(第2.2节);
• 验证CB和CBS的好处的实验结果(第3节)。

2. 方法

在第2.1节中，我们首先提出了对ViT架构进一步密集交互的需求。然后，我们在第2.2节中提出了一个简单的轻量级模块和一种显式地向vit注入密集交互的技术。

2.1 动机

自关注操作使vit在不限制每层空间范围的情况下进行空间交互。尽管与cnn相比，此类操作固有的丰富度，但我们研究了注意力的密度。第一个实验考察了ViTs是否需要进一步的空间相互作用。在第二个实验中，我们测量了注意力的分层熵，以检验vit倾向于学习空间相互作用的哪些特征。

**vit是否渴望更多的空间连接?ViTs中的多头自注意(MSA)和多层感知器(MLP)分别负责空间和通道相互作用。我们研究了添加哪个块(MSA或MLP)更能提高vit的性能。我们在ImageNet-1K上训练了300个epoch的8层ViT，在最后一层插入一个额外的MSA或MLP层。额外的参数数和flop数几乎相等。在图2中我们画出了训练损失和top-1/5 准确率。我们观察到，与额外的MLP相比，额外的MSA可以降低训练损失和提高验证精度。这表明，给定额外参数和flop的固定预算，ViT架构似乎更倾向于具有额外的空间交互，而不是通道交互。
MSA学习哪种类型的空间相互作用?在这个实验中，我们研究了MSA特别喜欢的空间相互作用类型。了解互动的类型将指导我们如何提高注意力表现。以往的研究关注的是MSA的远程依赖的有效性，而我们关注的是MSA的密度。我们通过熵的透镜根据深度来衡量注意力的分散。当注意力稀疏时，注意力的熵较低。另一方面，当注意力密集时，它的熵就高。图3a显示了viti - ti /-S/-B中每个MSA层的头部和令牌平均熵值的趋势(Dosovitskiy等人，2021;Touvron et al .， 2021b)。我们观察到，对于给定N = 197的注意图，其熵值趋向于更高，高达4.4，趋向于最大熵值−p1 N log 1 N≈5:3。为了提供对熵值的直观解释，我们从预训练的ViT-S中可视化训练样本的注意图。为了便于比较，我们还将低熵值的注意力可视化。图3b描述了根据熵值的注意图密度的程度，其中在熵值与观察到的平均熵4.4一样高的注意图中观察到密集分布的权重。值得注意的是，vit中的大多数注意力具有如此高的熵值;这表明MSA倾向于学习密集的相互作用。

**
最陡的梯度周围均匀的注意。密集相互作用的极端形式是均匀分布。为了研究注意机制的均匀分布解的难易性，我们深入研究了softmax函数的特点。简而言之，我们表明梯度幅度在诱导均匀输出的输入周围是最大的。我们进一步形式化这个直觉。
这表明，均匀的注意力可以被单个梯度步骤打破，这意味着它可能是最不稳定的注意力类型，至少从优化的角度来看是这样。

结论。我们已经研究了MSA层中相互作用的密度。我们发现，进一步的空间连接比进一步的通道交互更有利于ViT模型。MSA层倾向于学习具有更高熵的密集相互作用。考虑到学习密集交互的难度，ViT对密集交互的偏好是惊人的:MSA层的梯度随着更密集的注意图而更陡峭。密集的注意力图很难学习，但似乎对vit至关重要。

显式广播上下文

我们已经观察到额外空间交互的好处，以及从优化的角度来看，统一注意力是最具挑战性的注意力。观察结果促使我们设计一个互补的操作，明确地提供统一的注意力。我们通过CB模块的广播上下文来实现这一点。

背景广播(CB)。给定一个N个令牌$X\in R<sup>{N\times d}$的序列，我们的CB模块将平均池令牌提供给令牌，如下所示:

其中$x_i\in R</sup>d$是x中的第i个令牌。图4a说明了我们的CB模块。CB模块位于多层感知器(MLP)块的末端。图4b显示了我们的CB模块的整体ViT架构。我们在3.1节中的分析表明，无论其位置如何，CB的插入都会提高ViTs的性能。正如我们将看到的，当它被插入到MLP块之后时，性能的提高是最显著的。

集中注意力会有所帮助。为了检验均匀注意的有效性，我们进行了将均匀注意直接注入多头自注意层的初步实验。我们设计了三种直接注入均匀关注层的方法:1)将多头自关注头中的一个替换为CB(记为Replc)，减少了替换头对应的参数数量;2)调整Replc的参数数量。与原始的ViT(记为Comp.)相比较的方法，以及3)我们将CB作为并行的额外头部附加到多头自我注意(记为Att.)。表1给出了按参数个数排序的top-1精度。这三种方法的准确率分别提高了0.2%p、0.4%p和0.7%p。与上述把CB添加到MLP的方法相比，这些设置更直接融合密集交互到MSA中。因此，结果明确地告诉我们向vit注入密集相互作用的广泛好处。
计算效率。CB模块在PyTorch (Paszke等人，2019)，Tensorflow (Abadi等人，2016)和JAX (Bradbury等人，2018)等深度学习框架中使用一行代码实现:X = 0.5 * X + 0.5 * X.mean(dim=1, keepdim=True)。它不会增加参数的数量。它只会为推理和训练带来微不足道的额外操作。
维度缩放虽然我们关注的是上述无参数的方法，但我们提出了一种计算效率高的技术，即通过引入最小数量的参数来缩放维度。所提出的CB将密集交互注入到所有信道维度。令牌的某些通道维度需要密集的交互，而其他通道维度则不需要。我们引入维度标度权值Λ$\in R^d$，为每个维度选择性地注入均匀的注意力，如下所示:

其中为元素积。CBS引入的参数很少:为ViT-S增加了0.02%的参数量。
与SE的比较。SE模块(Hu et al .， 2018)与CB有一定的相似之处:两者都是神经网络架构的模块化附件，并使用池化操作。然而，SE (Hu et al .， 2018)旨在通过利用池来构建通道描述符、两个FC层和一个sigmoid函数来建模通道相互依赖性。不仅动机和操作的精确集合不同，而且计算需求也不同。参见表2中对它们的额外计算成本的比较。CB是一个比SE便宜得多的模块。最后，我们比较了SE和CB模型在ImageNet-1K上的性能，见表3。这两个模块都提高了ViT模型的性能，但CB的改进更大。考虑到计算成本和精度的提高，与SE相比，CB证明是有利的。

3. 实验

在第3.1节中，我们试验了放置模块的位置。在第3.2和3.3节中，我们评估了图像分类模块(ImageNet-1K (Russakovsky等人，2015))和语义分割模块(ADE20K (Zhou等人，2017);2019))的任务。在第3.4节中，我们展示了鲁棒性基准测试的结果，包括遮挡、ImageNet-A (Hendrycks等人，2021)和对抗性攻击(Goodfellow等人，2014)。最后，我们进行了消融研究，如熵分析、相对距离、标度权重分析和MSA中不同头部的准确性。

3.1 在VIT中插入CB的位置

我们研究了相对于VIT体系结构的主要块(MSA和MLP)， CB的最佳位置。我们用定位在MLP, MSA和两者上的模块训练vits，并在ImageNet-1K上进行验证。如表4所示，无论区块如何，CB都提高了性能，但在MLP区块中实现的准确率比在MSA区块中或两者中都高0.4%p。值得注意的是，无论位置如何，添加CB模块都可以提高ViTs的性能。对于本文的其余部分，我们选择MLP作为CB模块的默认位置。

现在，我们研究了CB模块在MLP块中的最佳位置。MLP块由两个全连接(FC)层和高斯误差线性单元(GELU)非线性激活函数组成(Hendrycks & Gimpel, 2016)。为了简单起见，我们省略了激活函数，我们有三个可能的CB位置:< Front > - FClayer - < Mid > - FClayer - < End >。我们用CB位于前面、中间和后面的ViTS进行训练，并在ImageNet-1K上验证。表5显示了计算成本和top-1精度的结果。与普通的ViT-S相比，中端和端提高了0.6%的准确率。由于MLP层的通道尺寸比前端和端扩展了四倍，因此Mid的计算成本是End的四倍。我们得出结论，在MLP结束时插入CB往往会产生最好的结果

3.2 图像分类

设置。ImageNet-1K (Russakovsky et al .， 2015)由128m张训练图像和50K张验证图像组成，其中包含1K个类。我们用CB模块在训练集上训练vit，并在验证集上报告准确性。在没有其他大数据集的情况下，在ImageNet-1K上训练vit (Deng et al .， 2009;Sun et al .， 2017)，我们采用了DeiT (Touvron et al .， 2021b)设置后的强正则化。我们应用了随机调整大小的裁剪、随机水平翻转、Mixup (Zhang等人，2018)、CutMix (Yun等人，2019)、随机擦除(Zhong等人，2020)、重复增强(Hoffer等人，2020)、标签平滑(Szegedy等人，2016)和drop-path (Huang等人，2016)。我们使用AdamW (Loshchilov & Hutter, 2019)，其beta值为(0.9,0.999)，学习率为10−3·(批大小)=1024，权重衰减为0:05。使用单周期余弦调度来衰减300个总epoch的学习率。我们基于PyTorch (Paszke等人，2019)和timm库(Wightman, 2019)在8个V100 gpu上实现。我们使用torchprofile库来计算flop的数量。更多的细节和额外的实验可以在附录中找到。

结果。我们在ImageNet-1K上用我们的模块训练了viti - ti /-S/-B (Dosovitskiy等人，2021)，并评估了ImageNet-1K、ImageNet-V2 (Recht等人，2019)和ImageNet-Real (Beyer等人，2020)的有效性。我们遵循来自DeiT的viti - ti /-S规范(Touvron等人，2021b)。如表6所示，与普通vit相比，我们的模块CB和CBS提高了性能。CB不改变参数的数量，而CBS增加了一些参数。与原始的flop相比，我们的模块的计算成本可以忽略不计。实验结果表明，该算法对图像分类是有效的。

3.3 语义分割

我们在ADE20K数据集上验证了我们的语义分割方法(Zhou et al .， 2017;2019)，由20K训练图像和5K验证图像组成。为了公平比较，我们采用XCiT (El-Nouby et al .， 2021)和Swin Transformer (Liu et al .， 2021)的方案。我们采用了UperNet (Xiao et al .， 2018)，并训练了40K次迭代或160K次更长的训练。超参数与XCiT相同:批大小为16,AdamW的beta值为(0.9,0.999)，学习率为6·10−5，权衰减为0.01，多项式学习率调度。我们将viti - ti /-S/-B的头部尺寸分别设置为192、384和512。表7显示了40K和160K训练设置下的结果。我们观察到，使用CB的viti - ti /-S/-B在40K迭代时将mIoU分别提高1.0、0.4和0.8，在160K迭代时将mIoU分别提高0.1、0.6和0.6。类似地，CBS改进了mIoU，除了ViT-S中的160K迭代。注入上下文可以改善语义分割。

3.4 评估模型稳定性

我们评估了CB和CBS在中心遮挡(Occ)、ImageNet-A (Hendrycks等人，2021)和对抗性攻击(Goodfellow等人，2014)方面的鲁棒性。对于Occ，我们对每个验证图像的中心112 × 112块进行零掩码。ImageNet-A是resnet50集合无法识别的具有挑战性的测试图像的集合。我们采用快速符号梯度法(FGSM)进行对抗性攻击。表8显示了健壮性基准测试的结果。Occ、ImageNet-A和FGSM的CB分别增加1.0、2.2和5.1;CBS分别为0.7、0.1和0.6。

3.5 消融实验

注意熵和相对距离根据深度。我们在2.1节中观察到，ViT模型中习得注意的熵趋于高。由此，我们假设vit可能受益于明确的均匀注意力注入。我们现在检查我们的CB模块是否降低了学习密集交互的自注意负担。我们比较了有和没有我们的CB模块的ViT模型之间的注意图的熵。图5a显示了有和没有我们的CB模块的ViT-S上的分层熵值。CB模块的插入显著降低了熵值，特别是在较深的层中。CB似乎减轻了MSA块的表示负担，并使其专注于稀疏交互。

我们还观察到，在更深的ViT层中存在注意力熵增加的趋势，如图5a所示。受此启发，我们尝试将CB仅插入到ViT模型的较深层中而不是将它应用到每一层的默认模式。我们把这种选择性插入方法记为CBy。如表9所示，CBy的准确率比普通的viti - ti /-S/-B分别高出1.2%p、0.9%p和0.3%p。与CB的viti - ti /-S/-B相比，CBy还进一步提高了top-1的精度，分别提高了0.2%p、0.3%p和0.1%p。我们仍然认为最实际的解决方案是将CB应用于每一层。然而，如果性能非常重要，那么也可以选择性地将CB插入到层的子集中，特别是更深的层，以获得最佳性能。

我们计算空间相互作用的相对距离，看看CB是否影响空间相互作用的范围。我们定义距离如下:

，其中N为空间标记的个数，aij为第i个标记与第j个标记之间的注意权值，pi为第i个标记的归一化距离。我们排除了自交互的情况来分析其他令牌的交互。如图5b所示，ViT-S和CB具有相似的趋势。CB保持了空间相互作用的范围。
维度缩放分析。我们分析了CBS中尺度权重Λ的大小，以确定根据深度需要统一注意力的趋势。我们测量90%和10%的分位数之比，|λ0.1|=|λ0.9|。比率告诉我们缩放权值的高低有多少相似。我们还根据深度计算缩放权重的平均值。平均值与统一注意力的重要性有关。如图6所示，比值和平均值随深度增加而增加。这意味着上层比下层更喜欢密集的相互作用。所得结果与上述熵分析的观察结果相吻合。

准确度根据正面次数而定。随着注意头数量的增加，MSA可以模拟标记之间丰富的空间交互。为了检验MSA中头部数量与空间相互作用之间的关系，我们通过调整MSA的头部数量来训练有和没有CB的viti - ti。如图7所示，精度差距随着头部数量的增加而减小。因此，我们提出的模块在正面次数较少的情况下比正面次数较多的情况下更有效。

其他上下文广播方法。CB添加平均令牌，但可以采用其他聚合方法，而不是在池和上下文视图中进行平均，全局最大池化是另一种广泛使用的方法来生成上下文。ViT体系结构中的类令牌是上下文信息源的另一种替代方案;类预测最终使用类令牌执行。表10列出了使用max池或类令牌在ImageNet-1K验证集上的前1精度。最大池化使ViT-S的精度提高了0.3%p，但比我们在CB中平均的精度低。类标记的替换甚至降低了0.6%的准确率。我们怀疑这是因为添加了类令牌后，类令牌的值保持不变。我们从2.1节的初步实验中得到的平均值优于其他简单的上下文建模。

4. 相关工作

transformers。自从transformers(Vaswani et al .， 2017)的开创性工作以来，它一直是自然语言处理(NLP)领域的标准架构。Dosovitskiy等
(2021)率先在视觉领域使用变形金刚，他们的视觉变形金刚(ViTs)。ViTs的工作方式几乎与原始的transformer相同，其中ViTs对输入图像的非重叠块进行标记，并在其上应用transformer架构。

有人试图通过将vit与cnn进行对比来理解包括MSA在内的vit的算法行为(Raghu et al, 2021;Cordonnier et al, 2020;Naseer等，2021;Park & Kim, 2022;Tuli et al .， 2021)。Raghu等人(2021)通过经验证明，与cnn相比，全球信息的早期聚合和更大的有效接受域(Luo等人，2016)。Naseer等人(2021)展示了vit对各种干扰的高度鲁棒性行为，包括闭塞、分布转移、对抗性和自然扰动。有趣的是，他们将这些有利的特性归因于MSA在ViTs中的大而灵活的感受场及其相互作用。同样，也有研究将MSA的有效性归因于许多视觉场景理解任务中的全局交互(Carion等人，2020;Strudel et al, 2021)。在这项工作中，我们研究了注意力密度的作用.

注意模块 全局背景对于捕捉对视觉场景的整体理解至关重要Torralba (2003);Rabinovich et al (2007);Shotton et al (2009);Wang et al (2018);曹等(出现)，这有助于视觉识别。为了捕获全局上下文，模型需要设计成具有足够大的接受域来交互和聚合所有本地信息。现有技术已经提出通过更深入地增强cnn的交互范围(Simonyan & Zisserman, 2014;他等人，2016)或通过扩大接受域(Yu & Koltun, 2016;Dai et al .， 2017;Wang et al, 2018)。Hu等人(2018)通过池化来压缩空间维度，以捕捉全球背景。Cao等人(即将出现)注意到，无论查询位置如何，非局部块的注意图都是相似的，并提出了一个全局上下文块。与我们的关注点不同，上述方法关注的是具有额外可学习参数的cnn。

5. 结论

我们仔细研究了ViTs的空间相互作用，特别是在密度方面。我们的探索受到一系列初步观察的激励，这些观察表明ViT模型更倾向于密集的相互作用。顺便说一下，我们还表明，至少从优化的角度来看，均匀注意可能是最难学习的注意。学习密集交互的偏好与优化难度不协调。它引导我们通过一个简单的模块手动引入进一步的密集交互:上下文广播(CB)。CB插入到ViT模型的中间层，将平均池令牌添加到令牌中。此外，我们提出了具有维度缩放的CB，以选择性地注入密集相互作用。我们的简单模块可以全面提高ViT在图像分类和语义分割方面的性能基准。CB只需要1行代码，几个flop和零参数来使用这个模块。我们希望这个实用的模块可以进一步提升您的ViT模型的性能。