【CVPR2023】Neighborhood Attention Transformer

最新推荐文章于 2025-04-25 21:06:48 发布

shanks66

最新推荐文章于 2025-04-25 21:06:48 发布

阅读量932

点赞数 41

分类专栏： CVPR论文 Attention 文章标签： transformer 深度学习人工智能

本文链接：https://blog.csdn.net/weixin_59422604/article/details/147460603

版权

CVPR论文同时被 2 个专栏收录

30 篇文章

订阅专栏

Attention

6 篇文章

订阅专栏

邻域注意力变换器

代码链接：
https://github.com/SHI-Labs/Neighborhood-Attention-Transformer
在这里插入图片描述
图 1：自注意力、（移位）窗口自注意力和我们的邻域注意力的注意力范围示意图。自注意力允许每个令牌关注所有内容。窗口自注意力将自注意力划分为不重叠的子窗口，随后是移位窗口自注意力，它允许窗口外的交互，这对于感受野扩展是必要的。邻域注意力将注意力定位在每个令牌周围的邻域内，引入局部归纳偏差，保持平移等变性，并允许感受野增长而无需额外操作

摘要

我们提出了邻域注意力（NA），这是首个用于视觉领域的高效且可扩展的滑动窗口注意力机制。NA是一种逐像素操作，它将自注意力（SA）定位到每个像素的最近邻像素上，因此与SA的二次复杂度相比，它具有线性的时间和空间复杂度。滑动窗口模式使NA的感受野能够在无需额外像素移位的情况下增长，并且与Swin变换器的窗口自注意力（WSA）不同，它保持了平移等变性。我们开发了NATTEN（邻域注意力扩展），这是一个带有高效C++和CUDA内核的Python包，它能让NA的运行速度比Swin的WSA快40%，同时内存使用量减少25%。我们进一步提出了邻域注意力变换器（NAT），这是一种基于NA的新型分层变换器设计，它可以提升图像分类和下游视觉任务的性能。NAT的实验结果极具竞争力；NAT-Tiny在ImageNet上的top-1准确率达到83.2%，在MSCOCO上的平均精度均值（mAP）为51.4%，在ADE20K上的平均交并比（mIoU）为48.4%，与类似规模的Swin模型相比，ImageNet准确率提高了1.9%，COCO的mAP提高了1.0%，ADE20K的mIoU提高了2.6%。为了支持更多基于滑动窗口注意力的研究，我们开源了项目并发布了模型检查点。

1. 引言

多年来，卷积神经网络（CNNs）一直是不同应用场景中计算机视觉模型事实上的标准架构。AlexNet在ImageNet上展示了其有效性，随后许多其他架构如VGG、ResNet和EfficientNet等也纷纷效仿。另一方面，变换器最初是作为用于自然语言处理（NLP）的基于注意力的模型被提出的，旨在利用语言的顺序结构。它们是BERT和GPT构建的基础，并且仍然是NLP领域的前沿架构。

2020年末，视觉变换器（ViT）被提出作为一种仅使用变换器编码器对图像补丁的嵌入空间进行操作的图像分类器，主要用于大规模训练。随后出现了许多其他方法，试图提高数据效率，最终使这类类似变换器的模型在ImageNet-1K分类（不基于大规模数据集如JFT-300M进行预训练）中达到了最先进的水平。

这些高性能的类似变换器的方法都基于自注意力（SA），这是原始变换器的基本构建块。SA在嵌入维度上具有线性复杂度（不包括线性投影），但在令牌数量上具有二次复杂度。在视觉领域中，令牌数量通常与图像分辨率呈线性相关。因此，对于严格使用SA的模型（如ViT），更高的图像分辨率会导致复杂度和内存使用量呈二次增长。这种二次复杂度使得这类模型难以轻松应用于下游视觉任务，如目标检测和分割，因为在这些任务中图像分辨率通常比分类任务大得多。另一个问题是，卷积受益于局部性和二维空间结构等归纳偏差，而点积自注意力从定义上讲是一种全局的一维操作。这意味着其中一些归纳偏差必须通过大量数据或先进的训练技术和增强方法来学习。

