Neighborhood Attention Transformer

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 MSCOCO 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.

1. Introduction

Convolutional neural networks (CNNs) [19] have been the de facto standard architecture for computer vision models across different applications for years. AlexNet [18] showed their usefulness on ImageNet [10], and many others followed suit with architectures such as VGG [26], ResNet [17], and EfficientNet [27]. Transformers [31] on the other hand, were originally proposed as attention-based models for natural language processing (NLP), trying to ex-ploit the sequential structure of language. They were the ba-sis upon which BERT [11] and GPT [2], [2]3,[2]4 were built, and they continue to be the state of the art architecture in NLP.