Swin Transformer V2: Scaling Up Capacity and Resolution

In plain words

Large language models got much bigger and kept improving, but image models lagged behind. This paper takes the Swin Transformer, a vision model that looks at an image in local windows, and makes it much larger: about 3 billion parameters, able to read images up to 1,536 by 1,536 pixels.

Scaling it up exposed three problems, and the paper fixes each one. First, the bigger model's internal values grew unevenly from layer to layer until training became unstable. A small change to where the normalisation layers sit, plus a different way of computing attention scores, calms this down. Second, a model trained on small images needs its position information stretched when it sees larger windows. Instead of stretching a lookup table by interpolation, the authors learn a small network that produces position biases for any offset, computed on a logarithmic scale so the stretch stays modest. Third, labelled images are scarce at this scale, so the model is first pretrained without labels by hiding parts of each image and predicting them.

The result is a model that sets new records on four vision benchmarks: image classification, object detection, semantic segmentation and video action recognition. The authors also point out that they used far fewer labelled images than the large Google vision model they compare against.

TL;DR

• Swin Transformer V2 is a scaled, more stable version of the Swin vision transformer, evaluated up to 3B parameters and 1,536 px inputs §3.5 §4.2.
• Three changes: residual post-norm plus scaled cosine attention for stability §3.2; a log-spaced continuous position bias for transferring between window sizes §3.3; SimMIM self-supervised pretraining to reduce labelled data §3.4.
• Author-reported headline results for the 3B model are 84.0% ImageNet-V2 top-1, 63.1 COCO test-dev box AP, 59.9 ADE20K val mIoU and 86.8% Kinetics-400 top-1 tab. 2 tab. 3 tab. 4 tab. 5.
• Relevant if you need a dense vision backbone that transfers across input resolutions, or are scaling a vision transformer and see training divergence.
• The official code is MIT-licensed, but the 3B model's configs and its 70M-image pretraining set are not in the repository or public .

Problem

Language models had grown more than a thousandfold in a few years, while the largest vision models were about 1 to 2 billion parameters and were used only for image classification §1. Three obstacles block scaling a Swin-style vision backbone: (1) training becomes unstable as the model grows, and the largest model diverges §3.1 fig. 3; (2) a model pretrained at low resolution transfers poorly to larger windows when position biases are interpolated bi-cubically §3.1 tab. 1; (3) large models need very large labelled datasets, such as JFT-3B, which the authors avoid §1 §3.4. High-resolution fine-tuning also exhausts GPU memory unless the training implementation changes §3.5.

Method

Swin Transformer V2 keeps the original Swin stage, block and channel layout §3.6. Three architecture changes are made fig. 1, and one change is made to pretraining (item 4).

1. Residual post-norm §3.2. The original Swin block applies layer norm before each sub-layer (pre-norm). The authors move the layer norm to the output of each residual branch, before the branch is added to the main path. This keeps the main path's amplitude from growing with depth §3.2 fig. 2. In the largest models an extra layer norm is placed on the main branch every six blocks §3.2 §3.6.

2. Scaled cosine attention §3.2 eq. 2. The dot product of queries and keys is replaced by a cosine similarity divided by a learnable scalar τ. τ is separate for each head and layer, and the authors state it is kept above 0.01 §3.2. Position bias is added to the scaled cosine. The original attention with a parameterised bias is eq. 1.

3. Log-spaced continuous position bias (Log-CPB) §3.3 eq. 3 eq. 4. The learned bias table of Swin V1 is replaced by a small meta network G, a two-layer MLP with a ReLU in between by default §3.3. G takes the relative offset (Δx, Δy) between two positions and returns the bias for every head. Before G sees the offset, each coordinate is mapped to sign(x) · log(1 + |x|) eq. 4. Because G accepts any offset, a pretrained model can be run with a new window size without retraining §3.3. For inference, the biases can be computed once per relative position and stored as ordinary parameters §3.3. The authors give a worked case: moving from an 8×8 to a 16×16 window, the raw offset range extrapolates by a ratio of 8/7 ≈ 1.14, while the log-spaced range goes from ±2.079 to ±2.773, a ratio of 0.33 of the original range §3.3.

4. Self-supervised pretraining with SimMIM §3.4. The 3B model is first pretrained with SimMIM §3.4 (see 2111.09886) on the ImageNet-22K-ext set of about 70 million images, then pretrained again with classification labels on the same images §4.2. The authors report that this uses 1/40 of the labelled images used by JFT-3B §3.4.

Memory measures §3.5. (a) ZeRO stage-1 optimizer state sharding, via DeepSpeed; the authors report that a 3B model with AdamW and fp32 states needs 48 GB without it §3.5. (b) Activation checkpointing; the authors report up to 30% slower training §3.5. (c) Sequential instead of batched self-attention in the first two stages, for 1,536 px inputs with 32×32 windows on 40 GB A100s §3.5.