因此，人们提出了局部注意力模块来缓解这些问题。独立自注意力（SASA）是最早将基于局部窗口的注意力应用于视觉的方法之一，其中每个像素关注其周围的一个窗口。它的显式滑动窗口模式与相同卷积的模式相同，周围有零填充，并进行简单的二维光栅扫描，因此保持了平移等变性。SASA旨在取代ResNet中的卷积，并且被证明比基线模型有显著改进。然而，作者指出，由于缺乏类似于卷积的高效实现，SASA在速度方面存在限制。另一方面，Swin是最早的基于局部自注意力的分层视觉变换器之一。它的设计和移位窗口自注意力使其能够轻松应用于下游任务，因为它们使计算可行，同时通过注入额外的偏差提升了性能。然而，Swin的局部注意力首先将自注意力应用于不重叠的窗口，然后移动窗口，这样做的动机是像SASA这样的滑动窗口方法存在吞吐量瓶颈。HaloNet使用了一种光晕机制，一次对像素块进行局部自注意力计算，而不是逐像素计算。他们这样做的一个关键动机也是注意到缺乏高效的滑动窗口注意力机制。
在这里插入图片描述
图 2：单个像素的邻域注意力（NA）与自注意力（SA）的查询 - 键 - 值结构示意图。SA 允许每个像素关注其他所有像素，而 NA 将每个像素的注意力定位在其自身周围的邻域内。因此，每个像素的注意力范围通常彼此不同。

在这项工作中，我们重新审视了显式滑动窗口注意力机制，并提出了邻域注意力（NA）。NA将SA定位到每个像素的最近邻，这不一定是围绕像素的固定窗口。这种定义上的改变使所有像素都能保持相同的注意力范围，否则在零填充的替代方法（如SASA）中，角落像素的注意力范围会减小。随着邻域大小的增加，NA也会接近SA，在最大邻域时与SA等效。此外，与分块和窗口自注意力不同，NA具有保持平移等变性的额外优势。我们开发了NATTEN，这是一个带有高效C++和CUDA内核的Python包，在实际应用中，它使NA的运行速度比Swin的WSA更快，同时使用的内存更少。我们构建了邻域注意力变换器（NAT），它在各种视觉任务中都取得了有竞争力的结果。

综上所述，我们的主要贡献如下：

提出邻域注意力（NA）：一种简单灵活的显式滑动窗口注意力机制，它将每个像素的注意力范围定位到其最近邻域，随着范围的增加接近自注意力，并保持平移等变性。我们在复杂度和内存使用方面将NA与自注意力、窗口自注意力和卷积进行了比较。
为NA开发高效的C++和CUDA内核，包括分块NA算法，该算法使NA的运行速度比Swin的WSA快40%，同时内存使用量减少25%。我们将它们发布在一个新的用于显式滑动窗口注意力机制的Python包NATTEN中，提供易于使用的、支持自动求导的模块，可以插入到任何现有的PyTorch管道中。
引入邻域注意力变换器（NAT），这是一种基于NA的新型高效、准确且可扩展的分层变换器。我们展示了它在分类和下游任务中的有效性。例如，NAT-Tiny在ImageNet上仅用43亿次浮点运算（GFLOPs）和2800万个参数就达到了83.2%的top-1准确率，在MS-COCO上的边界框mAP为51.4%，在ADE20K上的mIoU为48.4%，显著优于Swin变换器和ConvNeXt。

在这里插入图片描述
图 3：ImageNet - 1K 分类性能与计算量的关系，气泡大小表示参数数量。NAT 在分类任务中使用更少的 FLOPs 和相似数量的参数，优于 Swin 变换器和 ConvNeXt。

2. 相关工作

在本节中，我们简要回顾原始的自注意力（SA）、一些著名的视觉变换器和类似变换器的架构、一些著名的基于局部注意力的视觉变换器，以及一个最近的CNN，它为基于注意力的模型提供了最新的基线。

2.1 自注意力

瓦斯瓦尼（Vaswani）等人将缩放点积注意力定义为对查询和一组键值对的操作。计算查询和键的点积并进行缩放，应用Softmax对输出进行归一化以得到注意力权重，然后将其应用于值。它可以表示为：
$V)=softmax\left(\frac{Q K^{T}}{\sqrt{d}}\right) V \tag{1}$
其中 $d$ 是嵌入维度。自注意力对相同输入的线性投影同时作为查询和键值对应用点积注意力。在变换器中，通常应用多头注意力和自注意力的变体。多头注意力在不同的嵌入上多次应用点积注意力，从而形成注意力头。给定输入 $\in \mathbb{R}^{n ×d}$ ，其中 $n$ 是令牌数量， $d$ 是嵌入维度，此操作的复杂度为 $O(n^{2} d)$ ，注意力权重的空间复杂度为 $O(n^{2})$ 。

2.2 视觉变换器

