DiNA：扩张邻域注意力 Transformer

摘要

Transformer 正迅速成为跨模态、跨领域和跨任务中应用最广泛的深度学习架构之一。在计算机视觉领域，除了持续发展的纯 transformer 架构，分层 transformer 也因其优越的性能和在现有框架中易于集成而受到广泛关注。这类模型通常采用局部化的注意力机制，如滑动窗口的 Neighborhood Attention（NA）或 Swin Transformer 的 Shifted Window Self Attention。尽管这些方法在降低自注意力的二次复杂度方面效果显著，但局部注意力削弱了自注意力机制中两个最重要的特性：长距离依赖建模能力和全局感受野。本文中，我们提出了 Dilated Neighborhood Attention（DiNA），这是对 NA 的一种自然、灵活且高效的扩展，能够在不增加额外计算成本的前提下捕捉更多的全局上下文并指数级扩展感受野。NA 的局部注意力和 DiNA 的稀疏全局注意力互为补充，因此我们进一步提出 Dilated Neighborhood Attention Transformer（DiNAT），这是一种基于二者构建的新型分层视觉 transformer。DiNAT 的多个变体在 NAT、Swin 和 ConvNeXt 等强大基线模型上均实现了显著提升。我们的大型模型在 COCO 目标检测中比 Swin 同类模型快且高出 1.6% 的 box AP，在 COCO 实例分割中高出 1.4% 的 mask AP，在 ADE20K 语义分割中高出 1.4% 的 mIoU。在与新框架结合使用时，我们的大型变体成为 COCO（58.5 PQ）和 ADE20K（49.4 PQ）上最新的 state-of-the-art 全景分割模型，并成为 Cityscapes（45.1 AP）和 ADE20K（35.4 AP）上的实例分割 state-of-the-art 模型（未使用额外数据）。它在 ADE20K 上还达到了语义分割任务 state-of-the-art 专用模型的表现（58.1 mIoU），在 Cityscapes 上排名第二（84.5 mIoU）（未使用额外数据）。为了支持并鼓励计算机视觉及其他方向上的相关研究，我们开源了该项目：https://github.com/SHI-Labs/Neighborhood-Attention-Transformer。

1. 引言

Transformer [42] 对人工智能研究做出了重要贡献，最初应用于自然语言理解 [12,34]，随后由于其基于自注意力的通用架构，也被扩展到语音 [14] 和视觉 [13,32] 等其他模态。这一成功激发了在视觉领域中基于注意力机制模型的研究热潮，涵盖从骨干网络 [35,41] 到更具体的应用，如图像生成与密度建模 [6,32]、目标检测 [3]、图像分割 [20,43] 等多个方向。

Vision Transformer（ViT）[13] 是最早将 transformer 作为卷积神经网络（CNN）[19,22,23] 替代方案的重要尝试之一，而 CNN 一直是视觉领域的事实标准。ViT 将图像视为 patch 序列，并使用一个纯 transformer 编码器对图像进行编码和分类。它在大规模图像分类任务中展现出可与 CNN 相媲美的性能，进一步推动了 transformer 架构作为 CNN 替代者的视觉研究热潮 [38,39]。

Vision transformer 和 CNN 的不同不仅体现在架构和构建模块上，还体现在它们对数据的处理方式上。CNN 通常会在模型中逐步下采样输入数据，并构建分层特征图。这种分层设计对于视觉任务至关重要，因为物体的尺度变化较大，而高分辨率特征图对于分割等密集任务尤为关键。另一方面，transformer 以模型中维度固定著称，因此，纯 ViT 为了缓解自注意力的二次计算成本，从一开始就对输入进行了激进的下采样，这反过来限制了纯 ViT 作为密集视觉任务骨干网络的应用。

尽管将纯 ViT 应用于密集视觉任务的研究仍在持续 [17,24]，但分层视觉 transformer 的研究迅速占据主导地位 [29,44]，并持续增长 [15,28]。这类分层 transformer 模型的一个关键优势是其可以轻松集成到现有的分层视觉框架中。受 CNN 架构启发，分层视觉 transformer 通常由多个（一般为 4 个）transformer 编码器层组成，各层之间有下采样模块，并且初始下采样不再那么激进（即 1/4 而非 1/16）。然而，在分层 transformer 中，若早期层使用无约束的自注意力，其计算复杂度和内存使用量仍会随输入分辨率以二次级数增长，从而在高分辨率图像上变得难以处理。因此，分层 transformer 通常采用某种局部注意力机制。

