【論文5分まとめ】Swin Transformer V2: Scaling up capacity and resolution

この記事は、論文の内容を5分くらいで読めるようにまとめた記事です。そのため、前提となる知識や関連研究に関する説明は大幅に省略しています。
基本的には筆者の備忘録ですが、面白そうと思ったら是非ご自身でも読んでみてください。

概要

Swin Transformerをスケールアップした際に生じる問題を特定し、改善したバージョン2にあたるSwin Transformer V2を提案している。

書誌情報

• Liu, Ze, et al. "Swin Transformer V2: Scaling Up Capacity and Resolution." arXiv preprint arXiv:2111.09883 (2021).
• https://arxiv.org/abs/2111.09883

ポイント

Swin Transformer V2のV1からの変更点は、大きく分けて以下の2つに分けられる。

• ネットワーク構造の変更
• 大規模モデルのための訓練方法の変更

ネットワーク構造の変更

ネットワーク構造の変更は、まとめると下図のようになる。

Normalizationの位置変更

Swin Transformer（下図の-Pre）のサイズを大きくしていくと、深いレイヤーでのactivationが非常に大きくなってしまい、訓練が困難になってしまう。
この問題を回避するために、Layer Normalizationの位置をAttentionの前・MLPの前から、Attentionの後・MLPの後へと変更する。これによって、深い層でのactivationの大きさが控えめに抑えられる（下図の-Post）。

類似度計算の改善

Swin TransformerにおけるAttentionは、以下のように表される。

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

ここでは、パッチ$Q, K$ペアの類似度を内積によって得ている。
著者らは、いくつかのブロックにおけるAttentionマップが、支配的なピクセルに偏ってしまうということを発見した。そこで、この問題を解決するために、$Q, K$ペアの類似度を以下の式のように内積からCosineへと変更する。

$$\operatorname{Sim}\left(\mathbf{q}_{i}, \mathbf{k}_{j}\right)=\cos \left(\mathbf{q}_{i}, \mathbf{k}_{j}\right) / \tau+B_{i j}$$

Cosineを採用することにより、値域を一定の範囲内に抑えられるため、Attentionマップに極端な偏りを生じさせないようにできる。スケールパラメータ$\tau$は0.01よりも大きい値が使用される。

Relative Position Biasの改善

Swin TransformerにおけるAttentionは、以下のように表される（再掲）。

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

V1ではバイアス$B$は、Window内のパッチ間の相対的に位置に応じたパラメータ$\hat{B}\in\mathbb{R}^{(2 M-1) \times(2 M-1)}$によって与えられる。

Swin Transformerをスケールさせるためには、Windowサイズが大きくなっても安定して訓練ができる必要がある。
下表の最初の列はImageNet-1Kの教師データでモデルををスクラッチ学習した時の精度、2列目以降の4列は、そこからWindowの大きさと画像の解像度を大きくしていった時のファインチューニングなしの精度、ファインチューニングありの精度を表す。

Position Biasをパッチ間の相対位置に応じた個別のパラメータで表現していた場合（Parameterized position bias）、Windowが大きくなると、新しいバイアスパラメータを導入する必要があり、ファインチューニングなしでは精度は劣化する一方となる。また、ファインチューニングしたとしても、入力画像の解像度の上昇による精度向上は頭打ちとなってしまうということが示されている。
この問題を解決するために、Relative Position Biasを連続関数化するCPB（continuous position bias）を導入し、低解像度での訓練時には生じなかった値域の相対位置に対しても、それなりに妥当なバイアスを与えられるようにしたい。

まず、素朴な方法として$Q, K$の相対位置$(\Delta x, \Delta y)$を入力するとバイアス$B(\Delta x, \Delta y)$を出力する関数$\mathcal{G}(\Delta x, \Delta y)$を訓練することを考える。単純なMLPによって実装でき、推論時はバイアスは毎回計算する必要はなく、事前に計算しておけば良い。
これのように相対位置からバイアスを求められる連続関数の導入をLinear-Spaced CPBと呼ぶことにする。これによって、上表の2行目のように、スケール時の精度劣化をある程度抑えられ、ファインチューニング時も画像とWindowの解像度に応じて精度が向上する。

