PaperlayerView exact revision on arXiv

Source-linked AI summary

Neighborhood Attention Transformer

Ali Hassani, Steven Walton, Jiachen Li, Shen Li, Humphrey Shi

arXiv:2204.07143v5cs.CVcs.AIcs.LG

TL;DR

Vision self attention is powerful but becomes computationally and memory intensive as image resolution grows. The paper introduces Neighborhood Attention, NATTEN, and the NAT architecture to provide efficient local attention while preserving translational equivariance. NAT reports strong results across ImageNet classification, MS-COCO detection, and ADE20K segmentation, while NATTEN improves practical speed and memory use relative to Swin’s WSA.

  • Problem

    Self attention has quadratic complexity in vision token count, increasing computation and memory at high resolution and lacking convolutional locality and spatial inductive biases.

  • Method

    The paper proposes pixel-wise Neighborhood Attention with growing local receptive fields, implements it through NATTEN kernels, and builds the hierarchical NAT transformer.

  • Results

    NAT outperforms Swin Transformer and ConvNeXt in image classification and matches or exceeds them on downstream vision tasks, with NAT-Tiny reaching 83.2% ImageNet top-1 accuracy.

  • Takeaways & Limitations

    Neighborhood Attention provides a scalable sliding-window alternative that preserves translational equivariance, while NAT demonstrates competitive performance across vision tasks.

Abstract

from arXiv · show

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 .

1. Introduction

Vision Transformers offer strong performance but self attention scales quadratically with token count, motivating efficient local alternatives. The paper proposes Neighborhood Attention and NAT, supported by NATTEN kernels, with competitive classification and downstream results.

  • Motivation: Self attention has quadratic complexity and memory usage in the number of vision tokens, making high-resolution downstream tasks difficult.Vision token count typically grows with image resolution, while convolutions provide locality and spatial inductive biases that self attention lacks.
  • Neighborhood Attention: Neighborhood Attention localizes each pixel’s attention to its nearest neighborhood while preserving translational equivariance.Its neighborhood definition gives pixels an identical attention span and approaches self attention as neighborhood size grows.
  • NATTEN: NATTEN provides efficient C++ and CUDA kernels for sliding-window attention, allowing NA to run up to 40% faster than Swin’s WSA while using up to 25% less memory.The package provides autograd-supported modules that can be integrated into PyTorch pipelines.
  • NAT Results: NAT-Tiny reaches 83.2% top-1 accuracy on ImageNet, 51.4% box mAP on MS-COCO, and 48.4% mIoU on ADE20K.The model uses 4.3 GFLOPs and 28M parameters and is reported to outperform Swin Transformer and ConvNeXt.
  • NAT Results: NAT outperforms Swin Transformer and ConvNeXt on ImageNet-1K classification with fewer FLOPs and a similar number of parameters.The comparison is presented as classification performance versus compute, with bubble size representing parameter count.

2. Related Works

Related vision architectures balance attention’s flexibility against computational cost and inductive biases. Neighborhood Attention extends local attention with pixel-wise neighborhoods, avoiding shifted windows while supporting comparisons across classification, detection, and segmentation.

  • Self Attention: Self attention applies dot-product attention to projected inputs and has O(n^2d) computational complexity with O(n^2) attention-weight space complexity.Here, n is the number of tokens and d is the embedding dimension.
  • Vision Transformers: Vision Transformer uses a Transformer encoder over embedded image patches, but its competitiveness on medium-scale datasets was associated with large-scale pretraining and limited convolutional inductive biases.The cited discussion contrasts ViT’s patch-based design with CNNs’ inherent inductive biases.
  • Local Attention: SASA applies sliding-window attention with zero padding and preserves translational equivariance, but its implementation suffered high latency despite lower theoretical cost.Its window pattern resembles convolution and was intended to replace convolutions in CNNs.
  • Swin and ConvNeXt: Swin’s WSA partitions inputs into windows and pairs with shifted WSA to enable out-of-window interactions, while ConvNeXt provides a strong CNN baseline.ConvNeXt was designed as a fair comparison baseline for convolutional and attention-based models.
  • Neighborhood Attention: Neighborhood Attention localizes each query’s receptive field without cyclic shifts and is evaluated against Swin on classification, detection, and semantic segmentation.The paper reports that NA maintains translational equivariance and can run faster while using less memory through NATTEN.

