[論文レビュー] Neighborhood Attention Transformer

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

Advanced Neural Network Applications被引用数 5

ひとこと要約

この論文は、自己注意を各画素の最近傍画素に局所化する新しいスライディングウィンドウ自己注意メカニズムであるネイバーヒードアテンション（NA）を紹介している。これにより、線形時間・空間計算量を実現するとともに、並進等変性を維持する。NAはImageNetでトップ1正解率83.2%を達成し、画像分類および下流のビジョンタスクにおいてスウィントランスフォーマーとコンブレクストを上回る性能を発揮する。最適化された$π$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 .

研究の動機と目的

• 高解像度入力におけるビジョントランスフォーマーの標準自己注意の2次的計算量の問題を解決すること。
• 並進等変性と局所的なインダクティブバイアスを維持しながら、効率的でスケーラブルかつ微分可能なスライディングウィンドウ自己注意メカニズムを開発すること。
• SASA やスウィンのWSAのような先行の局所自己注意手法の非効率性と実装上のボトルネックを、最適化されたカーネル設計によって克服すること。
• NAを活用した階層的トランスフォーマー・アーキテクチャ、ネイバーヒードアテンショントランスフォーマー（NAT）を設計し、画像分類および密度予測タスクで最先端の性能を達成すること。
• 効率的なスライディングウィンドウ自己注意の研究を広く促進するために、$π$ATTEN$ライブラリおよびモデルのチェックポイントを公開すること。

提案手法

• NAは、各トークンが自身の最近傍ピクセルにのみ注目するピクセル単位のアテンションメカニズムを定義し、動的で局所的な受容 field を形成する。
• このメカニズムは、各ピクセルの固定サイズの近傍領域内で重み付き和を計算する微分可能な演算として定式化されており、グローバルなアテンション計算を回避する。
• 近傍領域は空間カーネル（例：3×3、5×5）によって定義され、局所領域内でのクエリとキーの類似度に基づいてアテンション重みが計算される。
• 著者らは、$π$ATTEN$と呼ばれる独自のC++およびCUDAカーネルライブラリを用いてNAを効率的に実装し、高速な推論と低メモリ使用量を実現している。
• NATは重複する畳み込みパッチエンコーディングと、残差ブロックを用いたより深く細いアーキテクチャを採用し、表現力の向上を図っている。
• スウィンと同様にダウンサンプリングとマルチスケール特徴量学習を組み合わせた階層的設計を採用しているが、WSAの代わりにNAを採用することで、効率性と性能が向上している。

実験結果

リサーチクエスチョン

• RQ1 ピクセル単位のスライディングウィンドウ自己注意メカニズムは、グローバル自己注意と同等の性能を維持しながら、線形計算量を達成できるか？
• RQ2 NAは、ウィンドウベースやブロックベースのアテンションメカニズムと比較して、並進等変性と局所的なインダクティブバイアスをより効果的に保持できるか？
• RQ3 NAの効率的な実装は、WSA や SASA といった既存のアテンションメカニズムを、速度と精度の両面で上回れるか？
• RQ4 階層的トランスフォーマー・アーキテクチャにおいてWSAをNAに置き換えることで、画像分類および密度予測ベンチマークでの性能がどの程度向上するか？
• RQ5 追加の演算を必要とせずに、NAの受容 field の成長はスウィンのシフトウィンドウ機構と比較してどの程度か？

主な発見

• NAT-TinyはImageNetでトップ1正解率83.2%を達成し、同程度のモデルサイズのスウィン-Tinyと比較して1.9%の向上を示した。
• MS-COCOオブジェクト検出では、NAT-Tinyが51.4%のmAPを達成し、スウィン-Tinyの1.0%上回った。
• ADE20Kセマンティックセグメンテーションでは、NAT-Tinyが48.4%のmIoUを達成し、スウィン-Tinyと比較して2.6%の向上を示した。
• $π$ATTEN$ライブラリのおかげで、NAはスウィンのWSA実装と比較して最大40%速く、最大25%少ないメモリを使用する。
• アブレーションスタディの結果、スウィン-TinyのWSAをNAに置き換えることで、ベースラインから正解率が0.5%向上した一方、SASAはわずかに低下した。
• NAT-Tinyにおける7×7カーネルサイズが、性能と効率の最良のバランスを実現しており、ImageNet正解率83.2%、スループット1537 imgs/secを達成した。