Model sizes §3.6. All sizes use four stages.

• SwinV2-T: C = 96, blocks {2, 2, 6, 2}.
• SwinV2-S, -B, -L: C = 96 / 128 / 192, blocks {2, 2, 18, 2}.
• SwinV2-H: C = 352, blocks {2, 2, 18, 2}; 658 million parameters.
• SwinV2-G: C = 512, blocks {2, 2, 42, 4}; 3 billion parameters. C is the channel count of the first stage §3.6.

Core equations, as written in the paper:

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

\text{Sim}(\mathbf{q}_i, \mathbf{k}_j) = \cos(\mathbf{q}_i, \mathbf{k}_j) / \tau + B_{ij} eq. 2

B(\Delta x, \Delta y) = \mathcal{G}(\Delta x, \Delta y) eq. 3

\widehat{\Delta x} = \text{sign}(x) \cdot \log(1+|\Delta x|), \quad \widehat{\Delta y} = \text{sign}(y) \cdot \log(1+|\Delta y|) eq. 4

Block-level sketch, following the description above. Window partitioning is Swin V1's §3.1; the code's shifted-window handling is not shown.

# Residual post-norm Swin V2 block  [§3.2] [fig. 1]
# Input:  features (B, H*W, C); window partitioning as in Swin V1  [§3.1]
# Output: features (B, H*W, C)

def swin_v2_block(features):
    # Attention branch: normalise its output before the residual add  [§3.2]
    attended = window_attention(features)
    features = features + norm_1(attended)
    # MLP branch with the same post-norm placement  [§3.2]
    features = features + norm_2(mlp(features))
    return features

# Scaled cosine attention with continuous bias  [eq. 2] [eq. 3]
# Input:  queries, keys, values (N, d) per head, N = Wh*Ww
# Input:  tau (learnable scalar per head); bias (N, N) per head
# Output: (N, d)

def window_attention(queries, keys, values, tau, bias):
    # Unit-normalise queries and keys, so their product is a cosine  [eq. 2]
    queries = normalize(queries)  # (N, d)
    keys = normalize(keys)        # (N, d)
    # Divide by tau, which is kept at or above 0.01  [§3.2] [eq. 2]
    tau = max(tau, 0.01)
    logits = (queries @ keys.T) / tau  # (N, N)
    # Add the position bias for every pair of positions  [eq. 2]
    logits = logits + bias  # (N, N)
    weights = softmax(logits, axis=-1)
    return weights @ values

# Log-spaced continuous position bias  [eq. 3] [eq. 4]
# Input:  offset (dx, dy) for one pair of positions in a window
# Output: one bias value per head

def continuous_bias(offset, meta_network):
    # Map each coordinate to log space, keeping its sign  [eq. 4]
    log_offset = sign(offset) * log(1 + abs(offset))  # (2,)
    # Two-layer MLP with ReLU in between, the default G  [§3.3]
    return meta_network(log_offset)

Hyperparameters

Values are as written in the paper unless the Where column names the code. "Not stated" means the paper does not give it. Learning rates in the code are per 512 samples and are scaled by batch size main.py. The code values shown are converted to the paper's batch size.

Ablation pretraining (ImageNet-1K, 256², window 8²) app. A1

Fine-tuning for ablations app. A1

SwinV2-B and SwinV2-L, ImageNet-22K pretraining app. A2

SwinV2-B and SwinV2-L, ImageNet-1K fine-tuning app. A2

SwinV2-B and SwinV2-L, downstream tasks app. A2

SwinV2-G (3B) app. A2

Architecture

Results

Author-reported. The setups differ across rows, and the authors state they did not align model complexity or pretraining data across comparisons §4.2.

Headline results (SwinV2-G, 3B)

SwinV2-B and SwinV2-L against Swin V1 (ImageNet-22K pretraining, ImageNet-1K)

Ablations (ImageNet-1K pretraining unless stated)

Paper vs code

Code paths are relative to the repository root at code_commit. The repository is the one the paper's abstract links [abstract] and is MIT-licensed LICENSE.