3. Method

Neighborhood Attention localizes self-attention to each pixel’s nearest neighbors, retaining linear scaling while preserving translational equivariance and allowing the receptive field to grow. NATTEN supplies optimized kernels for practical speed and memory efficiency, and NAT organizes NA into a hierarchical vision transformer.

  • Neighborhood Attention: NA assigns each pixel an attention span over its nearest neighboring pixels, rather than a fixed zero-padded window.This gives pixels an identical attention span and distinguishes NA from SASA’s padded windows.
  • Neighborhood Attention: NA preserves translational equivariance and expands its receptive field without Swin’s manually shifted attention variant.Dynamic neighborhoods allow neighboring interactions while avoiding blocked windows and their shifted counterpart.
  • Neighborhood Attention: As neighborhood size k grows, NA approaches self-attention and equals it at maximum neighborhood size.This convergence does not hold for SASA because zero padding remains around the input.
  • Tiled NA and N ATTEN: NATTEN provides C++ and CUDA kernels, tiled NA, half-precision support, 1D/2D operation, and autograd-compatible PyTorch integration.Tiled NA loads non-overlapping query tiles into shared memory and can reduce latency by up to an order of magnitude versus a naive implementation.
  • Neighborhood Attention Transformer: NAT uses four hierarchical levels of NAT Blocks, downsampling feature maps by half while doubling depth between levels.Each block combines multi-headed NA, an MLP, pre-module LayerNorm, and skip connections; the hierarchy supports feature pyramids for downstream tasks.

4. Experiments

Experiments evaluate NAT across classification, detection, instance segmentation, and semantic segmentation, alongside attention-pattern and architecture ablations. NAT generally outperforms or matches comparable Swin and ConvNeXt models, with efficiency advantages in several settings.

  • 4.1. Classification: NAT variants consistently outperform comparable Swin and ConvNeXt models on ImageNet-1K classification, with competitive throughput and memory usage.NAT-Mini exceeds Swin-Tiny by 0.5% with fewer parameters, higher throughput, and lower memory usage.
  • 4.2. Object Detection and Instance Segmentation: NAT-Tiny outperforms Swin-Tiny and ConvNeXt counterparts on MS-COCO with both Mask and Cascade Mask R-CNN, using slightly fewer FLOPs than comparable models.NAT-Mini exceeds Swin-Tiny with Mask R-CNN but falls slightly short with Cascade Mask R-CNN.
  • 4.2. Object Detection and Instance Segmentation: NAT-Base can outperform its Swin counterpart on MS-COCO while also achieving higher throughput.NAT-Small is competitive with Swin while being faster.
  • 4.3. Semantic Segmentation: NAT-Tiny significantly outperforms ConvNeXt-Tiny on ADE20K while being slightly more efficient, whereas NAT-Base performs on-par with Swin-Base.NAT-Small outperforms Swin-Small on single-scale performance while matching multi-scale performance.
  • 4.4. Ablation Study: Replacing Swin’s attention with NA improves accuracy by almost 0.5%, while replacing it with SASA causes a slight decrease.The ablation also attributes gains to overlapping convolutional downsampling and the deeper-thinner NAT configuration.

5. Conclusion