多索维茨基（Dosovitskiy）等人提出了一种基于变换器的图像分类器，仅由一个变换器编码器和一个图像标记器组成，名为视觉变换器（ViT）。以前的工作，如DETR，探索了CNN - 变换器的混合模型用于目标检测。而ViT提出了一种仅依赖于单个非重叠卷积层（补丁和嵌入）的模型。ViT主要在私有JFT - 300M数据集上进行预训练，并且在许多基准测试中表现优于最先进的CNN。然而，也有人指出，当ViT在中等规模的数据集（如ImageNet - 1K和ImageNet - 21K）上进行预训练时，它不再能取得有竞争力的结果。这归因于缺乏CNN固有的归纳偏差，作者认为在大规模训练中这种偏差被超越了。虽然这有效地证明了ViT在中等规模训练中的劣势，但它提供了经验证据，表明基于变换器的模型在更大规模上优于CNN。ViT为许多后续的视觉变换器和一般的基于注意力的模型铺平了道路，这些模型将其应用于中等规模的学习，甚至在更小的数据集上进行小规模学习。图夫龙（Touvron）等人通过探索数据效率扩展了对视觉变换器的研究。他们的数据高效图像变换器（DeiT）模型通过最小的架构变化和使用先进的增强和训练技术推动了ViT的发展。他们的努力凸显了基于变换器的图像分类器在中等规模数据领域的真正潜力，并启发了更多人采用他们的训练技术。

2.3 局部注意力

独立自注意力（SASA）是最早的滑动窗口自注意力模式之一，旨在取代现有CNN中的卷积。它的操作类似于带零填充的卷积，通过对特征图进行步长采样来提取键值对。作者报告说准确率有显著提高，但观察到尽管理论成本较低，但实现过程中存在高延迟问题。这种注意力模式也在语言处理的Longformer（滑动窗口注意力）中被采用，后来在视觉Longformer（ViL）中也被采用。虽然Longformer和ViL的实现与SASA不同，但由于计算开销，它们仍然无法扩展到更大的窗口和模型。此外，填充导致的角落情况下感受野减小的问题也没有得到解决。刘等人引入了窗口和移位窗口（Swin）注意力，作为基于非滑动窗口的自注意力机制，对特征图进行分区并分别对每个分区应用自注意力。此操作与SASA具有相似的理论复杂度，但它可以通过批矩阵乘法轻松并行化。移位变体在常规窗口注意力之后进行，顾名思义，它移动分区以允许窗口外的交互，这对于感受野的增长是必要的。他们提出的模型Swin变换器是最早的分层视觉变换器之一。它生成类似金字塔的特征图，在增加深度的同时减少空间维度。多年来，这种结构在CNN中被广泛使用，这就是为什么Swin可以轻松与其他网络集成，应用于下游任务，如检测和分割。Swin在ImageNet - 1K分类中优于使用卷积教师模型的DeiT。此外，Swin变换器成为MS - COCO目标检测和ADE20K语义分割中的最先进方法。瓦斯瓦尼等人提出了HaloNet，旨在通过用一种新的分块注意力模式取代SASA来避免其速度问题。他们指出，虽然这种改变放宽了平移等变性，但在速度和内存方面可以提供合理的权衡。HaloNet的注意力机制由三个阶段组成：分块、光晕和注意力。输入特征图被分块为不重叠的子集，作为查询。随后，提取“光晕化”的相邻块，作为键和值。然后对提取的查询和键值对应用注意力。HaloNet被证明在降低成本（与SASA相比）和提高性能方面都很有效，特别是在与网络中的卷积层结合使用时。许多工作在采用WSA方面效仿Swin，例如RegionViT，其中在每个局部自注意力层中插入一个区域令牌，以引入全局上下文。这项工作和HaloNet都表明，研究社区对滑动窗口注意力模式失去兴趣的部分原因是它们被认为效率低下。我们旨在通过引入NATTEN来改变这一现状。

2.4 新的卷积基线

刘等人提出了一种受Swin等模型影响的新CNN架构，称为ConvNeXt。这些模型不是基于注意力的，并且在不同的视觉任务中优于Swin。从那以后，这项工作成为了卷积模型和基于注意力的模型进行公平比较的新CNN基线。

我们提出邻域注意力，通过设计将感受野定位到每个查询周围的窗口，因此不需要像Swin使用的循环移位等额外技术。此外，邻域注意力保持平移等变性，而在HaloNet和Swin等方法中，为了提高效率牺牲了平移等变性。我们通过NATTEN Python包证明，邻域注意力的运行速度比Swin等方法更快，同时使用的内存更少。我们引入了一种具有这种注意力机制的分层类似变换器的模型，称为邻域注意力变换器，并展示了它在图像分类、目标检测和语义分割方面与Swin相比的性能。

