Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

In plain words

A Transformer looks at every part of its input and lets each part borrow from all the others. That works well for words, but a picture has far more small pieces than a sentence has words, and checking every piece against every other gets very expensive as the picture grows. Images also contain objects of very different sizes, which a fixed-size chunking of the picture handles poorly.

Swin Transformer keeps the attention local. The picture is cut into small patches, and each patch looks only at the patches in its own small square window. The network then merges neighbouring patches as it goes deeper, so the picture is described at several zoom levels, much like the feature maps of a classic convolutional network. To let neighbouring windows talk to each other, every second layer moves the window grid by half a window, so the new windows straddle the borders of the old ones.

The result is a general-purpose image backbone. Its cost grows in a straight line with the picture's size, and its multi-scale output can feed detectors and segmenters as well as classifiers. The authors report better numbers than the previous best on COCO object detection and ADE20K semantic segmentation.

TL;DR

Swin Transformer is a hierarchical vision Transformer for classification, detection and segmentation. Each block computes self-attention only inside non-overlapping M x M windows (M = 7 by default), and every second block shifts the windows by floor(M/2) so information crosses window borders. Patch merging gives four stages at 1/4, 1/8, 1/16 and 1/32 resolution, and attention cost is linear in image size for fixed M. Author-reported: 81.3% ImageNet-1K top-1 for Swin-T at 224 px, 58.7 box AP on COCO test-dev and 53.5 mIoU on ADE20K val for Swin-L. Implement from the code and this card together, not from the paper alone (see Paper vs code).

Problem

Plain Vision Transformers such as ViT use one token scale and global self-attention §1 §3.2. Two gaps follow from that design. First, visual entities vary a lot in scale, and tokens of one fixed size do not fit dense tasks such as detection §1. Second, global self-attention costs grow with the square of the token count, so high-resolution images and per-pixel prediction become intractable §1 §3.2, eq. 1. The paper's premise is that a convolution-like pyramid of feature maps, with attention limited to local windows, fixes both gaps while staying a single architecture for all three tasks §1. The paper also argues that local windows alone block information flow between neighbouring windows, which motivates the shift §3.2.

Method

Overall architecture

An RGB image is split into non-overlapping 4 x 4 patches, each flattened to a 48-dimensional vector and linearly projected to C dimensions §3.1. Four stages follow fig. 3 tab. 7:

1. Stage 1 runs Swin blocks on H/4 x W/4 tokens (56 x 56 for a 224 px input).
2. Each later stage starts with patch merging: a 2 x 2 neighbourhood is concatenated into 4C channels, normalised, and linearly reduced to 2C §3.1 swin_transformer.py. Resolution halves and channels double.
3. Stages 2, 3 and 4 run the same block pattern at strides 8, 16 and 32 §3.1 tab. 7.

Classification takes the final feature map, applies global average pooling and a linear classifier app. A2.1 swin_transformer.py.

Swin block

Each block replaces the standard multi-head self-attention module with a window module, keeping the rest of a Transformer block unchanged. A LayerNorm sits before each attention and MLP module, a two-layer MLP with GELU follows the attention, and each module has a residual connection §3.1 fig. 3b swin_transformer.py swin_transformer.py.

Consecutive blocks alternate between two window layouts §3.2:

\begin{aligned} \hat{z}^{l} &= \text{W-MSA}(\text{LN}(z^{l-1})) + z^{l-1} \\ z^{l} &= \text{MLP}(\text{LN}(\hat{z}^{l})) + \hat{z}^{l} \\ \hat{z}^{l+1} &= \text{SW-MSA}(\text{LN}(z^{l})) + z^{l} \\ z^{l+1} &= \text{MLP}(\text{LN}(\hat{z}^{l+1})) + \hat{z}^{l+1} \end{aligned} eq. 3

Here W-MSA uses the regular partition and SW-MSA the shifted one §3.2.

Window attention and relative position bias