The paper introduces Neighborhood Attention and NATTEN for efficient sliding-window attention, then develops NAT and evaluates its attention mechanisms and kernel sizes.

  • 5. Conclusion: Neighborhood Attention localizes self attention to each pixel’s nearest neighborhood, providing linear complexity, local inductive biases, and translational equivariance.The paper contrasts NA with blocked and window self attention, and notes that its window size grows toward self attention at corner cases.
  • 5. Conclusion: Table 6 compares attention mechanisms in models matched to Swin-T architecture, with SASA implemented using NATTEN.This setup isolates the effect of replacing the attention mechanism while giving SASA the same speed and memory implementation as NA.
  • 5. Conclusion: In the ablation, replacing SWSA with NA improves accuracy by almost 0.5%, whereas replacing it with SASA slightly decreases accuracy.The table uses Swin-T as the baseline and also reports gains from overlapping convolutions and the NAT configuration.
  • 5. Conclusion: NATTEN demonstrates that explicit sliding-window attention patterns can be efficiently parallelized for NA-based models.The authors present NATTEN to challenge the view that sliding-window patterns are not efficient or parallelizable.
  • 5. Conclusion: NAT is a hierarchical transformer design based on Neighborhood Attention, and the paper reports competitive image classification and downstream vision performance.The supplied passages identify NAT as the proposed model but do not provide its headline benchmark values.

A.1. Background

Sliding-window attention needs specialized implementations for scalable parallelization. NATTEN supplies a PyTorch extension that integrates Neighborhood Attention for token- and pixel-based models.

  • A.1. Background: 13% speed and approximately 9 times as much memory describe optimized plain-PyTorch NA relative to Swin on a 56 × 56 feature map.This example illustrates why a specialized implementation is needed for sliding-window attention to scale.
  • A.1. Background: 102% speed and approximately 20% less memory describe the same NA module with a naive CUDA implementation relative to Swin.The comparison uses the same 56 × 56 feature map and Swin reference as the plain-PyTorch result.
  • A.1. Background: 132% speed with no change in memory usage describes the same module using the Tiled NA algorithm relative to Swin.The paper points to additional relative speed and memory benchmarks in its figures.
  • A.1. Background: NATTEN provides PyTorch modules NeighborhoodAttention1D and NeighborhoodAttention2D for integrating NA into token- and pixel-based models.The extension supplies the implementation needed for PyTorch users to use NA in both settings.
  • A.1. Background: Each attention module uses linear projections for queries, keys, and values followed by a final linear projection.The paper describes this projection structure as standard for dot-product self-attention modules.

A.2. Naive CUDA Kernels

The paper implements NA with CUDA kernels and compares memory and throughput across implementations. Naive kernels improve over plain PyTorch but remain less optimized than tiled execution.

  • A.2. Naive CUDA Kernels: 80% to 130% speed relative to WSA+SWSA was achieved by the initial naive CUDA kernels despite their simple threadblock-based design.These kernels did not use shared memory or warp optimization and used kernel size 7 × 7.
  • A.2. Naive CUDA Kernels: Naive CUDA kernels divide computation across available threadblocks but read directly from global GPU memory.The paper identifies the lack of shared-memory and warp optimization as part of their non-optimal design.
  • A.2. Naive CUDA Kernels: Half-precision support requires CUDA’s vectorized half2 data type and operators to produce significant bandwidth improvement.Simply compiling kernels for half tensors is described as insufficient for substantial bandwidth gains.
  • A.2. Naive CUDA Kernels: NATTEN’s layer-wise memory usage is consistently lower than Swin’s, whereas plain PyTorch NA develops a quickly growing memory footprint.The comparison is shown in Figure I and specifically distinguishes implementations with and without NATTEN.
  • A.2. Naive CUDA Kernels: Torch-based, naive, and tiled NA implementations are compared with WSA+SWSA using relative throughput on a single A100 GPU.The plain PyTorch implementation runs out of memory for resolutions 448^2 and higher.

A.4. Tiled Neighborhood Attention