3. 方法

在本节中，我们介绍邻域注意力，这是一种考虑视觉数据结构的自注意力定位方法。与自注意力相比，这不仅降低了计算成本，还引入了类似于卷积的局部归纳偏差。我们表明，在限制自注意力方面，NA是比之前提出的SASA更好的替代方案，同时理论成本相同。然后我们介绍我们的Python包NATTEN，它为CPU和GPU加速提供了NA的高效实现。我们讨论了该扩展的新颖之处，以及它如何在使用更少内存的同时超过Swin的WSA和SWSA的速度。最后，我们介绍我们的模型邻域注意力变换器（NAT），它使用这种新机制代替自注意力。此外，NAT采用了类似于Swin的多级分层设计，这意味着特征图在不同级别之间进行下采样，而不是一次性全部下采样。与Swin不同，NAT使用重叠卷积对特征图进行下采样，而不是非重叠（补丁）卷积，已有研究表明，重叠卷积通过引入有用的归纳偏差可以提高模型性能。

3.1 邻域注意力

Swin的WSA可以被认为是现有用于限制自注意力以降低二次注意力成本的最快方法之一。它简单地对输入进行分区，并分别对每个分区应用自注意力。WSA需要与移位变体SWSA配合使用，SWSA移动那些分区线以允许窗口外的交互，这对于扩展感受野至关重要。然而，最直接的在局部限制自注意力的方法是允许每个像素关注其相邻像素。这导致大多数像素周围有一个动态移位的窗口，从而扩展了感受野，因此不需要手动的移位变体。此外，与Swin不同且与卷积类似，这种动态形式的受限自注意力可以保持平移等变性（我们在附录C中分析了包括我们自己的方法在内的不同方法的平移等变性）。受此启发，我们引入邻域注意力（NA）。给定输入 $\in \mathbb{R}^{n ×d}$ ，这是一个矩阵，其行是d维令牌向量，以及 $X$ 的线性投影Q、K和V，以及相对位置偏差 $B (i, j)$ ，我们定义邻域大小为 $k$ 的第 $i$ 个输入的注意力权重 $A_{i}^{k}$ ，为第 $i$ 个输入的查询投影与其 $k$ 个最近邻的键投影的点积：
$A_{i}^{k}=\left[\begin{array}{c} Q_{i} K_{\rho_{1}(i)}^{T}+B_{\left(i, \rho_{1}(i)\right)} \\ Q_{i} K_{\rho_{2}(i)}^{T}+B_{\left(i, \rho_{2}(i)\right)} \\ \vdots \\ Q_{i} K_{\rho_{k}(i)}^{T}+B_{\left(i, \rho_{k}(i)\right)} \end{array}\right] \tag{2}$
其中 $\rho_{j}(i)$ 表示 $i$ 的第 $j$ 个最近邻。然后我们将相邻值 $V_{i}^{k}$ 定义为一个矩阵，其行是第 $i$ 个输入的 $k$ 个最近邻的值投影：
$V_{i}^{k}=\left[\begin{array}{llll}V_{\rho_{1}(i)}^{T} & V_{\rho_{2}(i)}^{T} &... & V_{\rho_{k}(i)}^{T}\end{array}\right]^{T} \tag{3}$
邻域大小为 $k$ 的第 $i$ 个令牌的邻域注意力定义为：
$NA_{k}(i)=softmax\left(\frac{A_{i}^{k}}{\sqrt{d}}\right) V_{i}^{k} \tag{4}$
其中 $\sqrt{d}$ 是缩放参数。此操作在特征图的每个像素上重复进行。此操作的图示见图2和图VIII。

从这个定义很容易看出，随着 $k$ 的增加， $A_{i}^{k}$ 接近自注意力权重， $V_{i}^{k}$ 接近 $V_{i}$ 本身，因此邻域注意力接近自注意力。这是NA与SASA的关键区别，在SASA中，每个像素关注其周围的一个窗口，并且在输入周围进行填充以处理边缘情况。正是由于这种差异，随着窗口大小的增加，NA接近自注意力，而在SASA中由于输入周围的零填充，情况并非如此。

3.2 分块NA和NATTEN