Global self-attention over hw tokens costs 4hwC² + 2(hw)²C, while window attention with M² tokens per window costs 4hwC² + 2M²hwC eq. 1 eq. 2 (the paper omits softmax cost from both) §3.2. The window version is linear in hw for fixed M, which is why M is held at 7 §3.2 §3.3.

Within each window and head, attention is computed as

\text{Attention}(Q, K, V) = \text{SoftMax}(QK^T/\sqrt{d} + B)V eq. 4

where d is the query/key dimension and B is a learned bias taken from a table of size (2M−1) x (2M−1) indexed by the relative offset along each axis §3.2 swin_transformer.py swin_transformer.py. The paper reports this bias beats both no position encoding and absolute position embeddings, and that adding absolute embeddings to the input slightly hurts detection and segmentation §3.2 §4.4 tab. 4. Head dimension is fixed at d = 32 across all variants §3.3.

Shifted windows

Block l uses a regular partition that starts at the top-left pixel. Block l+1 displaces the partition by (floor(M/2), floor(M/2)) pixels §3.2 fig. 2. The paper's figure example uses an 8 x 8 map with M = 4, which gives 2 x 2 windows of 4 x 4 §3.2 fig. 2.

A shifted partition produces more windows, some smaller than M x M. Rather than pad each of them, the implementation rolls the feature map cyclically up and to the left, so the number of batched windows stays the same as in the regular partition. Tokens from different original regions then share a window, and a mask stops them from attending to each other §3.2 fig. 4.

# Shifted-window attention, one residual branch  [§3.2] [eq. 3] [fig. 4]
# tokens: (B, H*W, C), with H and W divisible by the window size M
def shifted_window_attention(tokens, height, width, window, shift):
    # Pre-norm, as in every Swin block  [§3.1] [eq. 3]
    normed = layer_norm(tokens)
    # Regroup the sequence as a grid  # (B, H, W, C)
    grid = reshape(normed, (batch, height, width, channels))
    if shift > 0:
        # Cyclic shift toward the top-left; nothing is padded  [§3.2] [fig. 4]
        grid = roll(grid, shift=(-shift, -shift), axes=(1, 2))
    # Cut the grid into nW windows of M*M tokens  [§3.2]
    windows = window_partition(grid, window)  # (B*nW, M*M, C)
    # Shifted layers only: mask pairs that came from different regions  [§3.2]
    mask = None
    if shift > 0:
        mask = shift_mask(height, width, window, shift)
    # Self-attention inside each window, with relative position bias  [eq. 4]
    windows = window_attention(windows, mask)
    # Undo the partition, then undo the cyclic shift  [§3.2]
    # [code: models/swin_transformer.py#L279-L284]
    grid = window_reverse(windows, window, height, width)
    if shift > 0:
        grid = roll(grid, shift=(shift, shift), axes=(1, 2))
    # Residual connection around the attention module  [eq. 3]
    output = reshape(grid, (batch, height * width, channels))
    return tokens + output

# Attention mask for the shifted layer  [§3.2]
# [code: models/swin_transformer.py#L222-L245]
# Every token gets a label for the region it came from before the shift;
# the 3 x 3 slice grid over H and W gives nine regions.
def shift_mask(height, width, window, shift):
    # Label each token by its (row slice, column slice)  # (1, H, W, 1)
    labels = region_labels(height, width, window, shift)
    # Same window partition as the features  # (nW, M*M)
    windows = window_partition(labels, window)
    # Pairs with different labels get -100, same-region pairs get 0  [§3.2]
    different = windows[:, None, :] != windows[:, :, None]
    return where(different, -100.0, 0.0)

