Swin Transformer

Swin Transformer（Shifted Window Transformerの略）は、小さな局所ウィンドウ内で自己アテンションを計算し、連続する層の間でウィンドウをシフトさせる階層型ビジョントランスフォーマーです。Microsoft Research AsiaのZe Liu氏らが2021年の論文Swin Transformer: Hierarchical Vision Transformer using Shifted Windowsで発表し、ICCV 2021の最優秀論文に贈られるMarr Prizeを受賞しました。局所ウィンドウにより、アテンションのコストはピクセル数に対して二次ではなく線形に増加します。また、パッチを段階的に統合することで、畳み込みニューラルネットワークのようなマルチスケールの特徴マップを生成します。これらの特長により、Transformerを汎用的なバックボーンとして、高解像度画像の分類だけでなく、物体検出やセグメンテーションにも実用化できました。

シフトウィンドウの仕組み#

Swinはまず、画像を4×4ピクセルのパッチに分割し、それぞれをtoken埋め込みに射影します。そのため、224×224画像はTinyモデルでチャネル数96の56×56トークングリッドになります。3,136個すべてのトークンに対するグローバルアテンションでは、各トークンをほかのすべてのトークンと比較する必要があります。Swinではグリッドを重複しない7×7トークンのウィンドウに分割し、各ウィンドウ内だけでアテンションを計算します。論文では、h×wグリッドとCチャネルの場合、グローバルアテンションの計算量を4hwC² + 2(hw)²C、M×Mトークンのウィンドウの場合を4hwC² + 2M²hwCとしています。後者はhwに対して線形です。

ウィンドウだけでは情報がやり取りされないため、ブロックを1つおきにウィンドウグリッドを半分の幅、つまり3トークン分シフトします。これにより、新しいウィンドウが以前の境界をまたぎます。ウィンドウ数を一定に保つため、実装では特徴マップを循環シフトし、元の画像で隣接していなかったトークン間のアテンションにマスクを適用します。各ブロックでは、学習可能な相対位置バイアスもアテンションスコアに加算します。

ステージ間では、パッチ統合層が隣接するトークンの各2×2グループを連結し、線形射影を行います。これにより解像度が半分になり、チャネル数が2倍になります。4つのステージは、入力解像度の1/4、1/8、1/16、1/32の特徴マップを出力します。これはResNetと同じストライドなので、Swinは特徴ピラミッドネットワークなど、既存の検出・セグメンテーションヘッドにそのまま組み込めます。torchvisionの実装では、両方の仕組みを確認できます。

import torch
from torchvision.models import swin_t

# Swin-T architecture; random weights are enough to inspect its structure
model = swin_t().eval()

# Each pair of blocks uses 7x7 windows, the second shifted by 3 patches
regular, shifted = model.features[1]
print(f"Window {regular.attn.window_size}, shift {regular.attn.shift_size} then {shifted.attn.shift_size}")

# features alternates stages with patch merging; stage outputs are (batch, height, width, channels)
x = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    for i, layer in enumerate(model.features):
        x = layer(x)
        if i % 2 == 1:
            print(f"Stage {i // 2 + 1}: {tuple(x.shape)}")

Window [7, 7], shift [0, 0] then [3, 3] と出力された後、ステージ出力として(1, 56, 56, 96)、(1, 28, 28, 192)、(1, 14, 14, 384)、(1, 7, 7, 768)が表示されます。各ステージでグリッドの幅と高さが半分になり、チャネル数が2倍になるため、密な予測ヘッドが必要とするピラミッドが得られます。

Swin TransformerとViT、CNNの比較#

初代ViTは画像を16×16のパッチに分割し、すべての層で単一の解像度を維持してグローバルアテンションを行います。画像分類には適していますが、検出やセグメンテーションで使う高解像度入力ではコストが高くなります。Swinは、局所性と特徴の階層構造という畳み込みネットワークの2つの事前バイアスを再導入しつつ、入力に応じて重みが変わるアテンション機構を維持します。Swinの局所ウィンドウは、各トークンがそれぞれの近傍に注目するスライディングウィンドウアテンションとは異なります。Swinのウィンドウはトークン間で共有される固定の区画であり、GPUの行列積に効率よく対応します。

公開時点で最大のモデルであるSwin-Lは、ImageNet-22Kで事前学習され、ImageNet-1Kでトップ1精度87.3%、COCO test-devでボックスAP 58.7およびマスクAP 51.1、ADE20KでmIoU 53.5を達成しました。物体検出とセマンティックセグメンテーションの各ベンチマークにおいて、公開済みの最高記録でした。その成功は畳み込みネットワークの研究も前進させました。MetaのConvNeXt（2022）は、Swinから取り入れた設計を用いてResNetを現代化し、同等以上の精度を達成しました。

Swinは、カスタムUltralyticsモデルのバックボーンとして利用できます。モデルYAMLファイル内のTorchVisionモジュールは、事前学習済みImageNet重みを使ってtorchvisionのswin_tを読み込みます。YOLO検出ヘッドは、返されるストライド32の単一特徴マップを使って学習します。

派生モデルと用途#

• Swin Transformer V2：2022年に設計を30億パラメーター、最大1,536×1,536ピクセルの画像へ拡張しました。学習を安定させるため、スケーリングしたコサインアテンションと、対数間隔の連続位置バイアスを使用します。
• 動画アクション認識：Video Swin Transformerはウィンドウを空間と時間の3次元に拡張し、動画クリップ内の動作を分類します。
• 超解像と画像復元：SwinIRはSwinブロックを使って、画像の拡大、ノイズ除去、JPEGアーティファクトの除去を行います。
• 医用画像解析：Swin-UnetとSwin UNETRは、SwinエンコーダーとU-Net形式のデコーダーを組み合わせます。Swin UNETRは、3D MRIボリューム内の脳腫瘍をセグメンテーションします。