华中科技 & 微软 提出 LaViT | 简化注意力机制，提升视觉Transformer效率 ！

视觉 Transformer （ViTs）的出现标志着计算机视觉领域的一个重大范式转变。ViTs通过自注意力模块捕获图像的全局信息，这些模块对分块图像标记进行点积计算。 尽管自注意力模块使ViTs能够捕捉长距离依赖关系，但计算复杂度与标记数量的平方成正比，这是ViTs实际应用的一个主要障碍。此外，深层ViTs中的自注意力机制也容易受到注意力饱和问题的影响。 因此，作者质疑在每个层中计算注意力分数的必要性，并提出了一种称为Less-Attention Vision Transformer（LaViT）的 Transformer ，它仅在每一阶段计算少量的注意力操作，并通过利用先前计算的注意力分数的注意力转换，在其他层中计算后续特征对齐。 这种新颖的方法可以缓解传统自注意力模块面临的两个主要问题：沉重的计算负担和注意力饱和。 作者提出的设计在效率和实施简易性方面具有优势，仅需要矩阵乘法，这在当代深度学习框架中得到了高度优化。此外，作者的架构在各种视觉任务上表现卓越，包括分类、检测和分割。

1 Introduction

近年来，计算机视觉领域经历了快速的增长和发展，这主要得益于深度学习的进步和大规模数据集的可用性。在众多突出的深度学习方法中，卷积神经网络（CNNs）[8]特别有效，在包括图像分类[8, 28]、目标检测[5, 22]和语义分割[1, 23]在内的广泛应用中表现出卓越的性能。

受到在自然语言处理中取得巨大成功的Transformers[27]启发，Vision Transformers（ViTs）[4]将每张图像划分为一组标记。这些标记随后被编码以产生一个注意力矩阵，该矩阵是自注意力机制的基本组成部分。自注意力机制的运算复杂性与标记数量的平方成正比，随着图像分辨率的提高，计算负担也变得更重。一些研究者尝试通过动态选择[9, 21]或标记剪枝[33]来减少标记冗余，以减轻注意力计算的负担。这些方法已经显示出与标准ViT相当的性能。然而，涉及标记减少和剪枝的方法需要精心设计标记选择模块，并可能导致无意中丢失关键标记。在这项工作中，作者探索了不同的方向，重新思考了自注意力机制。在[38]中提出的注意力饱和问题中，随着ViTs的层数逐渐加深，注意力矩阵往往保持较大程度的未变化，反映了前一层中观察到的权重分配情况。考虑到这些因素，作者提出了以下问题：

真的需要在网络的每个阶段从始至终一致地应用自注意力机制吗？

在本文中，作者提出通过引入Less-Attention Vision Transformer（LaViT）来修改标准ViT的基本架构。作者的框架，如图1所示，由Vanilla Attention（VA）层和作者提出的Less Attention（LA）层组成，以捕捉长距离关系。在每个阶段，作者仅在最初的几个Vanilla Attention（VA）层中计算传统的自注意力，并存储注意力得分。在后续层中，作者通过利用先前计算的注意力矩阵高效地生成注意力得分，从而减轻了与自注意力机制相关的平方计算开销。此外，作者在跨阶段的下采样过程中在注意力层内整合了残差连接，允许保留早期阶段学习到的关键语义信息，同时通过替代路径传递全局上下文信息。最后，作者精心设计了一种新的损失函数，以在转换过程中保持注意力矩阵的对角性。这些关键组成部分使得作者提出的ViT模型能够降低计算复杂度和注意力饱和度，最终在使用减少的浮点运算次数（FLOPs）和显著吞吐量情况下实现显著的性能提升。

为了验证作者提出方法的有效性，作者在各种基准数据集上进行了全面的实验，将作者模型的性能与现有的最先进的ViT变体（也包括最近的效率型ViT）进行了比较。实验结果表明，作者的方法在解决注意力饱和问题以及在视觉识别任务中实现卓越性能方面是有效的。

作者的主要贡献总结如下：

