NAT: スライディングウィンドウアテンションの革新

論文、コード、YouTube

• Ali Hassani, Steven Walton, Jiachen Li, Shen Li, Humphrey Shi, Neighborhood Attention Transformer, arXiv:2204.07143, 2022
• コード(NAT)(https://github.com/SHI-Labs/Neighborhood-Attention-Transformer)
• コード(NATTEN)(https://github.com/SHI-Labs/NATTEN)
• YouTube(https://www.youtube.com/watch?v=Ya4BfioxIHA)

掲載している画像は、本論文から引用となります。

概要

この論文では、Vision Transformerの新しいアテンション手法である、Neighborhood Attention(NA)を提案しています。NAは、以下の一番右の図のように、クエリ(濃い赤と青のピクセル)に対して、その周辺近傍に注目を向ける局所的なアテンション機構です。これは、赤のクエリが1ピクセル移動すると、注目するウィンドウ領域も一緒に移動するため、スライディングウィンドウアテンションとなります。NAを効率的に計算するために、C++およびCUDAカーネルを備えたPythonパッケージであるNATTEN(Neighborhood Attention Extension)も開発されました。NATTENは、Swin Transformerのウィンドウベースアテンションよりも最大40%高速で最大25%メモリ使用量を削減することができます。さらに、NAをベースにした階層型トランスフォーマーであるNeighborhood Attention Transformer(NAT)も提案しています。NATは、画像認識タスクで、Swin TransfomerやConvNeXtと同等かそれ以上の高い性能を実現しています。なお、NAを拡張したDiNAも存在します。

モデル	過去の記事
Swin Transformer	https://qiita.com/kinkalow/items/cb1024c2c9856ee1afca

NAがもたらす革新

グローバルアテンションの計算コストを削減するために、ローカルアテンションが提案されています。その初期例であるStand-Alone Self-Attention(SASA)は、NAと同じスライディングウィンドウアテンションですが、クエリが画像境界近傍にある場合、注目領域が異なります。動画でこれらの違いを確認できます。SASAの注目領域はこちらにあり、NAの注目領域はこちらにあります。以下の図でも示されていますが、SASAは畳み込みと類似しており、常にクエリが注目領域の中心にあるため、境界近傍では画像からはみ出る無駄な領域が発生します。これはゼロパディングによって処理されますが、情報のある注目領域が小さくなるというデメリットがあります。一方、NAは画像からはみ出る領域が発生することなく、すべてのクエリが同一の注目領域を持つように処理されます。さらに、NAはウィンドウサイズが大きくなるにつれてグローバルアテンションに近づく性質があります。それに加えて、NAはNATTENの開発により高速に実行でき、メモリ消費量も少なくすることができます。また、SASAとNAはtranslational equivarianceを保持していますが、これはSwin Transformerのローカルアテンションとは異なります。

SASA
　＊＊＊　　　　＊＊＊　　　　＊＊＊　　　　＊＊＊　　　　＊＊＊
　＊Ｑ＊ーー＋　＊Ｑ＊ー＋　＋＊Ｑ＊＋　＋ー＊Ｑ＊　＋ーー＊Ｑ＊
　＊＊＊　　｜　＊＊＊　｜　｜＊＊＊｜　｜　＊＊＊　｜　　＊＊＊
　　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜
　　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜
　　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋

　＊＊＊ーー＋　＊＊＊ー＋　＋＊＊＊＋　＋ー＊＊＊　＋ーー＊＊＊
　＊Ｑ＊　　｜　＊Ｑ＊　｜　｜＊Ｑ＊｜　｜　＊Ｑ＊　｜　　＊Ｑ＊
　＊＊＊　　｜　＊＊＊　｜　｜＊＊＊｜　｜　＊＊＊　｜　　＊＊＊
　　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜
　　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋

NA
　　Ｑ＊＊ー＋　＊Ｑ＊ー＋　＋＊Ｑ＊＋　＋ー＊Ｑ＊　＋ー＊＊Ｑ
　　＊　＊　｜　＊　＊　｜　｜＊　＊｜　｜　＊　＊　｜　＊　＊
　　＊＊＊　｜　＊＊＊　｜　｜＊＊＊｜　｜　＊＊＊　｜　＊＊＊
　　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜
　　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋

　　＊＊＊ー＋　＊＊＊ー＋　＋＊＊＊＋　＋ー＊＊＊　＋ー＊＊＊
　　Ｑ　＊　｜　＊Ｑ＊　｜　｜＊Ｑ＊｜　｜　＊Ｑ＊　｜　＊　Ｑ
　　＊＊＊　｜　＊＊＊　｜　｜＊＊＊｜　｜　＊＊＊　｜　＊＊＊
　　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜　｜　　　｜
　　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋　＋ーーー＋

translational equivariance

translational equivarianceとは、操作$f$がtranslation($T$)に対して$T(f(x))=f(T(x))$を満たす性質です。以下の図は、$T$を移動または回転、$f$を$\text{Sw}$または$\text{NA}^2$と選んだとき、translational equivarianceを保持するかどうかの実験結果を表しています。$\text{Sw}$はSwin TransformerのW-MSAとSW-MSAを順番に適用し、間に残差結合を挿入する操作を表します。一方、$\text{NA}^2$は2層のNAを適用し、間に残差結合を挿入することを表します。NAはスライディングウィンドウ特性によりtranslational equivarianceを保持できますが、Swin Transformerは保持できません。これは、Swin Transformerのウィンドウ分割により、オブジェクトが分割され、認識に影響が出る可能性があるためです。translational equivarianceを満たすことは、$f$が移動や回転に対して頑健であることを意味します。そのため、入力データに位置や回転のずれがあっても、NA操作は出力結果に影響が出ないようになります。

メソッド

Neighborhood Attention(NA)

NAは、ある位置$i$のクエリ$Q_i$と、その近傍位置$\rho_{1\le j\le k}(i)$($k$は近傍サイズ)におけるキー$K_{\rho_j(i)}$とバリュー$V_{\rho_j(i)}$、そして相対位置バイアス$B_{(i,\rho_j(i))}$を用いて計算されます。

\displaylines{
\begin{array}
  \boldsymbol{A}_i^k =
  \begin{bmatrix}
    Q_i K_{\rho_1(i)}^T + B_{(i,\rho_1(i))} \\
    \vdots \\
    Q_i K_{\rho_k(i)}^T + B_{(i,\rho_k(i))} \\
  \end{bmatrix}, \\
  \boldsymbol{V}_i^k=\biggl[V_{\rho_1(i)}^T, \cdots, V_{\rho_k(i)}^T\biggr]^T, \\
  NA_k(i) = \text{softmax}\Biggl(\frac{\boldsymbol{A}_i^k}{\sqrt{d}}\Biggr) \boldsymbol{V}_i^k
\end{array}
}

ここで、$\sqrt{d}$はスケーリングパラメータです。

NATTEN

NATTENは、C++とCUDAを使用して、NAのQK+BとAVの計算、およびそれらの逆伝播における勾配計算を、メモリ消費量を削減しつつ高速化するためのライブラリです。このライブラリは、半精度やフル精度、1次元から3次元までの実装を提供しています。NATTENには、実装のみを目的としたnaive NAと、空間をタイル状に分割し、クエリタイルとキータイルを共有メモリに確保することで高速化を図ったtiled NAがあります。以下の図は、Swin Transformerを基準にして、PyTorch-based NA、naive NA、およびtiled NAの相対メモリ使用量と相対計算速度を示しています。PyTorch-based NAは、特徴マップサイズが大きくなるにつれて、メモリ使用量と計算速度が著しく低下し、448x448以上ではメモリ不足のため測定できていません。naive NAは、小さな特徴マップサイズではSwin Transformerよりも計算速度が速いものの、サイズが大きくなると速度が低下します。一方、tiled NAは、どのサイズに対しても速度が速くメモリ使用量も少ないです。しかし、tiled NAは、2次元のQK計算にのみ使用され、各ヘッドにつき32次元であり、ウィンドウサイズが3、5、7、9、11、13に限定されているという制限があります。NATTENは現在も開発が進行中であり、本論文ではnaive NAとtiled NAが紹介されましたが、最新の論文では、制限を緩和し、高速に動作するGEMM NAやfused NAが提案されています。

Neighborhood Attention Transformer(NAT)

NATの全体的なアーキテクチャを以下に示します。

NATは、4つのステージで構成されています。最初のステージでは、2つの3x3Conv(ストライド=2)を用いて入力を埋め込み、入力サイズを4分の1に縮小します。その後、複数のNAT Blockが続きます。NAT Blockは、通常のTransformer BlockにおけるマルチヘッドアテンションをNAに変更したものです。2ステージから4ステージまでは、3x3Conv(ストライド=2)によって空間サイズを半分に減らし、チャンネル次元を2倍にします。その後、複数のNAT Blockが続きます。

計算複雑度

特徴マップ$h\times w\times d$で、$k\times k$ウィンドウサイズを用いたNAの計算複雑度は、以下の式で表されます。これはSwin Transformerの計算複雑度と同じです。

\displaylines{
&\Omega(\text{NA}) = 4hwd^2 + 2hwdk^2
}

実験

NATの有効性を検証するため、画像分類、物体検出、インスタンス/セマンティックセグメンテーションの画像認識タスクを行い、アブレーション解析も実施しています。
全てのタスクで、NATはConvNeXtとSwinと比較しています。画像分類では、NATは同等のパラメーター数とFLOPsでありながら、精度において両者を上回ります(下図)。物体検出とインスタンスセグメンテーションでは、同等または上回る性能を示し、セマンティックセグメンテーションでは、同等またはわずかに劣る場合もあれば、上回る場合もあります。

以下の表は、Swin-Tinyのアテンションブロックを、SASAとNAに置き換えたときの性能比較を示しています。NAに置き換えた場合が最も高い性能を発揮します。

以下の表は、NATモデルの効果検証のために、モデルを段階的にNATに近づけたときの性能を示しています。最初の変更では、1行目のパッチ埋め込みとパッチマージを、2行目のNATで使用される重なり合う畳み込み設計に置き換えています。これにより、精度が向上します。次に3行目で、モデルを深く薄くすることで、パラメータ数とFLOPsを削減しています。さらに、4行目と5行目では、Swin TransformerのアテンションパターンをそれぞれSASAとNAに変更しています。SASAでは精度が下がりますが、NAでは精度が向上します。

以下の表は、カーネルサイズの影響を分析するために、3x3から9x9までのカーネルサイズを試した実験結果を示しています。7x7を選んだ場合が最も高い精度を示しています。

結論

この論文では、画像認識のための新しいスライディングウィンドウアテンション機構であるNAを提案しています。NAは、線形計算量で動作し、translational equivarianceを保持し、ウィンドウサイズが大きくなるにつれてグローバルアテンションへ漸近し、境界近傍のアテンション範囲を制限しないという特徴を持ちます。さらに、NATTENの開発により、NAの高速化とメモリ使用量削減を実現しています。NAをベースに開発したNATは、Swin TransfromerとConvNeXtの両方よりも同等かそれらを凌駕する性能を達成しています。