Zenn
🐶

【論文5分まとめ】Swin Transformer

2022/01/26に公開
ディープラーニング
tech

概要

CNNのような階層的なネットワーク構造により、階層的な表現の獲得を可能にしたVision TransformerであるSwin Transformerを提案。Self-Attentionの範囲を固定サイズのWindow内に限定することで、画像サイズに対して線形の計算量で収まるような高効率なネットワーク構造を実現した。画像分類、物体検出、セマンティックセグメンテーションなどのさまざまなVisionタスクに適応可能なCNNのような柔軟性と、最先端の精度を実現している。

書誌情報

  • Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo; Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2021, pp. 10012-10022
  • 公式実装(https://github.com/microsoft/Swin-Transformer)
  • ICCV2021 BestPaper
  • Microsoft Research Asiaの研究

ポイント

全体像

まずは全体像を確認する。下図の(a)に全体像が示されている。

  • 入力画像を重なりの無いパッチに分割する。RGBのチャネル毎に4 \times 4ピクセルのパッチを得るので、まとめると48チャネルになる。Transformerへの入力長は\frac{H}{4}\times\frac{W}{4}になる。
  • 最初のステージでは、48チャネルをCチャネルへと変換するLinear Embeddingが使用される。
  • 階層的な処理を行うために、各ステージでPatch Merging処理が行われる。ここでは、2 \times 2のパッチが1つにまとめられ、一度チャネル数が4倍になった後に、Linear層によってチャネル数を半分にする。つまり、入力チャネル数の2倍のチャネル数がPatch Mergingの出力になる。これにより、ステージが進むたびにより広い領域の情報をマージしつつ、高次の特徴抽出処理が行われる。
  • 各ステージには複数のSwin TransformerBlockが含まれる。

Swin Transformer Block

次に、Swin Transformer Blockを確認する。

Swin Transformer BlockではMulti-head self attention(MSA)が使用されるが、重なりのWindow単位で区切られて適用される。Windowは標準的な位置で取る場合と、位置を半分シフトさせたShifted Windowの場合があり、Block毎に交互に適用される。これらは、W-MSAとSW-MSAと表記される（上図(b)）。Shited Windowによって、異なるWindow間の情報が効率的に混合されることが期待されている。

Swin Transformer Blockでは、LayerNorm, (S)W-MSAが適用されたのち、LayerNorm、MLPが適用される。MLPは2層のLinear層からなり、中間層のチャネル数は\alpha=4倍され、活性化関数としてGeLUが使用される。

W-MSAによるスケーラビリティ

ある時点での特徴マップにh \times w個のパッチが含まれているとする。このとき、ナイーブなMSAとW-MSAの計算の複雑性を比較すると以下のようになる。

\begin{aligned} &\Omega(\mathrm{MSA})=4 h w C^{2}+2(h w)^{2} C \\ &\Omega(\text{W-MSA})=4 h w C^{2}+2 M^{2} h w C \end{aligned}

通常のMSAはパッチ数hwの2次関数になるため入力特徴マップの解像度が大きくなると、計算の複雑さが爆発的に上昇してしまう。W-MSAでは、WindowサイズMは7であり、M^2は定数とみなされるので、計算の複雑さは解像度に対して線形に増加する。そのため、W-MSAを使用することでSwin Transformerはスケーラブルなアーキテクチャとなっている。

SW-MSAによるWindowのシフト

入力された特徴マップの解像度が8\times8で、64個のパッチからなるとする（下図左）。M=4として、1つのWindowにはM^2=16個のパッチを含む、重なりの無いWindowに分割することができる。

このWindowを\left(\left\lfloor\frac{M}{2}\right\rfloor,\left\lfloor\frac{M}{2}\right\rfloor\right)=(2, 2)シフトさせると、上図右のようになる。この状態では、元のWindowの大きさを保っているのは1つのみで、その他のWindowが8つでき、合計9個のWindowができることになる。この状態でそのまま処理を進めるのはやや複雑な処理となる。

そのため、SW-MSAにおいては、cyclic shiftという工夫を行う。下図のように、Windowをシフトしてはみ出した部分（薄く表示されているA, B, C）をWindowの反対位置へ移動させる。これにより、Windowの数やWindowあたりのパッチ数は一定となる。標準的なMSAでよく使用されるマスクを設定してあげることで、無関係なパッチ間でAttentionが生じないようにできるので、複雑な実装を回避することができる。

こうして得られた特徴マップの移動していた部分を本来の位置へと戻す（reverse cyclic shift）ことで、SW-MSAの処理が完了する。

Relative Position Bias

Window内の2つのパッチ（QとK）間の位置関係によってAttentionの大きさがある程度決まることが望ましいことは、直感的にわかる。そのため、2つのパッチの位置に応じたバイアスBをAttentionの計算に導入する。

\text { Attention }(Q, K, V)=\operatorname{SoftMax}\left(Q K^{T} / \sqrt{d}+B\right) V

Q, K, V \in \mathbb{R}^{M^{2} \times d}（d=32がデフォルト）なのでB\in \mathbb{R}^{M^2 \times M^2}になる。素朴に2つのパッチの絶対的な位置関係に対応するM^2 \times M^2次元のバイアスを訓練することもできるが、より本質的な方法として、パッチ間の相対的な位置関係を考慮したRelative Position Biasを導入する。

2つのパッチの相対的な位置関係のバリエーションはM^2 \times M^2ほど多く無い。例えば、Qが大きさMのWindow の左端、KがWindowの右端にあるとすると、横方向の相対位置はM-1である。逆に、Qが右端、Kが左端にあると、横方向の相対位置は-M+1となる。そのため、横方向の相対位置のバリエーションは2M+1通りあることになる。縦方向にも同様のことが言えるので、Q, Kの相対位置のバリエーションは(2M+1) \times (2M+1)となる。

以上のことから、訓練すべきバイアスを\hat{B} \in \mathbb{R}^{(2 M-1) \times(2 M-1)}とでき、訓練対象のパラメータを減らすことができる上に、精度も向上する。

アーキテクチャのバリエーション

各ステージのSwin Transformer Block数によって、以下のようなバリエーションを用意している。

  • Swin-T: C=96, layer numbers =\{ 2,2,6,2 \}
  • Swin-S: C=96, layer numbers =\{ 2,2,18,2 \}
  • Swin-B: C=128, layer numbers =\{ 2,2,18,2 \}
  • Swin-L: C=198, layer numbers =\{ 2,2,18,2 \}

Swin-T, Swin-Sは、複雑さで言えばResNet-50やResNet-101に対応する

実験

詳細については省略するが、ImageNetで分類、COCO2017で検出、ADE20kでセグメンテーションの実験をおこない、Ablation Studyをおこなっている。

以下の表は、Ablation Studyのひとつとして、SW-MSAやRelative Position Biasの有効性について検証した結果である。

Discussion