Swin Transformer — Windows, Shifts and a Pyramid for Dense Prediction

ViT and DeiT classify well, but they hand you one feature map at 1/16 of the input and pay for attention with the square of the number of tokens. A detector or a segmenter wants a pyramid of maps at strides 4, 8, 16 and 32, like a convnet gives it, at resolutions where a global ViT is simply not runnable. Swin's paper names the two things that separate pixels from words: visual entities "can vary substantially in scale", and images have "much higher resolution of pixels … compared to words in passages of text" [LIU21].

Swin's fix can be stated in three steps:

1. Start small and merge. 4×4-pixel patches instead of 16×16, and between stages, merge every 2×2 group of tokens into one. Four stages, four resolutions, a pyramid.
2. Attend inside windows. Self-attention runs within non-overlapping M×M windows (M = 7), so the cost grows linearly with the number of tokens instead of quadratically.
3. Shift the windows every other layer so information crosses window borders, and do the shift with a cyclic roll and a mask so it costs nothing extra.

Brush up before you start ;)

• Patch embedding and the single 14×14 map: An Image Is Worth 16×16 Words — ViT, Patch by Patch.
• Why n² kills images, local and blocked attention, sliding windows and 2-D relative positions: Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks. Swin's windows are not sliding windows; §4 is about the difference.
• The training recipe and the DeiT-S / DeiT-B baselines: DeiT — Training a Vision Transformer on ImageNet Alone.
• Why detection heads want strides 4 to 32: Object Detection — A Quick Read.

1. The Hierarchy: Patch Partition and Patch Merging

• Patch partition (§3.1 of the paper). Each token is a 4×4 patch, so 4 × 4 × 3 = 48 raw values, linearly embedded to C channels. At 224 px, stage 1 works on 56 × 56 = 3136 tokens; ViT-B at the same resolution has 196 [DOSO20].
• Patch merging sits at the start of stages 2, 3 and 4. Concatenate each group of 2×2 neighbouring tokens, LayerNorm, then a linear layer from 4C to 2C:

• Reading Eq. 1: tokens ÷ 4, resolution ÷ 2, channels × 2. It is the space-to-depth move, four neighbouring pixels stacked into one thicker pixel, followed by a projection that keeps only half of the stacked channels. Nothing else in the network changes resolution.
• The output sizes at 224 px are therefore 56², 28², 14² and 7², which is exactly the stride-4/8/16/32 pyramid an FPN plugs into [LIU21].

NOTE: "hierarchical" means resolution, not depth. ViT-L is deeper than Swin-T and still produces one 14×14 map. Patch merging, not more layers, is what makes Swin hierarchical.

The four sizes, all with head dimension d = 32 (the d_k of Query, Key, Value — Dot-Product Attention from Scratch), MLP ratio 4 and M = 7 (§3.3 and Table 7 of the paper) [LIU21]:

Table. 1 The Swin family (§3.3 and Tables 1 and 7 of [LIU21]). Swin-T is sized like ResNet-50 and DeiT-S, Swin-S like ResNet-101, Swin-B like ViT-B / DeiT-B.

Two more things that are not there: no class token (the stage-4 output is average-pooled, then a linear classifier, "as accurate as an additional class token", Appendix A2.1), and no absolute position embedding; §5 explains what replaces it.

2. Window Attention and the Cost of Looking Everywhere

Let's devise the cost equations step-by-step. Symbols: h and w are the token map's height and width (not pixels), C the channels (the D of An Image Is Worth 16×16 Words — ViT, Patch by Patch), M the window side in tokens (the k of Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks). The paper counts multiply-adds and leaves softmax out (footnote 3).

2.1 Global Attention, Eq. 1 of the Paper

Every attention layer does two kinds of work: four C×C projections (Q, K, V and the output) applied to every one of the hw tokens, and then $Q{K}^{\top }$ and $AV$, each hw × hw × C.

• The first term is linear in the number of tokens. The second is the problem: double the image side, four times the tokens, sixteen times the attention cost.

2.2 Window Attention, Eq. 2 of the Paper

Split the token map into non-overlapping M×M windows (padding at the bottom and right if hw is not divisible, footnote 4) and run ordinary multi-head self-attention inside each window only. There are hw/M² windows, each costing 2(M²)²C for its attention terms:

• Compare the second terms of Eq. 2 and Eq. 3: $(hw{)}^2$ became $M^{2}hw$. With M fixed, the whole thing is linear in the number of tokens. The projection term is identical in both.

Stage 1 of Swin-T at 224 px: h = w = 56, C = 96, M = 7.

• Projections: 4 · 3136 · 96² = 115.6 M multiply-adds, in both versions.
• Global attention term: 2 · 3136² · 96 = 1,888.2 M. Total Ω(MSA) ≈ 2.00 G per layer.
• Window attention term: 2 · 49 · 3136 · 96 = 29.5 M. Total Ω(W-MSA) ≈ 0.145 G per layer.
• The attention terms differ by hw / M² = 3136 / 49 = 64×; the whole layer by 13.8×.