Tiled Neighborhood Attention partitions queries into tiles and uses overlapping key tiles to compute outputs efficiently in shared memory. NATTEN provides optimized CPU and GPU implementations, while the current kernels retain hardware- and configuration-dependent constraints.

  • Tiled Neighborhood Attention: Tiled attention partitions inputs into non-overlapping tiles and computes outputs from values loaded into shared memory.Each thread reads adjacent global-memory cells, synchronizes, and then computes outputs using shared-memory values.
  • Tiled Neighborhood Attention: Tile size depends on kernel size, embedding dimension, GPU shared-memory capacity, and threadblock limits.Shared-memory tiling also introduces bank conflicts and unequal numbers of reads across inputs.
  • Tiled Neighborhood Attention: For a 7 × 7 kernel and 3 × 3 query tile, the required key tile is 9 × 9.The key-tile relation is tq + k − 1, giving 3 + 7 − 1 = 9 in this example.
  • Tiled Neighborhood Attention: NATTEN uses less memory than comparable WSA+SWSA models because NA avoids pixel shifts and masked attention, with positional biases fused into kernels.Figure III compares layer-wise memory usage between NAT and Swin.
  • Tiled Neighborhood Attention: NATTEN remains an early implementation, with planned kernel improvements, broader use-case support, and Hopper/CUDA 12 support.The authors also hope to implement Neighborhood Attention with implicit GEMM using packages such as CUTLASS.
  • Tiled Neighborhood Attention: NATTEN supports CPU training and inference through C++ kernels with AVX vectorization and multithreaded computation.The current version includes seven CPU kernels.

B.1. Ablation on RPB

The ablation examines relative positional biases and pixel shifts in WSA, using Table I to compare NA and WSA variants with and without positional biases.

  • B.1. Ablation on RPB: The study presents an ablation of relative positional biases and pixel shifts, with pixel shifts considered for WSA only.The ablation is reported in Table I.
  • B.1. Ablation on RPB: The supplied passage identifies the ablation factors but does not report the table’s numerical outcomes.No metric values are stated in the provided text.
  • B.1. Ablation on RPB: Table I compares NA and WSA with and without positional biases.Swin’s results are directly reported from the original paper.

B.2. Saliency analysis

The saliency analysis compares ViT-Base, Swin-Base, and NAT-Base on selected ImageNet validation images using output-derived salient maps.

  • B.2. Saliency analysis: Salient maps compare ViT-Base, Swin-Base, and NAT-Base on selected ImageNet validation images.The maps are created from the models’ outputs.
  • B.2. Saliency analysis: The selected images depict a Bald Eagle, Acoustic Guitar, Hummingbird, and Steam Locomotive.These are the stated ground-truth classes.
  • B.2. Saliency analysis: All listed images are correctly predicted except ViT’s Acoustic Guitar image, which is classified as Stage.The passage reports this prediction outcome for the visualization set.

C. Notes on translational equivariance

The paper defines translational equivariance as commuting with translation and analyzes how attention patterns affect it. NA largely preserves the property, while Swin’s partitioning breaks it and NA relaxes it at corners to maintain attention span.

  • Definitions: Translation in computer vision refers to shifting, and sometimes rotating, pixels.The paper uses T as a translation function in its equivariance discussion.
  • Definitions: A function f is equivariant to T when translating its output equals applying f after translating its input.This is expressed as T(f(x)) = f(T(x)).
  • Existing operations: Linear projections and convolutions are translationally equivariant, while self-attention is equivariant because its projections and attention weights preserve pixel-order changes.SASA also maintains translational equivariance through its sliding raster-scan pattern and window-dependent dynamic weights.
  • Swin attention: WSA and SWSA break translational equivariance because fixed partition boundaries change when objects are translated.An object can move from one WSA partition into two, altering the attention pattern.
  • Neighborhood Attention: NA preserves translational equivariance for the most part but relaxes it at corner pixels to maintain attention span.Repeated windows at edge pixels cause the corner-case relaxation.
  • Neighborhood Attention: NA’s repeated corner windows keep attention spans identical across the feature map and remain subsets of the original self-attention weights.The paper distinguishes this attention behavior from convolutional repeated windows, which produce repeated outputs.
Loading 2204.07143v5…