• 作者提出了一种新的ViT架构，通过重新参数化先前层计算的注意力矩阵来生成注意力得分。这种方法解决了注意力的饱和度和相关的计算负担。
• 此外，作者提出了一种新的损失函数，旨在在注意力重新参数化过程中保持注意力矩阵的对角性。作者认为这对于维护注意力的语义完整性至关重要，确保注意力矩阵准确反映了输入标记之间的相对重要性。
• 作者的架构在多种视觉任务（包括分类、检测和分割）中始终表现出优于多个最先进的ViT，同时具有相似或甚至减少的计算复杂度和内存消耗。

2 Related Work

作者将在两个广泛类别中讨论相关工作：一是解释深度神经网络预测的方法，二是这些解释在实际场景中的应用。

Vision Transformers

Transformer架构最初是为了机器翻译而引入的[27]，随后通过ViT[4]的发展被应用于计算机视觉任务。ViT的关键创新在于其能够捕获图像中远距离区域之间的长距离依赖关系，这是通过引入自注意力机制实现的。

受到ViT成功的启发，大量变体模型应运而生，每一种都是设计来改善原始架构固有的特定限制。例如，DeiT[25]通过引入蒸馏标记来提高训练时的数据效率。此外，CvT[32]和CeiT[36]将卷积结构整合到ViT框架中，以结合CNN（空间不变性）和ViT（长距离依赖建模）的优点。这些进展凸显了基于Transformer的架构在计算机视觉领域持续进化的趋势。

Efficient Vision Transformers

尽管ViT（视觉 Transformer ）非常有效，但它承受着巨大的计算负担。对高效视觉 Transformer 的研究通过包括分层下采样操作，标记减少，或轻量级架构设计[18, 19]来解决自注意力操作的二次成本。分层下采样操作通过逐渐减少各阶段的标记数量来解决自注意力的二次计算问题，并使ViT能够学习层次结构。另一个研究方向是引入标记选择模块以消除最不具意义的标记并减少计算负担。例如，[9, 21, 33]通过保留具有信息性的图像标记并删除注意力较小的标记来重新组织图像标记，以加速后续的MHSA和FFN计算。

Attention Mechanisms

视觉变换（ViTs）的关键组成部分是注意力机制，该机制计算所有 Patch 之间的成对交互，导致与输入大小成二次方复杂度。这个问题导致了沉重的推理计算，阻碍了ViTs在实际应用中的推广。

一些研究认为，通过利用稀疏注意力机制，可以减轻计算负担，这种机制根据 Patch 的相关性或邻近性有选择地关注一部分 Patch 。一个值得注意的方法是自适应稀疏标记剪枝框架[31]，它诱导了一个稀疏注意力矩阵，有效地解决了计算效率问题。此外，采用结构化稀疏模式等技术可以进一步降低计算复杂度，从而提高ViTs的整体效率。

另一个亟待解决的问题是在层深度增加时，注意力矩阵变化有限的问题，即注意力饱和。DeepViT[38]和CaiT[26]的研究中已经认识到了这个问题，它们报告称注意力饱和阻碍了深层ViTs捕捉额外语义信息的能力，甚至可能降低训练稳定性。因此，在ViTs中仔细设计自注意力机制以避免次优解是至关重要的。

3 Methodology

在本节中，作者首先回顾分层视觉 Transformer 的基本设计。然后，作者讨论其注意力机制的两大弱点，并 Proposal 在每个阶段可以使用较少的注意力来掌握自注意力。

作者通过利用前一层的存储注意力矩阵动态地重新参数化注意力得分，有效地缓解了注意力饱和的问题。

此外，作者还整合了残差连接，以促进早期阶段全局关系的传递。最后但同样重要的是，作者引入了一种新的损失函数，即保持对角线损失，以保持转换后注意力中的基本特性（即表示标记间的关系）。

Vision Transformer

让 表示一个输入图像，其中 表示空间分辨率， 表示通道数。作者首先通过将图像划分为 个块来标记化图像，其中每个块 () 的大小为 像素和 个通道。块大小 是一个超参数，决定了标记的粒度。通过使用步长和核大小等于块大小的卷积算子，可以提取块嵌入。然后每个块通过非重叠卷积投影到嵌入空间 ，其中 表示每个块的维度。