Now slide the input side up in the figure, and try a bigger M while you are there:

input side in pixels

224

window M (tokens)

7

• Push the input side up and the gap explodes: per stage-1 layer, global attention costs 2.0 / 16.7 / 127 / 2017 G at 224 / 384 / 640 / 1280 px, the windowed version 0.145 / 0.43 / 1.19 / 4.74 G. At detection resolutions the global version is not an option.
• The same arithmetic per stage at 224 px, attention term only: 64× at stage 1, 16× at stage 2, 4× at stage 3, and 1× at stage 4, where the 7×7 map is one window and W-MSA is global MSA. The win is in the early stages and at high resolution.
• You can see it in whole networks too: Swin-B goes from 15.4 G at 224 px to 47.0 G at 384 px (×3.05 for ×2.94 the tokens, linear), DeiT-B from 17.5 G to 55.4 G (×3.17, the quadratic part showing) [LIU21].

3. Shifted Windows

Windows alone have an obvious hole: a token never hears from anything outside its 7×7 square. The paper's fix is to move the partition between consecutive layers so that the new windows "bridge the windows of the preceding layer, providing connections among them" [LIU21].

3.1 The Block Pair, Eq. 3 of the Paper

Blocks always come in pairs. The first uses the regular partition (W-MSA); the second moves the partition by (⌊M/2⌋, ⌊M/2⌋) tokens, 3 for M = 7, and attends inside the shifted windows (SW-MSA).

• $z^l\in {ℝ}^{h\times w\times C}$ is the output of block l and ${\widehat{z}}^l$ the output of its attention sub-layer. MLP is two layers with GELU and 4× expansion, pre-norm, residual after each module: the transformer block with MSA swapped for (S)W-MSA and nothing else touched.
• Because blocks come in pairs, every stage depth in Table 1 is even: {2, 2, 6, 2}, {2, 2, 18, 2}.

3.2 Why Not Just Shift? The Cyclic Trick

Shifting the partition naively creates a problem the figure below makes visible: an h/M × w/M grid of windows becomes (⌈h/M⌉ + 1) × (⌈w/M⌉ + 1), and the ones along the edges are smaller than M×M. Padding them back up to M×M costs real compute: the paper's 2×2 → 3×3 example is 2.25× [LIU21].

The paper's answer is a cyclic shift: roll the whole token map toward the top-left by s = ⌊M/2⌋, apply the regular partition, and roll back afterwards. The wrapped-around strips land in the last row and column of windows, which now hold pieces from opposite sides of the image. A mask stops them from attending to each other:

• The mask is the −∞-before-the-softmax trick of the decoder in Attention Is All You Need — The Transformer, Block by Block, so the masked pairs simply do not exist for that layer. The result is identical to the shifted partition; the cost is the regular partition's.
• The window count stays ⌈h/M⌉ × ⌈w/M⌉: 64 at stage 1 instead of 81 (+27 %) [LIU21].

Move the shift s in the figure: at s = 0 the middle panel is the left one; at s = 2 count the middle panel's windows against the right one, and the last readout says how many of the rolled windows' logits the mask throws away.

tokens per side N

8

shift s (tokens)

2

NOTE: the cyclic shift needs the mask. Without it, a token on the bottom edge of the image would attend to a token from the top edge as if they were neighbours, and "cyclic shift" would no longer equal "shifted partition" in result, only in cost.

3.3 Shifted Is Not Sliding

The sliding-window attention of Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks gives every query its own neighbourhood. That is cheap on paper and slow on hardware, because every query has a different key set. In Swin, "all query patches within a window share the same key set, which facilitates memory access in hardware" [LIU21]; the partition moves once per layer, not the neighbourhood. The paper measured it (Table 5, V100, milliseconds per attention layer in stages 1 to 4, and images per second for Swin-T / S / B) [LIU21]:

Table. 2 Speed of the attention variants (Table 5 of [LIU21]): ms per layer per stage, then whole-network images per second on a V100.

• Shifting costs almost nothing: 2.8 → 3.0 ms in stage 1, 770 → 755 im/s for Swin-T.
• Cyclic beats padding by 13 / 18 / 18 % throughput for T / S / B.
• Against the naive sliding window, the cyclic shifted window is 40.8× faster in stage 1 (122.5 vs 3.0 ms) and 4.1× faster as a network; 2.5× and 1.5× against the hand-written kernel. And Table 6 of the paper shows the sliding-window Swin-T is no more accurate: 81.4 top-1 / 50.2 box AP / 45.8 mIoU against 81.3 / 50.5 / 46.1 for shifted windows [LIU21]. Same accuracy, a quarter of the time.

4. Relative Position Bias

Swin drops ViT's absolute position embedding and instead teaches each head how much to favour each offset. The attention inside one window, Eq. 4 of the paper, is the usual scaled dot-product with one addition:

