論文とコード

Xiaoyi Dong, Jianmin Bao, Dongdong Chen, Weiming Zhang, Nenghai Yu, Lu Yuan, Dong Chen, Baining Guo, CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows, arXiv:2107.00652, 2021
コード(https://github.com/microsoft/CSWin-Transformer)

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

概要

論文では、一般的な視覚認識タスクに適用できる汎用的なTransformerのバックボーンとして、CSWin Transformerを提案しています。Transformerのバックボーンを設計する際には、アテンションの範囲を適切に設定することが重要です。アテンションの範囲を広くすると、計算量が多くなり、狭くすると、モデルの性能が低下する可能性があります。この問題に対処するために、CSWin Transformerでは、特徴マップを等幅の水平ストライプと垂直ストライプに分割し、それぞれ異なるストライプ状のアテンションを、1つのセルフアテンション機構内で実行します。これにより、計算量を各ストライプの範囲内から増やすことなく、アテンションの範囲を、各ストライプを組み合わせた十字形のウィンドウ(Cross-Shaped Window)に拡大することができます。加えて、各レイヤーごとにストライプ幅を調整することで、アテンションの範囲をさらに広げ、わずかな計算負荷で、高い精度を維持できるモデルを構築することができます。また、CSWin Transformerは、既存の位置エンコーディングよりも局所的な位置情報をより良く処理できる、Locally-enhanced Positional Encoding (LePE)を新しく導入しています。LePEは、任意の入力画像に対応しているため、物体検出やセマンティックなどの下流タスクに効果的です。CSWin Transformerは、これらの設計と階層的な構造を組み込むことで、画像分類、物体検出、セマンティックセグメンテーションタスクにおいて、以前の最先端であったSwin Transformerよりも優れたパフォーマンスを発揮することができます。

全体的なアーキテクチャ

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

入力画像をトークン化するために、Convolutional Token Embeddingと呼ばれる重複する畳み込み層(カーネル7×7、ストライド4)を用います。この畳み込み層により、入力画像は$H/4\times W/4\times C$の特徴マップに変換されます。階層的な表現を生成するために、ネットワーク全体は4つのステージで構成されています。2つの隣接するステージの間には、トークンの数を減らし、チャネルの次元を2倍にするための、畳み込み層(カーネル$3\times 3$、ストライド$2$)が用いられます。したがって、特徴マップは、$i$番目のステージで$H/2^{i+1}\times W/2^{i+1}\times 2^{i-1}C$となります。各ステージは、$N_i$個の連続したCSWin Transformer Blockで構成されており、トークンの数を維持します。CSWin Transformer Blockは、通常のTransformer Blockと全体的に類似した構造を持っていますが、2つの違いがあります。(1) セルフアテンションの機構をCross-Shaped Window Self-Attentionに置き換えています。(2) 局所的な特徴を捉えるために、Locally-enhanced Positional Encodingをセルフアテンション操作に追加します。

Cross-Shaped Window Self-Attention

CSWinのセルフアテンションは、特徴マップを水平方向と垂直方向のストライプ状に分割し、各ストライプのセルフアテンションを1つのTransformer Block内で計算します。具体的には、マルチヘッドアテンションのヘッド$\lbrace h_1,\cdots, h_k \rbrace$を水平ストライプと垂直ストライプの2つのグループに分割し、各グループのヘッドにストライプのセルフアテンションを適用します(下図左)。

マルチヘッドのグループ化による処理は、アテンションの対象を全体に拡大し、より効率的にセルフアテンションを実現する方法です。これは、従来のセルフアテンションとは異なります。従来は、マルチヘッド全体に対して同じアテンション操作を適用するか、異なるアテンション操作を順次実行します(図参照)。実験結果によると、マルチヘッドのグループ化による処理は、従来の方法と比較して、精度が向上することが示されています。
上図のCriss-CrossとSequential Axialは、CSwinと密接に関連する先行研究です。しかし、これらの先行研究は、ストライプ幅が固定されているという制約があります。ストライプ幅を大きくすると、計算量は増加しますが、アテンションの範囲が広がり、精度が向上することが期待されます。CSwinは、ストライプ幅を任意に変更できるため、これらの先行研究をより汎用性の高いものに拡張しています。
ストライプ幅の調整は、モデルの性能と効率に大きな影響を与えるため、ストライプ幅を増やすと計算量がどれだけ増加するか把握しておくことが重要です。そこで、CSWinのセルフアテンションの計算複雑度を見ると、次の式で表されます。

\displaylines{
\Omega(\mathrm{CSWin}) = HWC * (4C + sw*H + sw*W)
}

ここで、$sw$はストライプ幅を表します。高解像度の入力の場合、初期ステージでは$H$と$W$が$C$よりも大きく、後期ステージでは$C$よりも小さくなる傾向にあります。このため、初期ステージでは小さな$sw$を、後期ステージでは大きな$sw$を選択することで、計算量を抑えながら、効果的にアテンションの範囲を拡大することができます。デフォルトでは、4ステージに対して$sw$を1、2、7、7に設定しています。実験結果によると、この設定は精度と計算量のバランスが取れていることが報告されています。また、この設定では、画像分類タスクでよく用いられる$224\times 224$の入力画像に対して、中間特徴マップのサイズが$sw$で割り切れるようにしています。

Locally-enhanced Positional Encoding

セルフアテンションは、入力データの順序に依存しないため、入力データの位置情報を無視します。この情報を追加するために、CSWin Transformerでは、局所的な位置情報を強化する手法として、locally-enhanced positional encoding (LePE)を導入しています。LePEは、Vを直接操作し、セルフアテンションモジュールと並行して動作します。式で書くと以下のようになります。

\displaylines{
\mathrm{Attention}(Q,K,V) = \mathrm{SoftMax}(QK^T/\sqrt{d})V+\mathrm{DWConv}(V)
}

ここで、$\mathrm{DWConv}$は、depth-wise convolutionであり、位置情報に関するバイアスを表します。3x3のカーネルを使用するため、入力要素の周辺情報に基づいて位置情報を生成します。
Transformer内で位置情報を追加する方法として、LePEの他に、Swin Transformerで用いられたrelative positional encoding(RPE)が知られています。RPEは、アテンションの計算時に追加される点がLePEと異なります。

\displaylines{
\mathrm{Attention}(Q,K,V) = \mathrm{SoftMax}(QK^T/\sqrt{d}+B)V
}

RPEでは、位置情報のバイアス$B$は学習後に固定されるため、テスト時の異なる入力に対して、局所的な情報の利用が制限される可能性があります。一方、LePEでは、異なる入力$V$に対して、depth-wise convolutionを通して、局所的な情報を取り入れます。この差により、LePEはRPEよりも精度が向上する可能性があると考えられます。
LePEと同様に、conditional positional encoding (CPE)も局所的な情報に基づいて位置エンコーディングを生成する方法です。しかし、CPEは、LePEとは異なり、Transformerに入力される前の特徴マップに対して位置情報を付与します。また、幅広く使用されているabsolute positional encoding (APE)も、Transformerに入力する前の特徴マップに対して位置情報を付与します。オリジナルのTransformerでは、APEは、多数の異なる周波数の正弦波を用いて生成されます。以下の図は、これらの位置エンコーディングの機構の違いを示しています。

位置エンコーディングには様々な方法が提案されていますが、LePEは、物体検出やセグメンテーションの下流タスクにおいて、高い精度を示す実験結果が報告されています。

実験

CSWin Transformerが汎用的な画像認識バックボーンとしての有効性を示すために、ImageNet-1Kによるクラス分類、COCOによる物体検出、ADE20Kによるセマンティックセグメンテーションの実験を行っています。また、CSWin Transformerの各コンポーネントを分析するために、包括的なアブレーションスタディも実施しています。ここではその一部を紹介します。

ImageNet-1Kによるクラス分類

次表は、クラス分類タスクにおけるCSWin Transformerと最先端のCNNおよびTransformerの性能を比較したものです。

CSWin Transformerは、入力画像のサイズに関わらず、従来の最先端のTransformerよりも大幅に性能が向上します。最先端のCNNと比較すると、CSWin Transformerは、smallとbaseの設定でEfficientNetと同等かそれ以上の結果を達成しながら、計算複雑度を削減することができます。

CSWinのセルフアテンション

次表は、CSWinのセルフアテンションが、特に下流タスクにおいて効率的かつ効果的であることを示しています。

CSWinのストライプ幅を$sw=1$として、従来の手法であるAxial-Attentionと比較すると、CSWinはImageNetではAxialよりもわずかに精度が優れており、下流タスクでは大きく上回っています。$sw=2$のCSWinはCriss-Cross Attentionよりもわずかに優れていますが、CSWinの計算速度は異なるタスクで2倍から5倍速くなっています。マルチヘッドをグループ化せず、水平ストライプと垂直ストライプのアテンションを逐次的に処理すると、CSWinの性能はすべてのタスクで低下します。このことは、グループ設計がはるかに効果的であることを示唆しています。

動的なストライプ幅

次図は、ストライプ幅と精度のトレードオフを示しています。

ストライプ幅を大きくすると、計算コストは増加しますが、Top-1の分類精度は最初は大きく向上します。しかし、ストライプ幅が十分に大きくすると、精度は低下します。デフォルト設定である[1,2,7,7]は、精度と計算コストの間で良好なバランスを実現しています。

位置エンコーディングの比較

次表は、LePEと他の位置エンコーディング手法(APE、CPE、RPE)の比較結果を示しています。また、位置エンコーディングを使用しない場合(No PE)と、すべてのTransformerブロックの前にCPEを適用した場合(CPE*)の結果も示しています。

LePEは、位置エンコーディングなしの場合と比較して、精度が向上します。分類タスクでは、RPEと同等の性能ですが、物体検出とセグメンテーションのタスクでは、他の位置エンコーディングよりも優れた性能を発揮します。

結論

CSWin Transformerは、効率的にアテンション範囲を拡大し、下流タスクに効果的な局所強化位置エンコーディングを導入した、新しいVision Transformerです。アテンション範囲を効率的に拡大するために、2つの工夫をしています。1つは、マルチヘッドをグループに分割し、各グループに水平方向と垂直方向の異なるストライプのアテンションを割り当たことです。もう1つは、ネットワークの深さ方向にストライプ幅を大きくして、わずかな計算コストでアテンションをさらに拡大したことです。これらの新しい手法を追加したことで、CSWin Transformerは最先端のパフォーマンスを達成しています。