邻域注意力 Transformer
计算机视觉与模式识别
2023-05-18 v5 人工智能
机器学习
摘要
我们提出邻域注意力（NA），这是首个面向视觉的高效且可扩展的滑动窗口注意力机制。NA 是一种逐像素操作，将自注意力（SA）局限于最近的相邻像素，因此相比 SA 的二次复杂度，具有线性的时间与空间复杂度。该滑动窗口模式使 NA 的感受野无需额外像素偏移即可增长，并且与 Swin Transformer 的窗口自注意力（WSA）不同，保持了平移等变性。我们开发了 NATTEN（Neighborhood Attention Extension），一个带有高效 C++ 与 CUDA 内核的 Python 包，使 NA 的运行速度比 Swin 的 WSA 快至多 40%，同时内存占用减少至多 25%。我们进一步提出邻域注意力 Transformer（NAT），一种基于 NA 的新型分层 transformer 设计，提升了图像分类及下游视觉任务性能。NAT 的实验结果具有竞争力：NAT-Tiny 在 ImageNet 上达到 83.2% 的 top-1 准确率，在 MS-COCO 上达到 51.4% mAP，在 ADE20K 上达到 48.4% mIoU，相比同等规模的 Swin 模型，ImageNet 准确率提升 1.9%，COCO mAP 提升 1.0%，ADE20K mIoU 提升 2.6%。为支持更多基于滑动窗口注意力的研究，我们开源了项目并在 https://github.com/SHI-Labs/Neighborhood-Attention-Transformer 发布了检查点。
引用
@article{arxiv.2204.07143,
title = {Neighborhood Attention Transformer},
author = {Ali Hassani and Steven Walton and Jiachen Li and Shen Li and Humphrey Shi},
journal= {arXiv preprint arXiv:2204.07143},
year = {2023}
}
备注
To appear in CVPR 2023. NATTEN is open-sourced at: https://github.com/SHI-Labs/NATTEN/