【Swin Transformer】今こそ押さえたいTransformer系画像認識モデル

画像認識 2024年03月22日

3つの要点
✔️ 近年コンピュータビジョンの研究でベースラインとしてよく用いられているSwin Transformerを解説
✔️ すべてのパッチと関連性(Attention)を計算するVision Transformerとは異なり，近傍のパッチをまとめたwindow内でAttentionを計算する
✔️ 異なるパッチサイズでAttentionの計算を行うため，様々なスケールの特徴が得られる

Swin Transformer: Hierarchical Vision Transformer using Shifted Windows
written by Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, Baining Guo
(Submitted on 25 Mar 2021 (v1), last revised 17 Aug 2021 (this version, v2))
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV); Machine Learning (cs.LG)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

Swin Transformerは2021年に発表されたTransformerベースの画像認識モデルです。近年のコンピュータビジョンの研究でベースラインとしてよく用いられているため，改めてどのようなモデルなのかを紹介したいと思います。

Swin Transformer以前のTransformerベースの画像認識モデルにVision Transformer(ViT)があります。画像を16×16の単語で構成される文章として捉え，自然言語処理で用いられていたTransformerを画像認識に初めて応用したものになります。

本記事で紹介する論文では文章と画像の違いを指摘し，ViTをより画像ドメインに適合させたSwin Transformerを提案しています。

文章と画像の違いとして以下の2点を挙げています。

画像中の視覚的要素は，単語トークンとは異なり，スケールが大きく変化する
文書中の単語よりも画像中のピクセルのほうが解像度が高い(情報量が多い)

これらの違いに対応するため，

異なるパッチサイズにおいてAttentionを計算する
より小さなパッチサイズでAttentionを計算する

ことを提案しています。

以下の図はこれらの点における，ViTとSwin Transformerの違いを表現しています。

小さいパッチサイズでAttentionを計算すると，きめ細かい特徴が取れる一方で，計算量が増大するという課題があります。

そこでSwin TransformerではShifted Window based Self-Attentionが導入されています。いくつかのパッチを1つのwindowとしてまとめ，そのwindow内のみでAttention計算を行うことで計算量を削減しています。

次のセクションでは，Swin Transformerの全体像から，Shifted Window based Self-Attentionを含めた細かな部分について見ていきます。

Swin Transformer

全体像

こちらがSwin Transformerの全体像となります。

はじめに入力された画像に対してPatch Partitionを行います。

Patch Pertition：4×4のピクセルを1つのパッチとして分割します。ViTでは16×16のピクセルを1つのパッチとしているため，Swin Transformerではよりきめ細かい特徴が抽出できます。

次にLinear Embeddingを行います。

Linear Embedding：パッチ(4×4×3ch)をC次元のトークンに変換します。Cはモデルのサイズによって異なります。

各パッチから得られたトークンに対して，Swin Transformer BlockでAttentionを計算し，特徴抽出を行います。Swin Transformer Blockの構成を以下に示します。

Swin Transformer Block：通常のTransformer Blockで採用されているMulti-head Self-Attention (MSA)をShifted Window based Self-Attention(W-MSAとSW-MSA)に置き換えています。次のセクションで詳しく説明します。それ以外は通常のTransformerとほとんど同じ構成になっています。

ここまでみてきたLinear EmbeddingとTransformer Blockの部分はStage1と呼ばれます。Stageは1から4までありますが，各Stageのパッチサイズが異なるため，様々なスケールでの特徴抽出が可能です。パッチサイズが異なるのはPatch Mergingにより，近傍のパッチを集約しているためです。

Patch Merging：各Stageにおいて，近傍(2×2)のパッチ(トークン)をまとめて1つのトークンとします。具体的には2×2のトークンを結合し，得られた4C次元のベクトルを線形層により2C次元とします。例えばStage2では，(H/4)×(W/4)×C次元から(H/8)×(W/8)×2C次元となります。

次の章ではSwin Transformer BlockでのAttention計算について詳しく見ていきます。

Shifted Window based Self-Attention

通常のTransformerとSwin Transformer BlockのAttention計算の違いについて計算量の観点から説明します。

通常のTransformerではすべてのトークン間のAttentionの計算を行います。h, wを画像における縦，横のパッチ数としたとき，計算量は以下の式のようになります。

一方，Swin Transformerでは複数のパッチで構成されるwindow内でのみAttentionの計算を行います。1つのwindowにはM×Mのパッチが含まれており，基本的にはM=7で固定されています。計算量は以下の式のようになります。

通常のTransformerではパッチ数(hw)の2乗に比例して計算量が増えます。一方，Swin TransformerではMの2乗に比例して計算量が増大しますが，M=7のため影響が小さく，パッチ数(hw)に対してもその増大は1乗オーダーに抑えられます。このため，Swin Transformerでは小さなパッチサイズでのAttention計算が可能となっています。

次にwindowへの分割方法について説明します。M×Mで画像を均等に分割するようにwindowが配置されています。このように配置された各windowにごとAttentionを計算するため，隣合うパッチだったとしても異なるwindowであればAttention計算がされないことになります。このwindowの境界の問題を解決するため，1度目のAttentionを計算（W-MSA : Window based Multi-Head Self-Attention）した後，windowをshiftさせ，もう一度Attentionを計算すること (SW-MSA : Shifted Window based Multi-Head Self-Attention)を提案しています。

以下の図のように，元のwindow分割から，([M/2], [M/2])ピクセルだけshiftさせます。

Efficient batch computation for shifted configuration

SW-MSAにおけるwindowはサイズが異なり，windowの数も増大します。そのため，愚直に処理を行うとW-MSAと比べて計算量が増大するという問題生じてしまいます。そこでSW-MSAでは実際にwindowの配置を変更するのではなく，cyclic shiftと呼ばれる方法を用いた疑似的な操作を行っています。

以下の図に示すように，画像全体を左上方向にshiftし，はみ出した部分を空いた部分に入れ込む処理 (cyclic shift) を行います。これを行うことで，W-MSAのwindow内でのAttention計算と同じように計算することができます。また，window内に隣接しないパッチが含まれる場合があるため，その部分はmask処理を行っています。最後の出力ではcyclic shiftと逆の操作 (reverse cyclic shift) を行うことで，パッチを元の位置に戻します。

Architecture Variants

Swin TransformerにはT, S, B, Lのサイズがあり，各ステージにおける次元数(dim)，ヘッド数(head)，ブロック数が以下の表のように異なります。

実験

画像認識タスクのImageNet-1K，物体検出タスクのCOCO，セマンティックセグメンテーションタスクのADE20Kで他のモデルとの比較を行っており，いずれも最高精度を達成しています。（実験結果の詳細は論文4章のTable1~3に記載があります。）

SW-MSAにおけるAblation Studyが行われており，いずれのタスクにおいてもW-MSAだけでなく，SW-MSAを導入したときのほうが精度が高いことが確認されています。

まとめ

本記事ではコンピュータビジョンの研究でベースラインとして使われているSwin Transformerについて紹介しました。

すべてのパッチ間でAttentionを計算するViTとは異なり，近傍のパッチをまとめたwindow内でのAttention計算とパッチの集約を繰り返すことで，様々なスケールでの特徴を取ることができます。また，すべてのパッチ間でAttentionを計算しないため，計算量が抑えられ，より小さなパッチサイズから特徴抽出が行える点も魅力だと言えます。

AIを最近学び始めたという方の理解の助けになれば幸いです。