[论文解读] Neighborhood Attention Transformer

Ali Hassani, Steven Walton+3 more|arXiv (Cornell University)|Apr 14, 2022

Advanced Neural Network Applications被引用 5

一句话总结

本文提出邻域注意力（NA），一种新型滑动窗口自注意力机制，将自注意力局限于每个像素的最近邻，实现线性时间与空间复杂度，同时保持平移等变性。NA 在 ImageNet 上达到 83.2% 的 top-1 准确率，在图像分类和下游视觉任务中优于 Swin Transformer 和 ConvNeXt，通过优化的 $π$ATTEN$ 软件包实现最高 40% 的推理速度提升和 25% 的内存占用减少。

ABSTRACT

We present Neighborhood Attention (NA), the first efficient and scalable sliding-window attention mechanism for vision. NA is a pixel-wise operation, localizing self attention (SA) to the nearest neighboring pixels, and therefore enjoys a linear time and space complexity compared to the quadratic complexity of SA. The sliding-window pattern allows NA's receptive field to grow without needing extra pixel shifts, and preserves translational equivariance, unlike Swin Transformer's Window Self Attention (WSA). We develop NATTEN (Neighborhood Attention Extension), a Python package with efficient C++ and CUDA kernels, which allows NA to run up to 40% faster than Swin's WSA while using up to 25% less memory. We further present Neighborhood Attention Transformer (NAT), a new hierarchical transformer design based on NA that boosts image classification and downstream vision performance. Experimental results on NAT are competitive; NAT-Tiny reaches 83.2% top-1 accuracy on ImageNet, 51.4% mAP on MS-COCO and 48.4% mIoU on ADE20K, which is 1.9% ImageNet accuracy, 1.0% COCO mAP, and 2.6% ADE20K mIoU improvement over a Swin model with similar size. To support more research based on sliding-window attention, we open source our project and release our checkpoints at: https://github.com/SHI-Labs/Neighborhood-Attention-Transformer .

研究动机与目标

• 为解决视觉变换器中标准自注意力的二次方复杂度问题，该问题限制了在高分辨率输入上的可扩展性。
• 开发一种高效、可扩展且可微分的滑动窗口注意力机制，保持平移等变性与局部归纳偏置。
• 通过优化核设计，克服先前局部注意力方法（如 SASA 和 Swin 的 WSA）的低效与实现瓶颈。
• 设计一种新型层次化变换器架构——邻域注意力变换器（NAT），利用 NA 实现图像分类与密集预测任务的最先进性能。
• 发布 $π$ATTEN$ 库与模型权重，推动高效滑动窗口注意力在视觉研究中的广泛应用。

提出的方法

• NA 定义了一种基于像素的注意力机制，其中每个 token 仅关注其最近邻的像素，形成动态的局部感受野。
• 该机制被表述为一种可微分操作，对每个像素周围固定大小的邻域计算加权和，避免全局注意力计算。
• 邻域通过空间核（如 3×3、5×5）定义，注意力权重通过局部区域内查询与键的相似性计算得出。
• 作者通过自定义的 C++ 与 CUDA 内核库 $π$ATTEN$ 高效实现 NA，实现高速推理与更低的内存占用。
• NAT 使用重叠卷积 patch 嵌入与更深更细的网络结构，配合残差块，增强表征能力。
• 模型采用层次化设计，包含下采样与多尺度特征学习，类似于 Swin，但将 WSA 替换为 NA 以提升效率与性能。

实验结果

研究问题

• RQ1 基于像素的滑动窗口注意力机制是否能在保持与全局自注意力相当性能的同时实现线性复杂度？
• RQ2 NA 是否比基于窗口或基于块的注意力机制更有效地保持平移等变性与局部归纳偏置？
• RQ3 NA 的高效实现是否能在速度与准确率上超越现有注意力机制（如 WSA 与 SASA）？
• RQ4 在层次化变换器架构中，将 WSA 替换为 NA 对图像分类与密集预测基准性能的提升程度如何？
• RQ5 NA 的感受野增长与 Swin 的移位窗口机制相比，是否无需额外操作即可实现更优效果？

主要发现

• NAT-Tiny 在 ImageNet 上达到 83.2% 的 top-1 准确率，比参数量相近的 Swin-Tiny 提升 1.9%。
• 在 MS-COCO 目标检测任务中，NAT-Tiny 达到 51.4% 的 mAP，比 Swin-Tiny 提升 1.0%。
• 在 ADE20K 语义分割任务中，NAT-Tiny 达到 48.4% 的 mIoU，比 Swin-Tiny 提升 2.6%。
• $π$ATTEN$ 库使 NA 的推理速度比 Swin 的 WSA 实现最高提升 40%，内存占用减少最多 25%。
• 消融实验表明，在 Swin-Tiny 中用 NA 替代 WSA 可使准确率提升 0.5%，而 SASA 则导致轻微下降。
• NAT-Tiny 中 7×7 的核大小在性能与效率之间达到最佳平衡，实现 83.2% 的 ImageNet 准确率与 1537 imgs/sec 的吞吐量。