Swin Transformer [29] 是最早的分层视觉 transformer 之一，它采用了 Window Self Attention（WSA）模块，随后接一个像素偏移的 Window Self Attention（SWSA），这两种方法都将 self attention 局部化于不重叠的子窗口中。这样做降低了 self attention 的计算成本，使其在分辨率上的时间和空间复杂度变为线性。SWSA 与 WSA 相同，只是在其之前引入了特征图像素的偏移，之后再进行反向偏移。这一步对模型性能至关重要，因为它允许跨窗口的交互，从而扩大感受野。Swin 的主要优势之一在于其高效性，因为像素偏移与窗口划分都是相对廉价且易于并行化的操作。此外，它几乎无需对 self attention 模块做出改变，使得实现更为简便。Swin 成为了多个视觉任务中的 state-of-the-art 方法，之后 Swin-V2 [28] 被提出以适应大规模预训练。

Neighborhood Attention Transformer（NAT）[15] 随后被提出，使用了一种基于滑动窗口的简单注意力机制，称为 Neighborhood Attention（NA）。与模仿卷积方式进行注意力计算的 Stand Alone Self Attention（SASA）[35] 不同，NA 将 self attention 局部化于每个 token 周围的最近邻区域，从定义上来说更接近原始 self attention，同时保持固定的注意力范围。此类逐像素的注意力操作曾被认为效率低且难以并行化 [29,35,41]，直到 Neighborhood Attention Extension [15] 的发布。借助这一扩展，NA 在实际运行中甚至可以比 Swin 的 SWSA 更快。NAT 在图像分类任务上显著优于 Swin，并在下游任务中也表现出竞争力，同时在架构略有差异的情况下，其扩展能力甚至超过了 Swin，实现了更快的运行速度。

尽管已有诸多关于采用局部注意力机制的分层视觉 transformer 的研究，但 self attention 中最重要的一些特性——包括全局感受野以及建模长距离依赖的能力——在局部化之后仍被削弱。

这引出了一个简单的问题：如何在分层视觉 transformer 中保持局部注意力带来的可计算性优势，同时避免其缺点？换句话说，理想的情况是，在保持线性复杂度的同时，保留 self attention 的全局感受野以及建模长距离依赖的能力。本文旨在回答这个问题，并通过扩展一种简单的局部注意力机制——Neighborhood Attention（NA）——提出 Dilated Neighborhood Attention（DiNA）：一种灵活而强大的稀疏全局注意力。将 NA 中的邻域扩张为更大的稀疏区域具有多个优势：1）能够捕捉更多的全局上下文；2）可以使感受野呈指数级增长，而非线性增长 [50]；3）不会带来额外的计算开销。为了验证 DiNA 的有效性，我们提出了 Dilated Neighborhood Attention Transformer（DiNAT），它不仅在下游任务表现上改进了现有的 NAT 模型，还在多个下游任务中以明显优势超越了强大的现代 CNN 基线模型，如 ConvNeXt [30]。

我们的主要贡献总结如下：

• 提出 DiNA，这是一种简单而强大的稀疏全局注意力模式，能够在不增加计算负担的前提下，使感受野呈指数级扩张，并捕捉更长范围的上下文。DiNA 同时保留了 NA 所引入的邻域对称性特征。它还可以适应更大分辨率，而无需扩大窗口大小。
• 分析基于卷积、局部注意力以及 DiNA 模型的理论感受野大小。
• 提出 DiNAT，一种结合了扩张与非扩张 NA 的新型分层视觉 transformer。DiNAT 在模型中引入逐层渐进的扩张变化，更优地扩展感受野，并促进由细到粗的特征学习。
• 在图像分类、目标检测和分割任务上对 DiNAT 进行了大量实验，发现其在多个下游任务中相比注意力机制和卷积基线均有显著提升。此外，我们还探索了各向同性与混合注意力变体、ImageNet-22K 预训练下的扩展实验，以及不同扩张值的影响。我们也在先进的分割框架中实现了图像分割任务的 state-of-the-art 表现。
• 扩展了 NAT TEN（NA 的 PyTorch CUDA 扩展），添加了扩张支持和 bfloat16 支持，使得相关研究可以推广到更多任务与应用中。