多头自注意力机制。 作者首先简要介绍vanilla自注意力机制，它处理嵌入的块并在多头自注意力块（MHSAs）的框架内工作。在第 个MHSA块中，输入 被投影到三个可学习的嵌入 。多头注意力旨在从不同的视角捕捉注意力；为了简单起见，作者选择 个头，其中每个头是一个维度为 的矩阵。第 个头注意力矩阵 可以通过以下方式计算：

和 分别是第 个头的注意力矩阵、 Query 和键。作者还把值 分割为 个头。为了避免由于概率分布的尖锐性导致的梯度消失，作者将 和 的内积除以 ()。注意力矩阵按以下方式连接：

在空间分割的标记之间计算的注意力可能指导模型关注视觉数据中最有价值的标记。随后，对相应的值 应用加权线性聚合：

下采样操作。 一些研究将分层结构引入ViTs，从CNN中分层架构的成功中汲取灵感。这些工作将Transformer块划分为 个阶段，并在每个Transformer阶段之前应用下采样操作，从而减少序列长度。在作者的研究中，作者使用具有核大小和步长设置为 的卷积层进行下采样操作。这种方法允许在每个阶段灵活调整特征图的尺度，从而建立一个模仿人类视觉系统组织的Transformer分层结构。

The Less-Attention Framework

作者网络架构的总体框架如图1所示。在每一个阶段，作者通过两个阶段提取特征表示。在最初的几个Vanilla Attention（VA）层中，作者执行标准的MHSA操作来捕捉整体的长距离依赖关系。随后，在接下来的Less-Attention（LA）层中，作者模拟注意力矩阵以减轻二次计算并解决注意力饱和问题，通过将线性变换应用于存储的注意力得分。在这里，作者将第个阶段的初始-th VA层Softmax函数之前的注意力得分表示为，它是通过以下标准过程计算的：

这里，和分别表示来自第个阶段的-th层的 Query 和键，跟随前一阶段的下采样。用来表示VA层的数量。在最初的Vanilla注意力阶段之后，作者摒弃了传统的二次MHSA，并对进行变换以减少注意力计算的数量。这个过程包括两个线性变换以及在它们之间的矩阵转置操作。例如，考虑第个阶段中-th层（，即LA层）的注意力矩阵：

\begin{split}&\mathbf{A}_{m}^{l}=\Psi(\Theta(\mathbf{A}_{m}^{l-1})^{ \text{T}})^{\text{T}},\;\;L_{m}^{\text{VA}}<l\leq l_{m},\\="" &\mathbf{z}^{\text{la},l}="\text{Softmax}(\mathbf{A}_{m}^{l})" \mathbf{v}^{l}.\end{split}="" \tag{5}

在这个背景下，和表示的变换是指具有维度的线性变换层。这里，，分别表示第个阶段的层数和VA层的数量。在这两个线性层之间插入转置操作是为了保持矩阵的相似性行为。这一步是必要的，因为单层中的线性变换是按行进行变换的，这可能会导致丧失对角线特征。

Residual-based Attention Downsampling

在分层ViTs的计算跨阶段时，通常会在特征图上采用下采样操作。虽然这种技术减少了标记数量，但它可能导致丢失重要的上下文信息。因此，作者提出，从前一阶段学习的注意力亲和力可能对当前阶段捕获更复杂的全局关系有利。受到ResNet [7]的启发，它引入了短路连接来解决特征饱和问题，作者也采用了类似的概念，并将其融入到作者架构中的下采样注意力计算中。通过引入短路连接，作者可以将内在偏见引入到当前的MHSA块中。这允许前一阶段的注意力矩阵有效地指导当前阶段的注意力计算，从而保留关键上下文信息。