过去，以逐像素的方式限制自注意力并未得到充分研究，主要是因为它被视为一种成本高昂的操作，需要底层重新实现。这是由于自注意力本身可分解为矩阵乘法，这是一种在加速器上易于并行化的操作，并且在计算软件（如LAPACK、cuBLAS、CUTLASS等）中有众多针对不同用例定义的高效算法。此外，大多数深度学习平台（如PyTorch）是基于此类软件（以及cuDNN等附加软件包）构建的。这对研究人员非常有帮助，因为他们可以使用矩阵乘法或卷积等操作的抽象，而后端会根据硬件、软件和用例决定运行哪种算法。这些平台通常还能处理自动梯度计算，这使得设计和训练深度神经网络变得非常简单直接。

由于NA（以及其他逐像素注意力机制，如SASA）的逐像素结构，以及NA中邻域定义的新颖性，在这些平台上实现NA的唯一方法是堆叠大量低效操作来提取邻域，将其存储为中间张量，然后计算注意力。这会导致操作速度显著变慢，内存使用量呈指数级增长。

为应对这些挑战，我们开发了一组高效的CPU和CUDA内核，并将它们打包成一个Python包——邻域注意力扩展（NATTEN）。NATTEN支持半精度计算，支持1D和2D数据，并且与PyTorch自动求导兼容。这意味着用户可以简单地将NA作为一个PyTorch模块导入，并集成到现有管道中。我们还要补充的是，SASA也可以通过这个包轻松实现，且无需更改底层内核（只需对输入进行零填充即可），因为它是NA的一个特殊情况，但反之则不成立。

在这里插入图片描述
图 4：NAT 相对于 Swin 的层间相对速度。两个内核大小为7的2次方的 NA 层比一对相同内核大小的 WSA 和 SWSA 层快 40%。延迟在单个 A100 GPU 上测量。NA 的 PyTorch 实现在分辨率为448×448及更高时会出现内存不足的情况。

NATTEN还包含我们的分块NA算法，该算法通过将不重叠的查询块加载到共享内存中来计算邻域注意力权重，从而最大限度地减少全局内存读取。与朴素实现相比，分块NA算法的延迟最多可降低一个数量级（技术细节见附录A），它能让基于NA的模型比类似的Swin模型运行速度快40%（见图4）。NATTEN已在https://github.com/SHI-Labs/NATTEN上开源。

3.3 邻域注意力变换器

NAT使用两个连续的3×3卷积层，步长为2×2对输入进行嵌入，使得空间大小变为输入的四分之一。这类似于使用4×4的补丁和嵌入层，但它利用重叠卷积而非非重叠卷积来引入有用的归纳偏差。另一方面，使用重叠卷积会增加成本，并且两个卷积层会引入更多参数。不过，我们通过重新配置模型来解决这个问题，从而实现更好的权衡。
在这里插入图片描述

图 5：我们的模型 NAT 的分层设计概述。该模型以卷积下采样器开始，然后进入 4 个连续的层级，每个层级由多个 NAT 块组成，这些块是类似变换器的编码器层。每层由多头邻域注意力（NA）、多层感知器（MLP）、每个模块前的层归一化（LN）和跳跃连接组成。在不同层级之间，特征图的空间大小减半，而深度翻倍。这使得通过特征金字塔更容易转移到下游任务。

NAT由4个层级组成，除最后一层外，每一层后面都跟着一个下采样器。下采样器将空间大小减半，同时将通道数量翻倍。我们使用3×3卷积层，步长为2×2，而不是Swin使用的2×2非重叠卷积（补丁合并）。由于标记器将输入下采样了4倍，我们的模型生成的特征图大小分别为 $\frac{h}{4}×\frac{w}{4}$ 、 $\frac{h}{8}×\frac{w}{8}$ 、 $\frac{h}{16}×\frac{w}{16}$ 和 $\frac{h}{32}×\frac{w}{32}$ 。这一改变的灵感来自以前成功的CNN结构，并且被其他分层注意力方法（如PVT、ViL和Swin变换器）所采用。此外，为了在更大的模型变体中保持稳定性，我们使用了LayerScale。图5展示了整体网络架构。表1总结了不同的NAT变体。
在这里插入图片描述

3.4 复杂度分析

在本小节中，我们进行复杂度和内存使用分析，在表2中比较了SA、WSA、NA和卷积。为简单起见，我们不考虑注意力头。给定形状为 $h \times w \times d$ 的输入特征图，其中 $d$ 是通道数， $h$ 和 $w$ 分别是特征图的高度和宽度， $Q K V$ 线性投影的计算量为3hwd²次浮点运算，这对三种注意力模式都是相同的。
在这里插入图片描述