尽管 DiNAT 的初步实验已在多个下游视觉任务上展现出显著改进，但其性能与应用远未止步于此。NA 的局部注意力与 DiNA 的稀疏全局注意力互为补充：它们能够同时保留局部性、建模长距离依赖、呈指数级扩展感受野，并保持线性复杂度。对 self attention 的限制可能还能通过避免其冗余交互（例如与重复、背景或干扰 token 的交互）来加快收敛速度 [26,36]。局部注意力与稀疏全局注意力的组合可能为视觉任务及更广泛领域带来强大助力。为支持该方向的研究，我们开源了完整项目，包括我们修改的 NAT TEN，其相较于朴素实现可将运行时间减少若干数量级 [15]。

2. 相关工作

我们简要回顾点积 self attention（DPSA）、Transformer [42] 和 Vision Transformer [13]。随后，我们介绍一些局部化 self attention 模块，如 SASA [35]、Swin Transformer 中的 SWSA [29] 以及 Neighborhood Attention Transformer 中的 NA [15]，并讨论它们的局限性，这些正是促使我们进行本研究的动机。最后，我们讨论稀疏注意力机制在语言处理 [1,37] 和视觉 [6,20] 中的既有应用。

2.1. 自注意力

Vaswani 等人 [42] 将点积注意力定义为 query 与一组 key-value 对之间的操作。query 与 key 的点积会先经过缩放，然后通过 softmax 激活函数生成注意力权重。这些注意力权重随后应用于 value：
$$Attention(Q,K,V)=softmax\left(\frac{Q{K}^T}{\sqrt{d}}\right)V(1)$$
其中 $\sqrt{d}$ 是缩放参数，$d$ 是 key 的维度。点积 self attention 只是上述操作的一种特殊情况，其中 query、key 和 value 都是同一输入的线性投影。给定输入 $X\in {\mathbb{R}}^{n\times d}$，其中 $n$ 是 token 的数量，$d$ 是嵌入维度，该操作的时间复杂度为 $O(n^{2}d)$，而注意力权重的空间复杂度为 $O(n^2)$（空间复杂度依赖于具体实现方式 [10]）。Vision Transformer（ViT）[13] 是最早将纯 transformer 编码器应用于视觉任务的工作之一，展示了大规模基于 self attention 模型的强大能力。后续工作通过对训练策略的最小修改 [38]、结构变更 [39]，以及对小数据场景的应用 [16] 等方式，进一步拓展了研究。由于其时间复杂度为二次方，许多研究尝试通过限制 attention 的范围来减少计算量，尤其是在处理如 NLP 中的长文档 [1] 或视觉中的高分辨率图像 [29] 等大规模输入时。限制 self attention 可以采用多种模式，其中之一就是局部化。

2.2. 局部注意力

Stand-Alone Self Attention (SASA). SASA [35] 是最早专为视觉模型设计的局部 attention 机制之一，早于 ViT [13] 数年。它将 key-value 对设置为特征图上的滑动窗口，因此为每个 query（像素）局部化 attention 到以其为中心的窗口中。这种操作可以轻松替换现有 CNN（如 ResNet）中的卷积，并在理论上甚至能降低计算复杂度。尽管其展现出一定潜力，作者发现该模块由于实现效率低下导致整体模型运行缓慢。因此，后续工作转向更高效的替代方法，如 HaloNet [41] 中的 block self attention 和 Swin [29] 中的 Window Self Attention。

Shifted Window Self Attention (SWSA). Liu 等人 [29] 提出了 Window Self Attention（WSA）及其变体 Shifted Window Self Attention（SWSA），并将其用于视觉任务的分层模型 Swin Transformer。他们指出，SASA 等滑动窗口方法效率低下，是他们发展 Window Self Attention 的动机之一。如其名称所示，SWSA 在 attention 操作前先对特征图进行像素偏移，然后在操作后再逆向偏移，以在不同层间创建不同的窗口划分方式，从而实现窗口间的交互，这对于扩大感受野至关重要（见图 3）。Swin 最初在目标检测与语义分割任务中成为新的 state of the art，并激发了后续研究将其拓展至生成 [52]、图像恢复 [25]、掩码图像建模 [49]、视频动作识别 [31] 等更多任务中。此外，后续版本 Swin-V2 [28] 凭借其最大规模模型，再次刷新了 SOTA。值得注意的是，Swin-V2 使用了更大的窗口尺寸来达成这一性能提升，这也带来了更高的时间复杂度和内存开销。

