点击下方卡片，关注
「AI视界引擎」
公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
在现代视觉Transformer设计中，Softmax注意力能够有效捕获长距离的视觉信息；然而，在处理高分辨率输入时会带来巨大的计算成本。
相比之下，线性注意力自然享有线性复杂度，并且具有扩展到更高分辨率图像的巨大潜力。
尽管如此，线性注意力表现不佳极大地限制了其在各种场景中的实际应用。本文通过新颖的理论分析，努力弥合线性注意力与Softmax注意力之间的差距，揭示了性能差异背后的核心因素。具体而言，作者从两个关键角度理解和缓解线性注意力的局限性：注射性质和局部建模能力。
首先，作者证明了线性注意力不具备注射性质，容易将不同的 Query 向量分配相同的注意力权重，从而导致严重的语义混淆，因为不同的 Query 对应相同的输出。
其次，作者证实了有效的局部建模对于Softmax注意力的成功至关重要，而线性注意力在这方面显得不足。上述两种基本差异显著影响了这两种注意力机制之间的差异，这一点通过作者在本文中进行的大量实验证明得到了验证。
此外，更多的实验结果表明，只要赋予线性注意力这两种特性，它可以在多种任务上超越Softmax注意力，同时保持较低的计算复杂度。
相关代码可在https://github.com/LeapLabTHU/InLine获取。
unset
unset
1 Introduction
unset
unset
近年来，Transformer和注意力机制在计算机视觉领域取得了前所未有的成功[7]。Softmax 注意力，又称点积注意力，展现出了卓越的表达能力，在各种视觉任务中达到了最先进的性能。然而，将 Softmax 注意力应用于视觉领域也面临着挑战。Softmax 注意力的二次复杂度导致当其应用于全局感受野时会产生不可承受的高计算成本。之前的研究工作[33, 19, 36, 6, 45]通过限制感受野或引入稀疏性努力降低计算复杂度。尽管这些方法有效，但它们不可避免地牺牲了 Softmax 注意力进行长距离建模和扩展的能力。
Softmax 自注意力的本质要求首先计算 Query 和键之间的点积
，然后通过归一化的分数
来计算，在序列长度
方面具有二次复杂度
。相反，线性自注意力放松了 Query
和键
之间的相似性评分，从 Softmax 变换为其他可以分解为核函数的形式。具体而言，线性自注意力用核函数
替换原始评分函数
为
。这种替代使得根据矩阵乘法的结合律改变计算顺序，即从
到
，从而将复杂度从
降低到
关于序列长度
。尽管如此，每一枚硬币都有两面。线性自注意力的效果通常不如 Softmax 自注意力 [3, 27, 40]，其较差的表达能力限制了其实用价值。
虽然有大量研究试图以不同方式缓解这一问题，但作者仍然不清楚导致线性与 Softmax 自注意力之间差距的关键因素。
在本文中，作者探讨了线性注意力和Softmax注意力之间的根本差异，并从两个洞察角度揭开这一课题的面纱：单射性质和局部建模能力。首先，作者将注意力视为一个映射 Query 到注意分数的函数。作者发现，该注意力函数的单射性质极大地影响了模型的表现。
具体而言，如果注意力函数不是单射的，不同的 Query 将诱导出相同的注意力分布，导致特征空间内的严重语义混淆。作者严谨的分析已经证明，Softmax 注意力函数是单射函数，而线性注意力函数则不是。因此，线性注意力容易受到语义混淆问题的影响，这大大限制了其表达能力。其次，作者对注意力权重分布的分析确认，Softmax 注意力的成功不仅仅依赖于其强大的长程建模能力。有效的短程建模也是实现最优结果的关键因素。
为了验证作者的分析，作者提出两种简单而有效的方法分别赋予线性注意力注入性质和局部建模能力。广泛采用的Swin Transformer [19] 架构被用来验证作者的发现。实验结果突显了这两种性质在线性注意力与Softmax注意力之间的差距中的重要性。此外，全面的实验表明，赋予这两种性质的线性注意力在多种任务中优于广泛使用的Softmax注意力。
作者的主要贡献和收获总结如下：
(1) 注射性是线性注意力和Softmax注意力之间的一个关键差异。尽管Softmax注意力具有注射性，但线性注意力的非注射性特性导致了语义混淆，并严重削弱了模型性能。据作者所知，作者的工作首次将注意力机制视为映射函数，并证明了其注射性属性的重要性。
(2) 局部建模对于注意力机制的有效性仍然至关重要，即使它以拥有大感受野和出色的长程建模能力而闻名。
(3) 作者挑战了线性注意力劣于Softmax注意力的观点，并通过上述两个属性证明，在保持较低计算复杂度的同时，线性注意力可以超越Softmax注意力。
unset
unset
2 Related Works
unset
unset
视觉Transformer和Softmax注意力机制。视觉Transformer[7]是最早将自注意力引入视觉领域的开创性工作。此后，注意力机制在各种视觉任务中取得了成功[7, 2, 24]。广泛使用的注意力机制是Softmax注意力机制[32]，也被称为点积注意力，它计算所有 Query 键对之间的相似性。虽然效果显著，但由于其计算复杂度呈平方级增长，在处理全局特征图时会导致不可管理的成本。因此，提出了多种方法以减少Softmax注意力机制的计算开销。PVT[33]通过下采样键值来降低计算复杂度。Swin Transformer[19]通过引入窗口注意力模式限制感受野。NAT[14]模仿卷积并在每个特征的邻域内计算注意力，而DAT[36]则提出了一种输入依赖的稀疏注意力模式。
线性注意力。与softmax注意力相比，线性注意力是一种具有自然线性复杂度
的另一种注意力范式。线性注意力通过使用核函数替代softmax，从而通过对计算顺序进行改变来降低计算复杂度至线性。然而，先前的研究[27, 3, 29, 40, 11]表明，线性注意力的表现明显劣于softmax注意力。
CosFormer归因于这一差异是由softmax注意力高效的重权机制导致，并提出使用余弦重权来增强线性注意力。Nystromformer[38]和SOFT[22]利用矩阵分解进一步逼近softmax操作。Efficient Attention[28]对queries和keys应用softmax函数。TransNormer[26]发现未加界的梯度和注意力稀释损害了线性注意力。FLatten Transformer[9]引入了一个聚焦函数以解决过度平滑的问题。
MLLA[10]借鉴了Mamba[8]的灵感以改进线性注意力。
尽管线性注意力和Softmax注意力取得了优雅的结果，但两者之间的根本差异原因仍然不清楚。在本文中，作者对线性注意力和Softmax注意力之间的差异进行了深入分析，并确定了高性能Softmax注意力的两个关键属性：注入性和局部建模能力。作者提供了理论证明和实验验证来支持作者的发现。
unset
unset
3 Preliminaries
unset
unset
基于此，Softmax 注意力机制 [32] 计算注意力权重，并计算输出为其值的加权和。
为了简化起见，作者省略了
中的
，因为作者可以等价地重新归一化
和
。这种注意力机制在现代视觉Transformer中非常成功。然而，它应当计算所有 Query 键对之间的相似性，从而导致
的复杂度。因此，使用全局感受野的Softmax注意力会带来巨大的计算成本。
线性注意力机制[15]被提出以通过具有线性复杂度
的方式高效地处理计算挑战。具体而言，
被替换为
，其中
是一个核函数。通过这种方式，线性注意力机制重新表述了式(1)为：
的形式表明，显式计算注意力权重
是不必要的。相反，作者可以利用矩阵乘法的结合性，将计算顺序从
变更为
，从而使计算复杂度降低到
。注入性质：设
是一个映射函数。作者称
为注入函数，当且仅当对任意
，若有
，则
。
unset
unset
4 Analysing the Gap between Linear and Softmax Attention
unset
unset
由于线性计算复杂度较低，线性注意力被视为解决高分辨率场景中softmax注意力计算挑战的有前途的解决方案。然而，先前的研究[27, 3, 40]表明，线性注意力的表示能力远低于softmax注意力，使其在实际应用中不可行。在本节中，作者从两个视角深入分析了线性注意力与softmax注意力之间的差距：注射映射和局部建模能力，并提供了理论证明和实验验证来增强对这两种注意力机制之间关键差异的理解。
4.1 Injectivity of Attention Function
作者首先定义Softmax函数和线性注意力的作用如下：
其中，
表示 Query，而
分别是式 (1) 和式 (2) 中的注意力分数。给定键
，
和
可以视为函数 fquery
映射每个
到其相应的 Softmax 和线性注意力分数
和
。然后，Softmax 和线性注意力针对
的最终输出可以表示为
和
。注入性质。在本文中，作者发现注意力函数的注入性质对模型性能有显著影响，这可能很大程度上解释了线性注意力和 Softmax 注意力之间的差距。具体来说，在温和的假设下，作者证明 Softmax 注意力函数
是注入的，而线性注意力函数
不是（命题 1 和 2，请参阅附录以获取完整证明）。因此，对于两个不同的 Query
和
，Softmax 注意力应该产生不同的注意力分布
，而线性注意力可能会产生相同值
。由于典型的不同 Query
代表不同的语义，线性注意力的非注入性质实际上会导致语义混淆，即
且
，使模型无法区分某些语义。
作者通过一个例子来更好地理解注入性质和混淆问题。如图1(a)所示，有四个共线且长度不同的向量。借助注入性质，Softmax注意力确保每个 Query 获得不同的注意力分数，从而生成更聚焦的注意力分布，特别是对于较长的 Query 。然而，当使用核函数
时，线性注意力无法区分强度不同但语义相同的 Query ，即具有不同长度的不同强度的共线 Query ，导致这四个 Query 得到完全相同的关注度分数。因此，线性注意力无法为更强的语义生成更聚焦的关注度得分，这可能解释了在[9]中讨论的缺乏聚焦能力的问题。当使用具有更强非线性的核函数时，线性注意力会遇到更为显著的混淆问题。例如，在图1(b)中，使用核函数
时，线性注意力会给方向和长度不同的四个 Query 分配完全相同的关注度分数。这种严重的语义混淆可以直接损害模型的性能。
实际模型中的混淆问题。虽然图1展示了混淆的概念，但作者也需要验证这种问题是否存在于实际模型中。因此，作者在Deit-T上进行了统计分析。在对ImageNet验证集进行推理时，作者计算了每张图像中混淆发生的次数（即
但
或
）。
由于两个浮点数向量在实际中很少严格相等，所以作者认为当它们之间差的L2范数小于1e-3时，它们可以被视为大致相等。
结果如图2所示。几乎所有的样本在使用Softmax注意机制的模型中没有遇到混淆，而在线性注意机制的模型中，有大量的样本出现了超过
次的混淆情况。这证明了在线性注意机制的实际模型中确实存在混淆问题。
根据表1 显示，引入混淆会导致性能明显下降，凸显了注意力函数注入性属性的重要性。因此，线性注意力的非注入性可能是导致其有限表达能力的一个关键因素。
使线性注意力函数成为一一对应的。作者提出了一种简单且有效的解决方案，使得线性注意力成为一个一一对应的函数。命题2（附录中证明）表明，在线性注意力中由于忽略了除法中的
，对于所有的
，
会获得相同的分数，从而导致非一一对应性。因此，作者简单地将线性注意力的规范化从除法改为减法，提出了作者的线性注意力（InLine），如下所示：
并且，与
相对应的注意力输出可以写为
。这一修改确保了注意力权重仍然相加为1，同时将注意力函数转换成了单射函数（参见命题3）。因此，单射线性注意力能够区分不同的 Query ，类似于 Softmax 注意力，并且不再受到混淆问题的影响。
命题 3（Inline 注意力是单射的）设
是一个单射映射。给定
，有
和
。对于所有
，当
时，有
。此外，类似线性注意力，Inline 注意力可以通过改变计算顺序以
复杂度来计算。
由于作者可以计算
和
，因此 InLine 注意机制的时间复杂度为
。考虑到多头注意力机制，该复杂度变为
，其中
和
分别是模型的通道维度和每个头的维度。
4.2 Local Modeling Capability
注意力机制以其大的感受野和出色的长程建模能力而著称。然而，作者发现有效的局部建模对于注意力的有效性至关重要。
在图3中，作者使用DeiT-T计算每个Query分配给局部3×3邻域的注意力值之和。每个DeiT-T注意力层包含总共14×14+1=197个Token，如果注意力分数随机分配，则每个 Query 的3×3邻域的期望注意力值为9/197。结果显示，所有三种注意力机制都倾向于更多关注每个Query的邻域，揭示了局部偏见，尤其是在浅层。值得注意的是，Softmax注意力分配了大量的注意力给局部窗口，表明其局部建模能力比两种其他注意力机制更强。图4提供了相应的可视化结果以进一步确认这一发现。作者认为，Softmax注意力表现更优的原因可能是它具有稳健的局部先验和强大的局部建模能力。为了验证这一假设，作者采用注意力 Mask 去除不同位置的Token，并评估它们对模型性能的影响。结果如表2所示。两个关键观察结果是：
- 去除局部Token显著降低了模型性能，而随机去除相同数量的Token则影响较小；
- 当去除局部Token时，Softmax注意力的表现比InLine注意力受损更严重。这些发现证明了局部建模对于两种注意力类型的重要性，并证实了Softmax注意力相对于InLine注意力的优势主要来源于其更强的局部建模能力。
基于作者的分析，增加局部偏置可能增强 InLine 注意力的表达能力。因此，作者采用 MLP 预测额外的局部注意力残差以增强 InLine 注意力。
具体地，对应于
的输出定义为：
其中，
表示 Inline 注意力函数，
是输入 Token 的平均值，
是预测的局部注意力残差，而
表示
的
局部邻域中的值。通过引入局部注意力残差项，作者显式地增强了 Inline 注意力的局部偏差。作者将引入了局部注意力残差的 Inline 注意力称为 Inline 注意力模块。由于局部残差项带来的计算成本仅为
，因此 Inline 注意力模块仍然保持线性复杂度
。
unset
unset
5 Empirical Study
unset
unset
在第4节中，作者分析了Softmax与线性注意力之间性能差距背后的核心因素，并提出了可能的解决方案。本节中，作者将进行实证验证，以全面验证这两项属性的重要性以及作者方法的有效性。
5.1 Implementation
作者利用流行的Swin Transformer架构[19]来探讨注入性和局部建模能力的影响。具体而言，作者将SwinT中的原始Softmax注意力替换为线性注意力，建立 Baseline 模型。
随后，作者依次引入注入性质和局部偏置以评估它们各自的影响力。为了全面验证InLine注意力模块的有效性，作者将其应用到包括DeiT[30]、PVT[33]、Swin[19]、CSwin[6]在内的四个先进且具有代表性的Transformer模型中，并使用Softmax注意力与各种最先进的方法进行广泛的比较。
5.2Datasets and Experiment Details
ImageNet图像分类。ImageNet-1K数据集包含128万张训练图像和5万张验证图像，共有1000个类别。为了进行公平比较，作者使用与基准模型相同的设置来训练作者的模型。作者使用AdamW优化器从零开始训练所有模型共300个周期，并采用余弦退化学习率策略，其中包含20个周期的线性Warm Up预热期。初始学习率为
，权重衰减率为0.05。增强和正则化策略包括RandAugment、Mixup、CutMix和随机擦除。在InLine-CSwin模型的训练过程中，使用了EMA策略，仿效CSwin的做法。
COCO目标检测。COCO [18]目标检测和实例分割数据集包含118,000张训练图像和5,000张验证图像。作者遵循相应 Baseline 模型的训练和测试策略，并采用预训练的InLine Backbone 进行实验。
ADE20K语义分割。ADE20K [44] 是一个广泛认可的语义分割基准，包含了20,000张训练图像、2,000张验证图像和150个语义类别。 Baseline 模型相同的设置被采用。
5.3 Empirical Analysis of Injectivity and Local Modeling
注射性特性。如表3所示，作者采用了四种不同的核函数
来验证注射性的效果。如第4.1节所述，当使用核函数
时，线性注意力无法区分具有不同强度的相同语义。通过引入InLine注意力解决了这一问题，准确率提高了2.5个百分点。当使用
时，线性注意力面临更为严重的语义混淆，引入注射性特性后，准确率显著提升，从70.2%提高到80.0%，提升了9.8个百分点。
这些明显的改进完全证明了注射性的重要性，并验证了作者提出的注射性线性注意力的有效性。作者也使用了两个不保证非负性的核函数。与[27]中的发现一致，没有非负性保证的情况下，线性注意力无法收敛，作者认为这是由于极端的语义混淆导致的。例如，当
时，线性注意力无法区分完全相反的语义，对
和
赋予相同的注意力分数。
局部建模能力。表4突显了局部建模能力的重要性。在左侧表中，作者将纯InLine注意力应用于Swin-T，并逐步增加窗口大小从
到
。由于InLine注意力具有线性复杂度，因此可以在保持相同计算成本的同时采用不同的窗口大小。较大的窗口大小会带来更大的感受野，通常与性能改进相关。
然而，实验结果显示，随着窗口大小的增大，模型性能并未提升。作者认为这可以归因于局部建模能力不足：较小的窗口大小限制了感受野，但引入了强烈的局部偏见，增强了局部建模；而较大的窗口大小虽扩大了感受野，但也进一步削弱了局部建模能力。为了验证这一点，作者在右侧表中应用了带有局部残差的InLine注意力，并呈现了结果。引入局部残差项后，观察到了显著的改进。此外，在引入局部残差项后，随着窗口大小的增加，性能表现持续提升，这强有力地支持了作者的分析。
5.4Main Results and Broad Comparisons
如表4所示，采用局部残差的InLine-Swin-T比Swin-T Baseline 取得了更好的效果，分数从81.3提升到了82.4。因此，作者想知道InLine注意力模块是否能在各种场景中优于广泛采用的Softmax注意力机制。为了验证这一点，作者进一步将其应用于几个代表性Transformer，并在图像分类、目标检测和语义分割等方面进行了全面比较。
图像识别任务。首先，作者将InLine注意力模块应用于DeiT [30]、PVT [33] 和 Swin Transformer [19]，并将结果展示在表5中。可以看出，用作者的方法替换Softmax注意力后，取得了显著的改进。例如，InLine-PVT-S在参数量仅为PVT-L的30%和FLOPs仅为40%的情况下，性能超过了PVT-L。
随后，作者将在CSwin Transformer [6]等先进的Transformer设计中应用该模块，并与各种最新的模型在ImageNet-1K上进行广泛的对比。如表6所示，作者的InLine-CSwin模型不仅在性能上优于CSwin，还超越了多种最新的CNN和Transformer设计。这些结果表明，作者的InLine注意力模块有可能成为广泛使用的Softmax注意力的一种更优越的替代方案。
推理吞吐量分析。作者提供了实际的速度测量值，详见图5。如图5(a)所示，InLine模型在准确性和延迟之间表现出明显更好的权衡。在图5(b)中，作者将窗口大小从
增加到
。由于Softmax注意力机制的二次复杂度，Swin-T的速度随着窗口的增大而急剧下降。相反，具有线性复杂度的InLine-Swin-T即使在较大窗口的情况下仍能展现出更高的速度，这可能是因为窗口分割导致的延迟降低所致。凭借全局的感受野，InLine不仅受益于高性能（见表4），还能实现快速的速度。此外，如图5(c)所示，在高分辨率场景下，InLine展现出显著的计算优势。
COCO目标检测。表7显示，在目标检测任务中，InLine注意力机制始终能够提高检测效果。例如，在相似的FLOPs条件下，InLinePVT-S的盒体AP值比PVT-T高出6.7%，而InLine-PVT-L在更少的FLOPs条件下，盒体AP值比PVT-M高出3.4%，这表明InLine注意力机制在线性复杂度下的优势，特别是在高分辨率场景中。
ADE20K语义分割。作者将模型应用于两个代表性的分割模型，SemanticFPN [16] 和 UperNet [37]。如表8所示，得益于注入性和有效的局部建模能力，Our InLine 在所有设置下都取得了更好的结果，并且计算成本明显更低。
与当前最佳的线性注意力设计进行比较。如表9所示，作者的简单InLine注意力设计在没有额外复杂功能的情况下，优于各种线性注意力方法。此外，作者的InLine注意力可以与其他先前的设计相结合以取得更好的效果，这留待未来的工作进一步探讨。例如，FLatten中的高级聚焦功能也可以应用于InLine注意力中。
5.5 AblationStudy
作者两大关键设计的有效性已在第5.3节中经过验证并进行了详细的分析。在表10中，作者提供了额外的结果以验证不同 Kernel 函数的影响。结果显示，作者的InLine注意力机制能够有效与不同的 Kernel 函数配合工作，进一步验证了该方法的有效性。ReLU和指数函数取得了稍微更好的效果。在本文中，为了简化起见，默认使用Identity()。
在本文中，作者探讨了导致线性注意力和Softmax注意力性能差距的核心因素。作者识别并验证了这两种注意力机制之间存在的两大基本差异：注入性质和局部建模能力。
注入性质表示注意力函数会根据 Query 的不同语义分配不同的注意力分数，反映了区分不同语义的能力。
使用不同的核函数时，线性注意力的非注入性质会导致各种语义混淆。
此外，尽管这些机制被认可具有强大的长程建模能力，但它们要想取得出色的效果仍然高度依赖有效的局部建模。详尽的经验验证无可争议地支持了作者的分析。
作者的研究还表明，在拥有上述两种性质的情况下，线性注意力可以在较低的计算复杂度下优于Softmax注意力。
unset
unset
参考
unset
unset
[0]. Bridging the Divide: Reconsidering Softmax and Linear Attention .
点击上方卡片，关注
「AI视界引擎」
公众号