SA具有二次复杂度，因为计算注意力权重和输出的计算量均为 $h^{2}w^{2}d$ 次浮点运算，且注意力权重的形状为 $h w \times h w$ 。Swin的WSA将查询、键和值划分为 $\frac{h}{k}×\frac{w}{k}$ 个形状为 $k \times k$ 的窗口，然后在每个窗口内应用自注意力，计算量为hwdk²次浮点运算。考虑到WSA的注意力权重形状为 $\frac{h}{k}×\frac{w}{k}×k^{2}×k^{2}$ ，其内存消耗为hwdk²。

在NA中， $A_{i}^{k}$ 的大小为 $h×w×k^{2}$ ，计算它的成本为hwdk²。 $V_{i}^{k}$ 的形状为 $h×w×k^{2}×d$ ，因此将注意力权重应用于它的成本为hwdk²。至于卷积，计算成本为 $hwd^{2}k^{2}$ ，内存使用量仅为 $d^{2}k^{2}$ 。表2中的总结表明，Swin的WSA和NA在理论上具有相同的计算成本和内存使用量。

4. 实验

我们通过在不同的视觉任务（如图像分类、目标检测和语义分割）中进行实验，展示NAT的适用性和有效性。我们还对不同的注意力模式以及NAT的设计进行了消融实验。附录B中提供了包括显著性分析在内的更多实验。

4.1 分类

我们在ImageNet-1K上训练了NAT的不同变体，以便与其他基于变换器和卷积的图像分类器进行比较。这个数据集仍然是少数用于中等规模图像分类的基准之一，包含大约128万张训练图像、5万张验证图像和10万张测试图像，分为1000个类别。我们使用常用的timm（Apache许可证v2）来训练NAT，并采用与我们对比的方法中使用的传统增强技术（CutMix、Mixup、RandAugment和随机擦除）和训练技术。我们遵循Swin的训练配置（学习率、迭代式余弦退火调度和其他超参数）。按照惯例，我们训练300个epoch，其中20个epoch用于热身学习率，其余的根据调度器进行衰减，并额外进行10个epoch的冷却训练。结果如表3所示，并在图3中进行了可视化。

我们观察到，NAT-Mini在参数更少、吞吐量更高和内存使用更低的情况下，比Swin-Tiny的准确率高出0.5%。对于其他三个变体，我们发现它们在参数数量和计算量相似的情况下，始终优于Swin和ConvNeXt的对应模型。虽然由于架构差异，我们的Small变体比其Swin对应模型略慢，但我们的Base变体比Swin-Base更快。

表 3：ImageNet - 1K 分类性能。所有模型在 224×224 分辨率下运行，无额外数据或预训练。峰值内存使用量和吞吐量在单个 NVIDIA A100 GPU 上以批量大小 256 进行测量。

模型	参数数量	FLOPs	吞吐量（图像 / 秒）	内存（GB）	Top - 1 准确率（%）
NAT - M	20M	2.7G	2135	2.4	81.8
Swin - T	28M	4.5G	1730	4.8	81.3
ConvNeXt - T	28M	4.5G	2491	3.4	82.1
NAT - T	28M	4.3G	1541	2.5	83.2
Swin - S	50M	8.7G	1059	5.0	83.0
ConvNeXt - S	50M	8.7G	1549	3.5	83.1
NAT - S	51M	7.8G	1051	3.7	83.7
Swin - B	88M	15.4G	776	6.7	83.5
ConvNeXt - B	89M	15.4G	1107	4.8	83.8
NAT - B	90M	13.7G	783	5.0	84.3

4.2 目标检测和实例分割

我们在MS-COCO上使用预训练在ImageNet上的NAT骨干网络训练Mask R-CNN和Cascade Mask R-CNN。我们遵循Swin的训练设置，使用mmdetection（Apache许可证v2），并采用相同的加速3×学习率调度进行训练。结果如表4所示。NAT-Mini在使用Mask R-CNN时优于Swin-Tiny，在使用Cascade Mask R-CNN时略逊于它，但计算量显著更少。NAT-Tiny在使用Mask和Cascade Mask R-CNN时均优于其Swin和ConvNeXt对应模型，不过与Swin对应模型相比，吞吐量略低。NAT-Small与Swin对应模型相比达到了有竞争力的性能，同时速度更快。NAT-Base甚至可以超越其Swin对应模型，并且吞吐量更高。
表 4：COCO 目标检测和实例分割性能。FLOPS 是针对输入分辨率 (1280, 800) 的。吞吐量在单个 NVIDIA A100 GPU 上以相同分辨率测量。