• CPB coordinate transform. Paper: sign(x) · log(1 + |x|) on raw offsets eq. 4. Code: offsets are divided by (pretrained window − 1) when a pretrained window is set, else by (window − 1); then multiplied by 8; then sign(x) · log₂(1 + |x|) / log₂(8) swin_transformer_v2.py. The worked example in §3.3 (±2.079 and ±2.773) follows the paper's formula, not the code.
• CPB output bounded. Paper: B = G(Δx, Δy) with no bound eq. 3. Code: 16 · sigmoid(G) swin_transformer_v2.py, so every bias lies in (0, 16).
• Bias recomputed per forward pass. Paper: the biases can be precomputed and stored for inference §3.3. Code: the table is passed through the MLP on each forward call swin_transformer_v2.py. Precomputation is not in this file.
• Residual gain initialisation. Not in the paper. Code sets the weight and bias of norm1 and norm2 in every block to zero at model init swin_transformer_v2.py, called from the model constructor swin_transformer_v2.py. Each residual branch therefore outputs zero at initialisation.
• Label smoothing. The config default is 0.1 config.py, but it applies only when Mixup is off main.py. The repository's SwinV2 configs keep Mixup on config.py, so label smoothing is not used there. The paper does not mention label smoothing.
• Learning-rate convention. Paper: absolute learning rates app. A1 app. A2. Code: per 512 samples, multiplied by batch size and world size over 512 main.py. The values agree at the paper's batch sizes: 1.25e-4 × 4096 / 512 = 1e-3, 5e-4 × 1024 / 512 = 1e-3, and 2e-5 × 1024 / 512 = 4e-5.
• Attention heads. Not stated in the paper. Base config: 4, 8, 16, 32 swinv2_base_patch4_window12_192_22k.yaml.
• Not in the repository. (a) ZeRO stage-1 via DeepSpeed (no DeepSpeed config found). (b) Sequential self-attention for the first two stages. (c) The extra layer norm every six blocks for H and G. (d) Configs for SwinV2-H and SwinV2-G, since configs/swinv2/ has tiny, small, base and large only. (e) The 70M-image ImageNet-22K-ext pretraining set, which the paper describes as privately collected §4.1. Activation checkpointing is present as a flag main.py and in the block module swin_transformer_v2.py.
• SimMIM details. The repository includes SimMIM configs for Swin V2 base simmim_pretrain__swinv2_base__img192_window12__800ep.yaml. Its 800-epoch schedule is not the 20-epoch stage used for SwinV2-G app. A2.

Pitfalls

• Text and table disagree on Log-CPB at window 24². The text gives 78.9% without fine-tuning §4.3; Table 1 gives 79.1% tab. 1. Use the table.
• SwinV2-B gain over V1. The text gives +0.8% for ImageNet-1K V1 §4.2; Table 2 gives 87.1 − 86.4 = +0.7 tab. 2.
• SwinV2-G pretraining is not reproducible from this release. The 70M-image set is private, the G configs are absent, and the appendix cites the self-supervised method with an anonymous key app. A2. The main text names SimMIM §3.4. Treat the G pretraining recipe as "SimMIM, details not stated here".
• Unstated values that change results. Attention head counts per size §3.6; CPB hidden width §3.3; τ initialisation §3.2; warm-up for ImageNet-1K fine-tuning app. A2; Mixup and CutMix alphas app. A1.
• Baseline stability. Post-norm diverges at the default learning rate; the paper uses one quarter of the default rate for that baseline tab. 7 caption.
• Window divisibility. Even window sizes keep the shift offset divisible by the window size app. A1 fn.. Odd sizes such as 7 also work, per the same footnote.
• Test-time window changes. The authors report +0.2 gains on ADE20K and Kinetics-400 from a larger test window §4.2, and a further gain on COCO from a different test window §4.2 tab. 3. They attribute these to Log-CPB; no control run isolates it.
• Memory and compute are the hidden costs. Without ZeRO, a 3B model with AdamW and fp32 states needs 48 GB for states alone §3.5. Activation checkpointing can slow training by up to 30% §3.5. The G-scale pretraining cost is an estimate of "< 0.5k" TPUv3 core-days tab. 2 fn..
• Comparisons are not controlled. The authors did not align complexity or pretraining data §4.2. Their ImageNet-1K V1 result is below CoAtNet-7 tab. 2.

Lineage

• Builds on: Swin Transformer arXiv:2103.14030 (no card here). V2 keeps its stage and block layout and changes the normalisation, attention scores and position bias §3.2 §3.3 §3.6.
• Builds on: SimMIM 2111.09886 for self-supervised pretraining §3.4.
• Related: continuous convolution work (SchNet, continuous convolution networks, relation-based continuous convolution) that uses a meta network on irregular offsets §2.
• Compared against: ViT-G, V-MoE and CoAtNet-7 (all JFT-3B) for ImageNet tab. 2; CopyPaste, SwinV1-L, YOLOR, CBNet, DyHead and SoftTeacher for COCO tab. 3; SwinV1-L, Focal-L, CSwin-L, MaskFormer, FaPN and BEiT for ADE20K tab. 4; ViViT, video SwinV1-L and TokenLearner for Kinetics-400 tab. 5.
• Superseded by: not known from wiki content.