• $B_{ij}$ is a scalar added to the logit between query i and key j, and it depends only on their offset $(\Delta y,\Delta x)\in [-M+1,M-1{]}^2$. The same offset gets the same bonus everywhere in the window: a convolution-like prior, learned.
• It is gathered from a small table $\widehat{B}$ of (2M − 1)² entries per head, and it is added inside the softmax, which is why it is a bias and not an embedding.

How small? For M = 7, $\widehat{B}$ is 13 × 13 = 169 numbers per head, where a full M² × M² table would need 2,401. Swin-T has 3·2 + 6·2 + 12·6 + 24·2 = 138 head-layers, so 23,322 parameters in total, 0.08 % of 29M; Swin-B has 376 head-layers and 63,544. For that, Table 4 of the paper gives (Swin-T, top-1 / COCO box AP / ADE20K mIoU) [LIU21]:

• No position encoding: 80.1 / 49.2 / 43.8.
• Absolute embedding (ViT's): 80.5 / 49.0 / 43.2.
• Absolute + relative: 81.3 / 50.2 / 44.0.
• Relative bias only: 81.3 / 50.5 / 46.1.

So the bias is worth +1.2 top-1 and +2.3 mIoU over nothing, and adding ViT's absolute embedding on top hurts the dense tasks (−0.3 AP, −2.1 mIoU). The paper's conclusion: "inductive bias that encourages certain translation invariance is still preferable for general-purpose visual modeling" [LIU21]. Taking the dot product away and keeping only the bias gives 79.3, so the content term still matters. The table can also be bicubically interpolated to a new window size at fine-tuning time.

The same Table 4 prices the shift itself: without it Swin-T scores 80.2 / 47.7 / 43.3, with it 81.3 / 50.5 / 46.1, i.e. +1.1 top-1, +2.8 box AP, +2.8 mIoU for 0.2 ms per stage-1 layer [LIU21].

5. Training

Swin trains with the DeiT recipe, same optimizer, schedule, batch and augmentations, with four edits: 20 warm-up epochs instead of 5, gradient clipping back on (max norm 1), stochastic depth raised to 0.2 / 0.3 / 0.5 for T / S / B, and no repeated augmentation or EMA, "which do not enhance performance" in Swin's hands, contrary to what DeiT found [LIU21].

6. Results

6.1 ImageNet-1k

Table. 3 ImageNet-1k only, no distillation; every column from Table 1a of [LIU21], V100 throughput included. Swin-B at 224 px is 83.5 in Table 1 and 83.3 in the paper's text and Table 8, a mismatch between the paper's versions; I quote the table.

• At equal cost: Swin-T 81.3 vs DeiT-S 79.8 (+1.5 at 4.5 vs 4.6 GFLOPs); Swin-B 83.5 vs DeiT-B 81.8 at fewer FLOPs (15.4 vs 17.5).
• With ImageNet-22k pre-training (Table 1b of the paper): Swin-B↑384 86.4 against ViT-L/16's 85.2 at 27 im/s, and Swin-L↑384 87.3 [LIU21].

6.2 Where the Pyramid Pays: COCO and ADE20K

Table. 4 COCO detection (Table 2 of [LIU21]). Compare only within a row group; the frameworks differ between groups.

• Dropping Swin-T in for ResNet-50 is worth +3.4 to +4.2 box AP across four different detectors.
• Against DeiT-S with a pyramid bolted on by deconvolution, Swin-T is +2.5 box AP, +2.3 mask AP and 47 % faster: the hierarchy is cheaper than faking one.
• On ADE20K with UperNet (Table 3 of the paper): Swin-T 46.1 and Swin-S 49.3 mIoU against DeiT-S 44.0 and ResNet-101 44.9; Swin-L pre-trained on 22k reaches 53.5 [LIU21].

7. What Came Next

Swin made the transformer a drop-in backbone, and it did so by importing convnet ideas: a stem of small patches, four stages with downsampling between them, a 7×7 neighbourhood, a translation-friendly position prior. The next post, Convolutions Strike Back — Was It the Attention or the Recipe?, follows that traffic in both directions. PVT, CvT and CoAtNet put pyramids and convolutions back into attention networks, and ConvNeXt takes Swin's macro design, its 7×7 kernel, its LayerNorm and its DeiT recipe, removes the attention entirely, and scores 82.1 against Swin-T's 81.3. Which raises the question that post is really about: how much of this was the attention, and how much was everything around it?

8. References

1. [LIU21] Liu, Lin, Cao, Hu, Wei, Zhang, Lin, Guo. Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. arXiv:2103.14030.
2. [DOSO20] Dosovitskiy et al. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. arXiv:2010.11929.
3. [TOUV21] Touvron, Cord, Douze, Massa, Sablayrolles, Jégou. Training data-efficient image transformers & distillation through attention. arXiv:2012.12877.