骨干网络	参数数量	FLOPs	吞吐量（FPS）	AP_b	AP_b50	AP_b75	AP_m	AP_m50	AP_m75
Mask R - CNN - 3x 调度
NAT - M	40M	225G	54.1	46.5	68.1	51.3	41.7	65.2	44.7
Swin - T	48M	267G	45.1	46.0	68.1	50.3	41.6	65.1	44.9
ConvNeXt - T	48M	262G	52.0	46.2	67.0	50.8	41.7	65.0	44.9
NAT - T	48M	258G	44.5	47.7	69.0	52.6	42.6	66.1	45.9
Swin - S	69M	359G	31.7	48.5	70.2	53.5	43.3	67.3	46.6
NAT - S	70M	330G	34.8	48.4	69.8	53.2	43.2	66.9	46.5
Cascade Mask R - CNN - 3x 调度
NAT - M	77M	704G	27.8	50.3	68.9	54.9	43.6	66.4	47.2
Swin - T	86M	745G	25.1	50.4	69.2	54.7	43.7	66.6	47.3
ConvNeXt - T	86M	741G	27.3	50.4	69.1	54.8	43.7	66.5	47.3
NAT - T	85M	737G	24.9	51.4	70.0	55.9	44.5	67.6	47.9
Swin - S	107M	838G	20.3	51.9	70.7	56.3	45.0	68.2	48.8
ConvNeXt - S	108M	827G	23.0	51.9	70.8	56.5	45.0	68.4	49.1
NAT - S	108M	809G	21.7	52.0	70.4	56.3	44.9	68.1	48.6
Swin - B	145M	982G	17.3	51.9	70.5	56.4	45.0	68.1	48.9
ConvNeXt - B	146M	964G	19.5	52.7	71.3	57.2	45.6	68.9	49.5
NAT - B	147M	931G	18.6	52.5	71.1	57.1	45.2	68.6	49.0

4.3 语义分割

为了展示NAT在语义分割方面的性能，我们在ADE20K上使用NAT骨干网络训练UPerNet。我们遵循Swin在ADE20K上的训练配置，并使用mmsegmentation（Apache许可证v2）。此外，按照标准做法，在训练时将输入图像随机调整大小并裁剪为512×512。结果如表5所示。值得注意的是，NAT-Mini优于Swin-Tiny，并且与ConvNeXt-Tiny非常接近。NAT-Tiny显著优于ConvNeXt-Tiny，同时效率也略高。NAT-Small在单尺度性能上优于Swin-Small，在多尺度性能上与之匹配。NAT-Base与Swin-Base性能相当，但略逊于ConvNeXt-Base。需要注意的是，NAT-Small和NAT-Base与它们的Swin和ConvNeXt对应模型相比，计算量更少，而性能处于同一水平。还值得一提的是，由于此任务中输入图像分辨率（512×512）不会使特征图能被Swin的窗口大小7×7整除，Swin在每次注意力操作之前都需要对输入特征图进行零填充来解决这个问题，这使得Swin的计算量比原本的差异更大。而NAT支持任意大小的特征图。
表 5：ADE20K 语义分割性能。FLOPS 是针对输入分辨率 (2048, 512) 的。吞吐量在单个 NVIDIA A100 GPU 上以相同分辨率测量。

骨干网络	参数数量	FLOPs	吞吐量（FPS）	mIoU
				单尺度	多尺度
NAT - M	50M	900G	24.5	45.1	46.4
Swin - T	60M	946G	21.3	44.5	45.8
ConvNeXt - T	60M	939G	23.3	46.0	46.7
NAT - T	58M	934G	21.4	47.1	48.4
Swin - S	81M	1040G	17.0	47.6	49.5
ConvNeXt - S	82M	1027G	19.1	48.7	49.6
NAT - S	82M	1010G	17.9	48.0	49.5
Swin - B	121M	1188G	14.6	48.1	49.7
ConvNeXt - B	122M	1170G	16.4	49.1	49.9
NAT - B	123M	1137G	15.6	48.5	49.7

4.4 消融实验