# Window attention with relative position bias  [eq. 4] [§3.2]
# windows: (B*nW, M*M, C); d = C / heads = 32  [§3.3]
def window_attention(windows, mask):
    # Project to queries, keys and values per head  [§3.1]
    queries, keys, values = project_qkv(windows)
    # Scale the dot products by 1/sqrt(d)  [eq. 4]
    scores = matmul(queries, transpose(keys)) / sqrt(head_dim)
    # Add one learned bias per head and per relative offset  [§3.2]
    scores = scores + relative_bias(bias_table)  # (heads, M*M, M*M)
    if mask is not None:
        # Shifted windows only: add the mask to every window  [§3.2]
        scores = add_mask(scores, mask)
    # Softmax over the keys, then mix the values  [eq. 4]
    return matmul(softmax(scores), values)

Patch merging

# Patch merging: halve resolution, double channels  [§3.1] [tab. 7]
# grid: (B, H, W, C) with H and W even
def patch_merging(grid):
    # Gather the four cells of every 2 x 2 neighbourhood  [§3.1]
    top_left = grid[:, 0::2, 0::2, :]
    bottom_left = grid[:, 1::2, 0::2, :]
    top_right = grid[:, 0::2, 1::2, :]
    bottom_right = grid[:, 1::2, 1::2, :]
    # Concatenate into 4C channels  # (B, H/2, W/2, 4C)
    merged = concat(
        [top_left, bottom_left, top_right, bottom_right],
        axis=-1,
    )
    # Normalise, then reduce 4C to 2C with a linear map  [tab. 7]
    return linear(layer_norm(merged), out_features=2 * channels)

Stage schedule: the block pattern alternates regular and shifted, with the shift set to zero on every even-indexed block swin_transformer.py.

Hyperparameters

Architecture variants

Shared settings: patch size 4 §3.1 config.py; head dimension 32 §3.3; query/key/value bias on config.py; absolute position embedding off §3.2 config.py; patch-merging LayerNorm tab. 7 swin_transformer.py. Stochastic depth rises linearly with depth swin_transformer.py.

ImageNet-1K classification

COCO detection and instance segmentation

ADE20K semantic segmentation

Results

Author-reported, single-crop unless noted §4 tab. 1 tab. 2 tab. 3.

Model size and FLOPs for the classification models: Swin-T 29M / 4.5G, Swin-S 50M / 8.7G, Swin-B 88M / 15.4G at 224² tab. 1. Throughput was measured with the timm repository on a V100, following DeiT tab. 1. Larger input sizes trade speed for accuracy: Swin-B at 224², 256², 320² and 384² reaches 83.3, 83.7, 84.0 and 84.5 top-1 tab. 8.

Paper vs code

Code: https://github.com/microsoft/Swin-Transformer at commit f82860bfb5225915aca09c3227159ee9e1df874d, MIT license, owner Microsoft Corporation. The paper's abstract links this repository. The repository holds the classification code only. Its README points to separate repositories for COCO detection and ADE20K segmentation README.md, which we did not clone or check, so the detection and segmentation settings above rest on the paper alone app. A2.2 app. A2.3.

The classification code matches the paper on the architecture: window size, head dimension, cyclic shift, mask construction, relative bias table, patch merging, and classifier head swin_transformer.py swin_transformer.py swin_transformer.py swin_transformer.py. The discrepancies below are in training settings.

