AxWin Transformer: 局所と大域を両立するAxWin Attentionによる高性能ビジョントランスフォーマー

論文

• Fangjian Lin, Yizhe Ma, Sitong Wu, Long Yu, Shengwei Tian, AxWin Transformer: A Context-Aware Vision Transformer Backbone with Axial Windows, arXiv:2305.01280, 2023

挿入している画像は、本論文からの引用になります。

概要

トランスフォーマーは、画像分類や物体検出などの視覚タスクで優れた性能を示しています。しかし、高解像度の画像が対象となるダウンストリームタスクでは、グローバルなセルフアテンションによる2乗の計算量が課題となります。この問題を解決する一つの方法は、特徴マップを局所ウィンドウ領域に分割し、各ウィンドウ領域でセルフアテンションを適用することです(下図(b))。この方法は、計算量を削減することができますが、長距離のグローバルなモデリング能力が不足します。一方、軸ベースのセルフアテンション(下図(c))は、ダウンストリームタスクにとってより効果的ですが、局所モデリング能力が不足しています。この問題を解決するために、論文ではAxial Window Self-Attention(AxWin Attention)を提案しています(下図(d))。AxWin Attentionは、局所特徴に焦点を当てるウィンドウ領域と、大域特徴をモデル化する軸方向領域で構成され、より豊かな文脈依存関係を獲得します。さらに、高解像度画像の文脈情報を充実させるために、マルチスケールのパッチ埋め込みによるダウンサンプリングを考案しています。これらの設計を基に構築されたAxWin Transformerは、分類タスクと物体検出およびセグメンテーションタスクで、Pale Transformerなどの最先端の手法を凌駕する性能を示しています。

全体的なアーキテクチャ

AxWin Transformerは、steam層、4つの階層ステージ、クラス分類ヘッドから構成されています(下図(a))。 Stem層は$3\times 3$ convolution(ストライド=2)と2つの$3\times 3$ convolution(ストライド=1)からなります。Stem層の後、各ステージは、ダウンサンプリングを行うMulti-Scale Patch Embedding(MSPE)と複数のAxWin Transformer Blockを適用します。入力画像は、MSPEで画像サイズを半分に縮小して特徴量を2倍にし、AxWin Transformer Blockでは画像サイズを維持します。最終的なクラス分類ヘッドは線形層です。

Multi-Scale Patch Embedding

MSPEは、上図(b)のように、入力特徴マップ$X\in R^{h\times w\times c}$に対して、複数のスケールの特徴を抽出し、それらを融合することで、より豊かな表現を生成する出力特徴マップ$Y\in R^{\frac{h}{2}\times \frac{w}{2}\times 2c}$を得ます。具体的には、MSPEは、4つに分岐し、各分岐枝$i$に対して、$i$個の$3\times 3$ depth-wise convolution(ストライド=2)を適用して、中間特徴マップ$X_i\in R^{\frac{h}{2^i}\times \frac{w}{2^i}\times 2c}$を得ます¹。$i$番目と$i+1$番目の分岐では、スケールが異なる解像度を一致させるために、低解像度特徴マップをバイリニア補完によって拡大します。その後、$3\times 3$ depth-wise convolution(ストライド=1)と$1\times 1$ convolutionを適用し、最後に残差結合によって融合させます。

AxWin Attention

AxWin Attentionは、特徴マップを小さな正方形の領域に分割したウィンドウ領域と、行方向と列方向を組み合わせた軸領域で、それぞれセルフアテンションを適用することで、局所と大域の両方の特徴を捉えるアテンション機構です。具体的には、特徴マップ$X\in R^{h\times w\times c}$が与えられた場合、まず、全結合層を用いて、Query($X_q$)、Key($X_k$)、Value($X_v$)を生成します。次に、上図(d)に示すように、チャンネル次元をウィンドウグループ$\lbrace X_{wq},X_{wk},X_{wv} \rbrace\in R^{h\times w\times \frac{c}{2}}$と軸グループ$\lbrace X_{aq},X_{ak},X_{av} \rbrace\in R^{h\times w\times \frac{c}{2}}$の半分に分割します。ウィンドウグループでは、Swin Transformerの先行研究と同様に、重複しない方法で$\lbrace X_{wq},X_{wk},X_{wv} \rbrace$をウィンドウに分割し、各ウィンドウでセルフアテンションを行います。軸グループでは、Pale Transformerの先行研究を参照して、$\lbrace X_{aq},X_{ak},X_{av} \rbrace$を、チャンネル次元で2分割し、行と列の2つのグループにします。行グループは、等間隔に並んだ行$\lbrace X_{aq}^r,X_{ak}^r,X_{av}^r \rbrace\in R^{s_{ar}\times w\times \frac{c}{4}}$を$\frac{h}{s_{ar}}$個含み、列グループは、等間隔に並んだ列$\lbrace X_{aq}^c,X_{ak}^c,X_{av}^c \rbrace\in R^{h\times s_{ac}\times \frac{c}{4}}$を$\frac{w}{s_{ac}}$個含みます。ここで、$s_{ar}=s_{ac}$であり、$s_{ar}$と$s_{ac}$はそれぞれ行と列の数を表します。分割完了後、チャンネル次元の結合操作を行い、セルフアテンションに適用するための$\lbrace X_{aq}^{rc},X_{ak}^{rc},X_{av}^{rc} \rbrace$を取得します²。最後に、ウィンドウグループと軸グループのセルフアテンションの出力を、チャンネル方向で結合します。