然而，在这种情境中直接将短路连接应用于注意力矩阵可能会带来挑战，主要是因为当前阶段与前一阶段的注意力维度存在差异。在这里，作者设计了一个注意力残差（AR）模块，它由深度卷积（DWConv）和卷积层组成，用于在下采样前一阶段的注意力图时保持语义信息。作者将前一阶段（第阶段）的最后一个注意力矩阵（在层）表示为，当前阶段（第阶段）的下采样初始注意力矩阵表示为。的维度为（表示第阶段的标记数量）。作者将多头维度视为常规图像空间中的通道维度，因此使用DWConv操作符（步长，核大小），作者可以在注意力下采样时捕捉标记之间的空间依赖关系。经过DWConv变换后的输出矩阵适合当前阶段的注意力矩阵大小，（在作者的情况下）。在注意力矩阵上进行深度宽度卷积后，作者执行卷积以跨不同头交换信息。

作者的注意力下采样在图2中说明，从到的转换可以表示为：

其中LS是[26]中引入的层缩放操作符，用于减轻注意力饱和。是第阶段第一层的注意力分数，它是通过将标准MHSA与方程4以及由方程6计算的残差相加得到的。

两个基本设计原则指导了作者的注意力下采样模块。首先，作者利用DWConv在下采样时捕捉空间局部关系，从而有效地压缩注意力关系。其次，利用操作交换跨 Head 的注意力信息。这个设计至关重要，因为它促进了从前一阶段到后续阶段的注意力有效传播。将残差注意力机制融入现有ViT主干只需进行少量调整，通常只需添加几行代码。值得强调的是，这种技术可以无缝应用于Transformer架构的各种版本。唯一的前提是存储来自前层的注意力分数，并相应地建立与此层的短路连接。这个模块的重要性将通过全面的消融研究进一步阐明。

Diagonality Preserving Loss

作者通过结合注意力转换算子精心设计了Transformer模块，旨在减轻计算成本和注意力饱和的问题。然而，一个紧迫的挑战仍然存在——确保转换后的注意力保持标记间的相互关系。众所周知，对注意力矩阵应用转换可能会损害它们捕捉相似性的能力，这主要是因为线性变换是逐行处理注意力矩阵的。因此，作者设计了一种替代方法，以确保转换后的注意力矩阵保留传递标记间关联所必需的基本属性。一个传统的注意力矩阵应该具备以下两个属性，即对角线性和对称性：

因此，作者设计了第层的对角线保持损失，以保持这两个基本属性：

这里，是Diagonality Preserving（对角线保持）损失，旨在保持方程8中注意力矩阵的性质。作者在所有转换层上添加了对角线保持损失和原始交叉熵（CE）损失[4]，因此作者训练中的总损失可以表示为：

其中是最后一层表示中的分类标记。

Complexity Analysis

作者的架构由四个阶段组成，每个阶段包含层。在连续的每个阶段之间应用下采样层。因此，传统自注意力的计算复杂度是，而相关的K-Q-V转换的复杂度为。相比之下，作者的方法在转换层中使用一个线性变换，从而避免了计算内积的需要。因此，作者转换层中注意力机制的的计算复杂度降低到，减少了倍。此外，由于作者的方法仅在Less-Attention层内计算 Query 嵌入，作者的K-Q-V转换复杂度也同样减少了倍。

在连续阶段之间的下采样层中，以的下采样率为例，注意力下采样层中DWConv的计算复杂度可以计算为。同样，注意力残差模块中卷积操作的复杂度也是。然而，重要的是要注意，每个阶段只发生一次注意力下采样。因此，与Less-Attention层实现的复杂度降低相比，这些操作引入的额外复杂度是微不足道的。

4 Experiments

在本节中，作者评估了作者的模型在两个基准数据集上的性能：用于分类的ImageNet-1K [8]，用于检测的COCO2017 [10]以及用于分割的ADE20K [37]。作者将作者的模型与其他在这些数据集上的最先进工作进行了比较，以展示其有效性和效率。此外，作者还进行了消融研究，以探究所提出模型中每个组件的必要性和贡献。这种分析对于了解每个部分的作用以及帮助确立作者方法的有效性提供了宝贵的洞见。

Architecture Variants

为了在保持相似计算复杂度的情况下与其他模型进行公平比较，作者建立了三个模型：LaViT-T、LaViT-S和LaViT-B。

详细的配置信息在表1中提供，作者遵循与PVT [29, 30]相同的网络结构，不同的是引入了Less-Attention Transformer编码器和跳跃连接注意力下采样。块的数目、通道数和头数会影响计算成本。