次に、更なる改善を目指し、以下のように相対位置を対数スケール化し、これらからバイアスを求めるようなLog-Scaled CPBを考案している。

$$\begin{aligned} &\widehat{\Delta x}=\operatorname{sign}(x) \cdot \log (1+|\Delta x|) \\ &\widehat{\Delta y}=\operatorname{sign}(y) \cdot \log (1+|\Delta y|) \end{aligned}$$

これにより、Windowサイズが大きくなることで生じる相対位置の値域の急激な変化をある程度抑えることができると期待されている。Log-Scaled CPBにより、ファインチューニングなしの時の精度劣化は最も抑えられ、ファインチューニングした時の精度向上も最も大きいことが確認されている。

大規模モデルのための訓練方法の変更

大規模なモデルの訓練のための工夫については、私個人の関心がやや薄い領域なので、さらっとまとめます。

GPUメモリ削減の工夫

V2の一番大きいモデルでは、30億ものパラメータを持つ事になる。
30億ものパラメータをもつモデルをナイーブに訓練しようとすると、最新のGPUでも、メモリが足りなくなる。そこで、本研究では以下のような工夫をおこなっている。

• Zero-Redundancy Optimizer (ZeRO)によるGPUメモリ消費量の削減
• Activation check-pointingによるGPUメモリ消費量の削減
• 高解像度画像を対象としたときに、self-attentionの計算部分で使用メモリが極端に大きくなってしまう問題を発見。self-attention部分だけバッチ処理でなく、シーケンシャルに処理することでメモリ使用量の増大に対処。

以上の工夫で、このようなモデルをNvidia A100-40G GPU上で、以下のような高解像度の訓練ができた。

• 入力解像度1536x1536のCOCO物体検出モデル
• 320x320x8のKinetics-400動作分類モデル

自己教師あり学習の導入

一般に、巨大なモデルの訓練には大規模なデータが必要になる。しかし、Swin Transformer V2の最も大きいバージョンであるSwinV2-Gに関しては、ノイジーなImageNet-22Kを5倍にデータ拡張し（7000万枚）、Masked Image Modeling手法のSimMIMによる自己教師あり学習も併用している。これによって、JFT-3Bほど巨大なデータセットは使用せずに、訓練することに成功している。

SimMIMは、サムネイルを見た限りでは以前まとめたMAEに近そうな手法のようですが、未確認です。
ImageNet-22Kを5倍にする方法についての詳細はよくわかりませんが、github上で質問が行われているので、いずれ明らかになることでしょう。

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

Swin Transformer V2は、V1と同様に、ベースとなるチャネル数とSwin Transformer Blockの数によってスケールできる。これに加え、入力画像解像度とWindowサイズを大きくすることができる。

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

V2では、V1よりもさらに大きなバリエーションを用意している。以下の2つに関しては、activationの爆発が生じやすいので、6レイヤーごとにLayer Normをメインブランチにも導入している。

• SwinV2-H: $C=352$, layer numbers $=\{2,2,18,2\}$
• SwinV2-G: $C=512$, layer numbers $=\{2,2,42,2\}$

実験

以下の4種類のタスクについて実験をしている。これらのタスクでstate-of-the-artを更新している。

• 画像分類：ImageNet-22Kで事前学習し、ImageNet-1K V1/V2でファインチューニングしてTop1精度を評価している。上述の通り、大きいモデルはSimMIMによる自己教師あり学習で事前学習している。
• 物体検出：COCOで評価しているが、大きいモデルではImageNet-22Kでの事前学習に加え、Object365も事前学習に使用している。
• セマンティックセグメンテーション：ADE20Kで評価
• 動画の行動認識：Kinetics-400で評価