Neighborhood Attention (NA). NA [15] 被提出为一种简单的滑动窗口 attention，它为每个像素局部化 self attention 到其最近邻居上。NA 在时间与空间复杂度以及参数数量上与 Swin 的 WSA 和 SWSA 相当，但它采用重叠滑动窗口操作，因此保留了平移等变性。虽然 NA 的滑动窗口模式与 SASA 类似，但其以最近邻为核心的 attention 公式使其成为 self attention 的一种直接限制，因此 NA 不同于 SASA，随着窗口大小的增大，可以逐渐逼近 SA。滑动窗口 attention 面临的一个主要挑战是缺乏高效的实现支持，因为现有的深度学习或 CUDA 库并不直接支持这类操作。因此，NA 被同时引入了 NAT TEN，即一个具备高效 CPU 与 GPU 内核的扩展，实现上使 NA 在速度和内存使用上优于如 WSA/SWSA 等模块。

模型 Neighborhood Attention Transformer (NAT) 在分层结构上与 Swin Transformer 相似。除了 attention 模块外，两者的主要区别在于：NAT 在下采样层中采用的是重叠卷积，而不是 Swin 中使用的 patch 操作。因此，为了在参数数量和 FLOPs 上保持与 Swin 的各个变体一致，NAT 模型采用了更深的网络，并使用了更小的 inverted bottlenecks。NAT 在图像分类任务中取得了优于 Swin 的结果，在下游任务中也有相当竞争力的表现。

尽管基于局部 attention 的模型因其对局部性的保持与高效性而能够在各种视觉任务中表现良好，但它们在捕获如 self attention 所具备的全局上下文方面仍显不足，而全局上下文同样对视觉任务至关重要。此外，局部化的 attention 机制相比 self attention 的完整感受野，拥有更小且增长缓慢的感受野，这一点类似于卷积。除了 self attention 之外，也有一些工作探索了在视觉任务中引入全局感受野的方法，其中包括但不限于 Non-local Neural Networks [45]。然而，具备不受限制的全局感受野的操作通常相比那些受限制（如局部或稀疏）的操作，会面临更高的计算复杂度。

2.3. 稀疏注意力

Child 等人 [6] 提出了 Sparse Transformers，除了支持更深层次的变体外，还采用了稀疏核注意力机制。通过这种机制，模型能够更高效地训练处理更长序列的数据。在稀疏注意力领域还有其他工作，如 Longformer [1]、Routing Transformers [37] 和 CCNet [20]，它们的共同特点是在不可避免需要处理较长 token 序列的情况下，通过降低 self attention 的计算成本，同时仍保留全局上下文。Longformer [1] 具体研究了结合带扩张和不带扩张的 1D 滑动窗口注意力，以及针对特定 token 的全局注意力，从而使模型能够处理长文档同时维持全局上下文。CCNet [20] 采用轴向注意力，通过引入全局上下文来提升语义分割头的表现，同时避免了无约束 self attention 的二次复杂度。最近，MaxViT [40] 探索了一种混合模型，结合了 MBConv、Window Attention [29] 和稀疏网格注意力，在 ImageNet 上取得了较高准确率，但该模型的复杂度较高，吞吐量低于 Swin [29]。

尽管这些非局部和稀疏限制的 self attention 机制表现出很大潜力，但在分层视觉 transformer 领域尚未得到充分研究。为了扩展局部感受野，并将全局上下文重新引入分层视觉 transformer，我们提出了 Dilated Neighborhood Attention（DiNA），这是对 NA 的扩展，通过增大步长使邻域覆盖更长范围，同时保持整体注意力跨度不变。DiNA 可作为一种稀疏且全局的操作，并在与仅局部的 NA 结合使用时最为有效。我们在图 4 中展示了感受野的示意图，将全连接层与卷积、扩张卷积进行比较，类似地也对 self attention、NA 和 DiNA 进行了比较。我们利用分层视觉 transformer Dilated Neighborhood Attention Transformer（DiNAT）提供了该论断的实证支持。

温馨提示：

阅读全文请访问"AI深语解构" DiNA：扩张邻域注意力 Transformer