• Gradient clipping: paper gives max norm 1 app. A2.1; code default is 5.0 config.py, used in training at main.py. None of the classification configs overrides it.
• 22K pre-training weight decay: paper gives 0.01 app. A2.1; the Swin-T and Swin-B 22K configs set 0.05 swin_tiny_patch4_window7_224_22k.yaml swin_base_patch4_window7_224_22k.yaml.
• 22K pre-training schedule: paper gives linear decay app. A2.1; the 22K configs set no scheduler, so the default is cosine config.py.
• 22K→1K fine-tuning learning rate: paper gives a constant 1e-5 app. A2.1. The code uses a base LR of 2e-5 with cosine decay and 5 warm-up epochs swin_tiny_patch4_window7_224_22kto1k_finetune.yaml. The linear scaling rule multiplies the base LR by batch/512 main.py, so batch 1024 gives a peak of 4e-5, not 1e-5.
• Stochastic depth for higher-resolution fine-tuning: paper gives 0.1 app. A2.1. The 384² configs set 0.5 for Swin-B trained on IN-1K swin_base_patch4_window12_384_finetune.yaml and 0.2 for Swin-B pre-trained on IN-22K swin_base_patch4_window12_384_22kto1k_finetune.yaml. Only the Swin-T 22K→1K config uses 0.1 swin_tiny_patch4_window7_224_22kto1k_finetune.yaml.
• Window size at 384²: the paper sets M = 7 by default §3.3 and does not say otherwise. The 384² configs use M = 12 swin_base_patch4_window12_384_finetune.yaml.
• Label smoothing: not stated in the paper, but the code uses 0.1 and folds it into the Mixup targets config.py build.py main.py.
• Fixed input size: the paper gives no restriction on input size. The code asserts that the input matches the built img_size swin_transformer.py and has no padding path, so the window partition assumes divisibility.
• Attention mask: the paper says a masking mechanism limits attention to each sub-window §3.2. The code adds −100 to the masked pairs, not −inf swin_transformer.py.
• Paper-internal disagreement: Swin-B at 224² is 83.5 in tab. 1 but 83.3 in §4.1 and in tab. 8.

Pitfalls

• Stage-4 window clamp. If the feature map is no larger than the window, the code sets the shift to 0 and shrinks the window to the feature size swin_transformer.py. At 224², stage 4 is 7 x 7, so its shifted block runs unshifted, with no mask. At 384² with M = 12, stage 4 is 12 x 12 and behaves the same. The paper does not mention this.
• Relative bias table. It has (2M−1)² rows per head per block. A model pre-trained at one window size and fine-tuned at another needs the table resized by bicubic interpolation §3.2 utils.py. The code does this when the sizes differ.
• Weight decay. The code excludes relative_position_bias_table and absolute_pos_embed from weight decay swin_transformer.py. The paper does not say so.
• Initialisation. The bias table starts from a truncated normal with std 0.02 swin_transformer.py. The paper gives no initialisation.
• Shifted-layer masks cost time. Cyclic shift and masking keep the window count fixed but add a roll, a mask and a reverse roll per shifted block. The paper reports a 13% to 18% speed-up of cyclic over naive padding §4.4 tab. 5. The fused window kernel is optional and off by default config.py.
• Speed numbers depend on the implementation. Latency results are on a V100 with built-in PyTorch ops. The paper says these ops are not all well-optimised, while ResNe(X)t uses cuDNN §4.2. Compare speeds only with the same stack.
• Linear complexity holds for fixed M. Cost is linear in image size only while M stays fixed §3.2. Changing the input resolution or the window changes the per-block cost and the mask.
• Unstated optimisation details. Minimum LR and warm-up LR for regular training are not stated in the paper. The code defaults are 5e-6 (minimum) and 5e-7 (warm-up), multiplied by batch/512 config.py main.py.
• Learning rate depends on batch size. The paper states the LR for its batch sizes. Code scales it linearly with the total batch main.py. Keep the batch size in your config or recompute the LR.
• Throughput tables are not comparable across papers. Throughput uses the timm repository on a V100, following DeiT tab. 1.

Lineage

• Builds on: the standard Transformer for its self-attention and layer design §3.2 (Vaswani et al., arXiv:1706.03762), and ViT for patch tokens and pre-training §1 §3.2 (arXiv:2010.11929). The regular ImageNet-1K training recipe follows DeiT (arXiv:2012.12877) app. A2.1.
• Compared against: ViT, DeiT, ResNet and ResNeXt, RegNet, EfficientNet, DetectoRS, Copy-paste, HTC++, SETR, UperNet and Performer §4 tab. 1 tab. 2 tab. 3 tab. 6.
• Superseded by: not stated in this paper.