AxWin Transformer Block

AxWin Transformer Blockは、上図(c)に示すように、条件付き位置エンコーディング(CPE)、AxWin Attention、内部畳み込みフィードフォワードネットワーク(ICFFN)の3つの主要モジュールから構成されています。CPEは、入力特徴量に対して動的位置埋情報を付与します。AxWin Attentionは局所的および大域的な文脈情報を取得します。ICFFNは、特徴量変換と局所情報抽出の2つの機能を、MLP(2つの全結合層とその間にGELU活性化関数)と$3\times 3$ depth-wise convolutionという2つの異なる処理を用いて実現しています(上図(e))。

実験

AxWin Transformerの汎化能力を示すため、画像分類、セマンティックセグメンテーション、物体検出、インスタンスセグメンテーションの各タスクで実験を実施しています。また、アブレーションスタディも実施しています。

視覚タスクにおける性能評価

次の表は、AxWin Transformerと既存手法のパフォーマンスを比較しています。表の上から順に、ImageNet-1Kによる画像分類、COCOによる物体検出とインスタンスセグメンテーション、ADE20Kによるセマンティックセグメンテーションの結果を示しています。



AxWin Transformerの変種は、すべてのタスクにおいて、他のTransformerの変種よりも高い精度を達成しています。また、PVTの変種を除き、他の手法よりもFLOPsを削減しています。さらに、パラメータ数は、他の手法とほぼ同等です。AxWin Transformerは、Pale TransformerとSwin Transformerのアイデアを組み合わせることで、両者を上回る精度と計算コストの削減を実現しています。

アブレーションスタディ

次の表は、異なるダウンサンプリング方法で実施したときの性能比較を表しています。MSPEモジュールは、計算量のわずかな増加だけで性能向上をもたらしています。

次の表は、アテンションの対象範囲を、軸上、ウィンドウ上、両方を組み合わせたAxWinに変更した際の性能比較を示しています。AxWinアテンションは、最も優れた結果を達成しています。

次の表は、分割サイズが異なる場合に得られる性能の比較を表しています。特徴マップを分割するときの、ウィンドウサイズ、軸の行サイズ、列サイズは、すべて同じ値で、特徴マップがそのサイズで割り切れない場合は、余った部分をパディングで埋めます。分割サイズが7または12にすると、性能が向上します。そのため、tiny、small、baseモデルの分割サイズは、それぞれ7、7、12に設定します。

結論

AxWin Attentionは、局所的な情報と大域的な情報を同時に考慮する、効率的なセルフアテンションメカニズムです。このメカニズムに基づいて開発されたAxWin Transformerは、従来の手法よりも優れた性能を示す、高性能なビジョントランスフォーマーバックボーンです。

1. チャンネル次元が$C$から$2C$に増加している理由は、論文には記載されていません。depth-wise convolutionはチャンネル次元を変化しないため、チャンネル次元の増加は、$1\times 1$ convolutionなどでチャンネル数を変換したものと考えられます。 ↩

2. $\lbrace X_{aq}^{rc},X_{ak}^{rc},X_{av}^{rc} \rbrace$は、空間次元ではなく、チャンネル次元で連結したトークンです。そのため、1つのトークンには、行と列の両方の情報が含まれています。このトークンをセルフアテンションに適用するアイディアは、Pale TransformerやCSWin Transformerなどの論文には記載されていなかったため、新しい方法であると考えられます。ただし、行グループと列グループを結合する場合、空間次元を同じにする必要があるため、$h s_{ac}=w s_{ar}$の条件が必要になります。 ↩