导读

自2020年以来，Transformer模型迅速抢占了CV、NLP中诸多任务的主流模型，主流研究大多认为Transformer模型中的注意力机制能够以较大的感受野建模输入Token之间的全局联系。然而，注意力机制的计算复杂度与Token的映射维度成线性正比，与Token数量成指数级正比，这使得Transformer模型大多受限于参数量大的问题。本文提出了Neighborhood Attention Transformer，一种由Neighborhood Attention组成的高效、准确、可扩展的新型分层Transformer。其核心在于将Swin中的局部Attention操作进行朝着卷积的方向更进一步扩展。实验表明，NAT的性能优于Swin Transformer和ConvNeXt，NAT-Tiny模型在ImageNet上仅用4.3 GFLOPs和28M参数就能达到83.2%的top-1精度。

贡献

正如前面提到的，尽管注意力机制带来了Transformer模型的优越性能，但它们同时也带来了巨大的参数运算量。此外，注意力机制集中关注于全局联系，无法像卷积操作一样捕获输入中局部信息、平移不变性等归纳偏置。

为解决这一问题，局部注意力机制已经被提出来。例如图1中给出的Swin模型，它使用一个窗口对局部像素进行切割，并在一个窗口内计算其注意力机制，接着，将窗口不断进行平移，从而得到了多个不同位置的局部归纳偏置。HaloNet同样采用了局部注意力模块，并发现注意力机制与卷积层的结合不仅可以带来最佳性能，还能在计算消耗和平移不变性等归纳偏置中取得较好的平衡。

本文从卷积层的角度将Swin中的局部Attention操作进行朝着卷积的方向更进一步扩展，提出了NAT：Neighborhood Attention Transformer。其主要贡献可以总结如下：

提出了一个neighborhood attention操作。顾名思义，其以query对应的位置为中心来设定局部窗口，从而提取key和value进行计算。这是一种概念上更加简单、灵活和自然的的注意力机制。
基于提出的neighborhood attention构建了一个完整的vision transformer模型。模型延续始终分层的金字塔结构，每一层跟着一个下采样操作来缩减一半的尺寸。不同于现有的Swin等模型中采用的等效于非重叠卷积的操作，这里使用小尺寸且带重叠的卷积操作来进行特征嵌入和下采样。同时对于Stem阶段，同样采用了重叠卷积的操作。这些设定也带来了良好的效果，虽然会引入更多的计算成本和参数量，但是通过合理配置模型的结构，仍然可以获得足够优异的表现。
分类、检测、分割上都获得了良好的表现，超过了现有的众多方法，包括Swin和ConvNeXt等。对于28M参数这个近似于ResNet50的量级，在ImageNet1K上都可以达到83.2%的Top1的分类表现效果优异。

方法

上图展示了自注意力操作和NA之间的不同之处：

自注意力操作允许每个token去和全局所有token之间进行交互；
NA中会强制要求k，v的选择来自于以q为中心的邻域，即这里同时强调了两点，一个是局部区域，一个是以q为中心。

同时也体现出了和Swin这类的直接将原始的全局注意力机制约束到局部窗口内的局部注意力形式上的差异：

二者相同之处在于对于每个q都是和一个局部的窗口内的k和v来计算attention计算，这也因此使得二者之间的理论上的计算量是一致的，都是有着线性的计算复杂度和内存使用。
但是，本文的NA却有着一点额外的优势，即直接将注意力操作范围限制到了每个像素的邻域，而这些领域的计算天然地实现了一种类似于卷积那样的“滑动窗口”式的带重叠区域的交互，而不需要再像Swin那样需要额外借助于偏移操作来引入窗口的交互。同时更实际的一点是，也不再限制输入必须要能被窗口大小整除。

我们可以通过一个更直观的动画理解NAT，如上图所示。从原本局部Attention操作的窗口内部计算标准Attention操作改为了特定点上的q和其局部范围内的kv计算Attention，即原本的固定等分的窗口变成了以q为中心的滑动窗口。通过这样的形式，对Attention操作引入更多的局部偏置属性，从而针对视觉任务获得更好的表现。