我们将Swin的注意力模式（WSA+SASA）与滑动窗口模式进行比较，即SASA（使用我们的NATTEN包实现，因此与NA具有大致相同的吞吐量和相同的内存使用量）和我们的NA。我们简单地替换Swin-Tiny中的注意力块，并在ImageNet-1K分类、MSCOCO目标检测和实例分割以及ADE20K分割任务上运行模型。结果如表6所示。
表 6：不同注意力机制的性能比较。所有展示的模型在架构上与 Swin - T 相同，仅注意力机制被替换。SASA 使用我们的 NATTEN 实现，因此与 NA 具有相同的速度和内存效率。

注意力	ImageNet	MSCOCO	ADE20K	参数数量	FLOPs	吞吐量（图像 / 秒）	内存（GB）
	Top - 1 准确率（%）	AP_B	AP_m	mIoU
SWSA	81.3%	46.0	41.6	45.8	28.28M	4.51G	1730
SASA	81.6%	46.0	41.4	46.4	28.27M	4.51G	2021
NA	81.8%	46.2	41.5	46.4	28.28M	4.51G	2021

另外，我们以Swin-Tiny为基线进行消融实验，研究NAT设计（卷积下采样和更深更窄的架构）的影响。我们逐步将模型转换为NAT-Tiny，并将结果列于表7中。
表 7：以 Swin - T 为基线对 NAT 进行的消融研究。通过重叠卷积（第 2 行）和我们的 NAT 配置（第 3 行），我们显著提高了 Swin 的分类准确率（第 1 行）。将 SWSA 替换为 NA（第 5 行）导致准确率提高近 0.5%，而将其替换为 SASA（第 4 行）导致准确率略有下降。SASA 使用我们的 NATTEN 实现，因此与 NA 具有相同的速度和内存效率。

注意力	下采样器	层数	头数	MLP 比例	Top - 1 准确率（%）	参数数量	FLOPs（G）	吞吐量（图像 / 秒）	内存（GB）
SWSA	Patch	2, 2, 6, 2	3	4	81.29	28.3M	4.5	1730	4.8
SWSA	Conv	2, 2, 6, 2	3	4	81.78	30.3M	4.9	1692	4.8
SWSA	Conv	3, 4, 18, 5	2	3	82.72	27.9M	4.3	1320	3.0
SASA	Conv	3, 4, 18, 5	2	3	82.54	27.9M	4.3	1541	2.5
NA	Conv	3, 4, 18, 5	2	3	83.20	27.9M	4.3	1541	2.5
我们首先用NAT中使用的重叠卷积设计替换补丁嵌入和补丁合并，这使得准确率提高了近0.5%。在采取第二步，通过使模型更深更窄来减小模型大小和计算量后，我们发现模型的准确率比第一步提高了约0.9%。然后，我们尝试将Swin中的WSA和SWSA注意力模式替换为SASA，发现准确率略有下降。然而，将WSA和SWSA替换为我们的NA后，准确率进一步提高了0.5%。

我们还在表8中进行了内核大小实验，尝试了从3×3到9×9的内核大小，以分析其对模型性能的影响。
表 8：不同内核大小下 NAT - Tiny 的性能

内核大小	ImageNet		MSCOCO		ADE20K
	Top - 1 准确率（%）	吞吐量	AP_b	AP_m	吞吐量	mIoU
3×3	81.4	2015 图像 / 秒	46.1	41.4	46.8fps	46.0
5×5	81.6	1810 图像 / 秒	46.8	42.0	45.5fps	46.3
7×7	83.2	1537 图像 / 秒	47.7	42.6	44.5fps	48.4
9×9	83.1	1253	48.5	43.3	39.4fps	48.1

5. 结论

在本文中，我们提出了邻域注意力（NA），这是首个用于视觉领域的高效且可扩展的滑动窗口注意力机制。NA是一种逐像素操作，它将每个像素的自注意力定位到其最近邻域，因此具有线性复杂度。它还引入了局部归纳偏差，并保持了平移等变性，这与分块（HaloNet）和窗口自注意力（Swin）不同。与SASA不同，随着窗口大小的增加，NA接近自注意力，并且在角落情况下不会限制注意力范围。

我们通过开发NATTEN挑战了“显式滑动窗口注意力模式效率低下或不可并行化”这一普遍观点。通过使用NATTEN，基于NA的模型甚至可以比现有的替代模型运行得更快，尽管后者主要基于在底层计算包之上构建的高度优化的深度学习库运行。我们还提出了邻域注意力变换器（NAT），并展示了这类模型的强大之处：NAT在图像分类任务中优于Swin变换器和ConvNeXt，在下游视觉任务中也优于或与它们具有竞争力。我们开源了整个项目，以鼓励在这一方向上进行更多研究。