Baselines

作者通过将作者提出的方法与各种卷积神经网络（CNNs）、视觉 Transformer （ViTs）和分层视觉 Transformer （hierarchical ViTs）进行比较，进行了彻底的实验评估。具体来说，以下基准被使用：

• 卷积神经网络： ResNet [7]，RegNet [20] 和 EfficientNet [24]。
• 视觉 Transformer ： ViT [4]，DeiT [25]，CvT [36]，DeepViT [38]，FocalViT [35] 和 SwinTransformer [13]。
• 高效视觉 Transformer ： HVT [17]，PVT [29]，DynamicViT [21]，EViT [9]，LiT [19]，EfficientViT [2] 和 PPT [33]。

Image Classification on ImageNet-1K

设置. 图像分类实验是在ImageNet-1K数据集上进行的。作者的实验协议遵循DeiT [25]中概述的流程，但模型本身除外。具体来说，作者采用了与DeiT相同的数据增强和正则化技术。作者使用AdamW优化器[16]从头开始训练作者的模型300个周期（包含5个周期的预热）。初始学习率设置为0.005，并根据余弦调度器进行调整。全局批量大小设置为1024，分布在4个GTX-3090 GPU上。在验证集上进行测试时，先将输入图像调整到256像素，然后进行224 x 224像素的中心裁剪以评估分类准确性。

结果. 作者在表2中展示了在ImageNet-1K上的分类结果。根据它们的计算复杂度，模型被分为三组：微小（约2G）、小型（约4G）和基础（约9G）。与显著减少的计算要求相比，作者的方法在与现有最先进的ViTs相比具有竞争力的性能。具体来说，在微小和小型模型群组中，作者的方法分别至少超过了所有其他现有模型0.2%和0.5%，同时保持了显著降低的计算成本，这是作者的主要关注点。在基础尺寸模型中，作者的架构采用了PVT的基础结构，但包含了Less-Attention组件，其性能优于两种基于PVT的模型（PVT-M和PVT-L）。此外，作者还将在几个高效的ViT设计与作者的架构进行比较。作者观察到，作者的结果在有效性和效率之间取得了更好的平衡。请注意，由于作者资源高效的Less-Attention机制，作者的设计需要降低计算成本，使得作者轻量级的模块成为在移动平台上实现ViT的一个吸引人的选择。

Object Detection on COCO2017

设置. 作者在COCO 2017 [10] 数据集上进行了检测实验。作者测试了RetinaNet [12]上模型的的有效性。作者遵循通常的做法，使用从ImageNet-1K获得的预训练权重初始化主干网络。此外，作者使用了AdamW [16] 优化器，并在8个GPU上以批量大小为16来训练网络。

结果. 表3展示了目标检测的结果。很明显，作者的LaViT模型在与CNN和Transformer的对应模型相比有显著的优势。具体来说，在使用计划时，作者的小型版本LaViT-T在可比设置下相对于ResNet实现了9.9-12.5 AP的提升，而小型版本LaViT-S相对于其CNN对应版本提高了8.1-10.3 AP。在使用计划时，这一趋势持续存在，作者的LaViT模型始终表现出具有竞争力的性能。特别值得一提的是，作者的架构在检测性能上始终优于Swin Transformer，同时训练负担较小。因此，在COCO2017上的结果再次确认了### ADE20K上的语义分割

设置.作者在ADE20K数据集上进行了语义分割实验，该数据集包含150个类别，训练有20,000张图像，验证有2,000张图像。作者用于分割的主干网络是Semantic FPN [11] 和 UperNet [34]。作者遵循[13]中建立的训练设置，并将图像大小调整为进行训练。作者训练UperNet 160k次迭代，SemanticFPN 80k次迭代。初始学习率设置为，使用多项式调度器进行学习率衰减。实验是通过在个GTX3090 GPU上使用批量大小为16进行的。

