[논문 리뷰] Neighborhood Attention Transformer

Ali Hassani, Steven Walton외 3명|arXiv (Cornell University)|2022. 04. 14.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 각 픽셀의 가장 가까운 이웃으로 자기주의 주목을 국소화하는 새로운 슬라이딩 윈도우 자기주의 메커니즘인 Neighborhood Attention (NA)을 소개한다. 이는 선형 시간 및 공간 복잡도를 보장하면서도 이동 불변성(translation equivariance)을 유지한다. 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와 같은 비효율성과 구현의 성능 저하 요인을 최적화된 커널 설계를 통해 극복하고자 한다.
• Neighborhood Attention Transformer (NAT)라는 새로운 계층적 트랜스포머 아키텍처를 설계하여 NA를 활용해 이미지 분류 및 조밀 예측 작업에서 최신 성능을 달성하고자 한다.
• 비전 분야의 효율적인 슬라이딩 윈도우 주목을 위한 연구를 널리 확산시키기 위해 $π$ATTEN$ 라이브러리와 모델 체크포인트를 공개하고자 한다.

제안 방법

• NA는 각 토큰이 오직 자신의 가장 가까운 이웃 픽셀들만 주목하도록 하는 픽셀 단위의 주목 메커니즘을 정의한다. 이는 동적이고 국소적인 수신 필드를 형성한다.
• 이 메커니즘은 전역 주목 계산을 피하기 위해 각 픽셀 주변의 고정 크기 영역에서 가중치 합을 계산하는 미분 가능한 연산으로 수식화된다.
• 이웃 영역은 공간 커널(예: 3×3, 5×5)을 통해 정의되며, 주목 가중치는 국소 영역 내에서 쿼리-키 유사도를 기반으로 계산된다.
• 저자들은 고속 추론과 낮은 메모리 사용량을 가능하게 하기 위해 고유의 C++ 및 CUDA 커널 라이브러리인 $π$ATTEN$을 활용해 NA를 효율적으로 구현했다.
• NAT는 겹치는 컨볼루션 패치 임베딩과 잔차 블록을 사용한 더 깊고 얇은 아키텍처를 채택하여 표현 능력을 향상시켰다.
• Swin과 유사한 계층적 설계를 채택했지만, WSA 대신 NA를 도입함으로써 효율성과 성능을 향상시켰다.

실험 결과

연구 질문

• RQ1 픽셀 단위의 슬라이딩 윈도우 주목 메커니즘이 전역 자기주의와 유사한 성능을 유지하면서도 선형 복잡도를 달성할 수 있는가?
• RQ2 NA는 윈도우 기반 또는 블록 기반 주목 메커니즘보다 이동 불변성과 국소적 인덕티브 바이어스를 더 효과적으로 유지하는가?
• RQ3 NA의 효율적인 구현이 WSA나 SASA와 같은 기존 주목 메커니즘보다 속도와 정확도 면에서 뛰어나게 성능을 높일 수 있는가?
• RQ4 계층적 트랜스포머 아키텍처에서 WSA를 NA로 대체할 경우, 이미지 분류 및 조밀 예측 벤치마크에서 성능 향상은 어느 정도 이루어지는가?
• RQ5 NA의 수신 필드 성장은 추가 연산 없이도 Swin의 시프트된 윈도우 메커니즘과 비교해 어떠한가?

주요 결과

• NAT-Tiny는 모델 크기가 유사한 Swin-Tiny와 비교해 ImageNet에서 83.2%의 top-1 정확도를 기록했으며, 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에서 WSA를 NA로 대체하면 기준 모델 대비 정확도가 0.5% 향상되었고, SASA 대비 약간의 하락이 관찰되었다.
• NAT-Tiny에서 7×7 커널 크기가 성능과 효율성의 최적 균형을 이룩했으며, ImageNet 정확도 83.2%와 1537 imgs/sec의 처리 속도를 기록했다.