结果.表4提供了分割结果的概览。作者的模型在Semantic FPN和UperNet上均优于Swin Transformer，分别实现了mIoU +2.6和+2.7的提升。在Semantic FPN测试中，作者的LaViT-S相对于 Baseline （PVT-S）实现了相对适度的mIoU +0.9的提升，但显著减少了计算量。当集成到UperNet架构中时，LaViT相对于各种主流模型实现了mIoU +2.7、+1.0和+1.4的显著提升。即使在使用测试时间增强时，这些有竞争力的结果也能保持。特别是，LaViT-S在mIoU上比Focal-T高出+1.4，在MS mIOU上高出+2.5。这些发现强调了LaViT在计算效率高的注意力机制框架内产生高质量语义分割输出的能力。

Ablation Study

注意力饱和。为了证明作者的Less-Attention模块在解决注意力饱和问题上的有效性，作者在图3中展示了注意力相似度比（由当前层和前一层注意力图计算出的余弦相似度）。作者使用两种 Backbone 网络进行比较，分别是ViT和PVT。在3a中，作者选择了具有25层且无层次结构的ViT架构。在3b中，作者以PVT-M作为 Baseline ，在第3阶段评估注意力相似度，该阶段由18层组成。两组结果清楚地表明，原始架构遇到了显著的注意力饱和问题。然而，通过融入作者的模块，有效地缓解了这一现象，使深层注意力能够发挥其预期作用。

Less-Attention模块的可扩展性。作者将Less-Attention模块扩展到各种ViT架构中，并在表5中报告了结果。将Less-Attention层整合到任何基础Transformer架构中都能提高准确度，同时减少计算需求。值得注意的是，当将模块整合到Vanilla ViT/DeiT架构中时，改进最为显著。这可能是因为Vanilla ViT/DeiT没有层次结构，因此会遇到较大的注意力饱和问题。此外，将作者的方法整合到DeepViT中，作者看到计算资源的最大减少。这些发现共同强调了作者的方法的可扩展性，表明将LA模块应用于现有ViT架构可以使其更加实用和可行。

各个组件的重要性。作者对所提出的模块使用ImageNet-1k数据集进行了消融研究，结果如表6所示。

对于两种网络（小型和微型），作者的模块对Transformer的训练是不可或缺的。将Less-Attention层替换为MHSA的 Baseline ，其预测准确度分别下降了0.5%和0.6%。此外，移除注意力残差模块（表示为“w/o AR”）会导致预测准确度下降0.2%和0.4%。最后，作者认为，为保留对角线上的相关性而增加的损失函数对于有效理解视觉数据中的语义信息至关重要。当仅依赖CE损失时，模型的预测会恶化。这可能是因为仅依赖变换矩阵来捕捉注意力可能会损害它们表达 Token 之间关联的能力。所有这些实验性发现共同强调了模型架构中每个组件的贡献。

Less-Attention选择。在Deep ViTs中，仔细选择Less-Attention的起始层至关重要。因此，作者设计了实验来选择网络架构中第二阶段的起始层，结果如表7所示。

如表所示，直接从第二层使用Less-Attention层会导致模型性能下降。这一现象可能是因为过度依赖第一层MHSA的语义。因此，在更深层次利用Less-Attention层可能可以缓解这一问题。此外，在相对较深的层次使用Less-Attention层不会对模型性能产生太大影响，但可能导致计算成本增加。这与作者减少计算开销的架构设计目标相矛盾。

5 Conclusion

为了减少昂贵的自注意力计算，作者提出了一种新模型，称为Less-Attention Vision Transformer (LaViT)。LaViT利用了多头自注意力（MHSA）块中计算的依赖性，并通过复用先前MSA块的注意力来跳过注意力计算。

此外，作者还引入了一个直接保持对角线损失，旨在促进注意力矩阵在表示标记之间关系时的预期行为。值得注意的是，Transformer架构有效地捕捉了跨标记关联，在保持参数数量和每秒浮点运算（FLOPs）的计算效率的同时，超越了基线性能。

全面的实验已经证实了我们的模型作为一个基础架构在多个下游任务中的有效性。

具体而言，所提出的模型在分类和分割任务上表现出了相对于先前Transformer架构的优越性，从而在分类和分割任务上取得了最先进的性能。

参考

[1].You Only Need Less Attention at Each Stage in Vision Transformers.