Vision Transformers: Applying Transformer Architecture to Computer Vision

Suppose the Ministry of Agriculture wants a model that scans Sentinel-2 satellite tiles over Punjab and flags every hectare of wheat showing early signs of yellow rust — a pixel-level segmentation task, not a whole-image label. Each tile is at least 1024×1024 pixels, because rust patches can be a few metres wide and a coarser tile would blur them into noise. A model built the way the original Vision Transformer was built — split the image into a fixed grid of patches, flatten them into a token sequence, and run full self-attention where every token looks at every other token — hits a wall almost immediately. At patch size 4, a 1024×1024 tile produces 256×256 = 65,536 tokens. Self-attention needs, per head per layer, an attention matrix of size 65,536 × 65,536 ≈ 4.29 billion entries. Stored in fp16, that is roughly 8.6 GB — for one attention head, in one layer, on one tile. Swin-T's first stage alone runs three attention heads across two transformer blocks before any patch merging even happens. The arithmetic makes the point without needing a GPU spec sheet: uniform global attention over high-resolution images is not merely slow, it is architecturally incompatible with dense prediction at real resolution.

This chapter is about the fix that made transformers usable as backbones for detection and segmentation, not just classification: the Swin Transformer (Liu, Ze, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo, "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows," ICCV 2021). Where the original ViT (Dosovitskiy et al., "An Image Is Worth 16×16 Words: Transformers for Image Recognition at Scale," ICLR 2021) keeps a single token resolution and full global attention through every layer, Swin restructures the computation along two axes at once: it restricts attention to small local windows so cost stops being quadratic in image size, and it merges patches after each stage so the network produces a pyramid of feature maps — exactly the shape that detection and segmentation heads (Mask R-CNN, FPN, UPerNet) were built to consume from CNN backbones. This is a genuinely different mechanism from the pretraining-objective story (masked prediction, self-distillation, DINOv2-style targets) — it is about how attention itself is computed and how a transformer produces multi-scale features, and it's the piece that made satellite-scale, dense-prediction vision transformers practical in production.

Why quadratic attention is the actual bottleneck

Recall the shape of standard multi-head self-attention (MSA) on a sequence of N tokens of dimension C: queries, keys and values are each an N×C matrix produced by a linear projection (cost ∝ NC²), and the attention scores QKᵀ form an N×N matrix (cost ∝ N²C), which is then used to weight V (another N²C). For an h×w grid of patches, N = hw, so the dominant term scales as (hw)²C — quadratic in the number of patches. The Swin paper writes this out precisely for one MSA module operating on an h×w token grid of channel width C:

Ω(MSA)   = 4hwC² + 2(hw)²C

The 4hwC² term is the cost of the four linear projections (Q, K, V, and the output projection); the 2(hw)²C term is the cost of computing QKᵀ and applying the result to V. For a 224×224 image cut into 4×4 patches, the first Swin stage has h = w = 56. Already hw = 3,136, so (hw)² ≈ 9.83 million — and that number appears multiplied by C in the dominant term. This is the term windowing is built to kill.

Window-restricted attention: making the quadratic term linear

Swin's fix is architecturally simple: partition the h×w grid into non-overlapping M×M windows (Swin's default is M = 7) and run self-attention only within each window, never across window boundaries. A window has M² tokens, so one window's attention cost is proportional to (M²)² = M⁴ — quadratic in M, not in hw. Since M is a small, fixed constant (7) regardless of image size, and the number of windows is hw/M², the total cost across all windows is:

Ω(W-MSA) = 4hwC² + 2M²hwC

The linear-projection term is unchanged (every token still gets projected once); the attention term has gone from 2(hw)²C to 2M²hwC — linear in hw instead of quadratic, with M² as a fixed multiplier. This is the whole trick: attention that used to cost O((hw)²) now costs O(hw), at the price of each token only seeing its own M×M neighbourhood instead of the whole image.

Worked example: how much does windowing actually save, stage by stage?

Swin-T's four stages use channel widths C = 96, 192, 384, 768, doubling each time patches are merged (explained below), while the spatial grid halves each time: 56×56 → 28×28 → 14×14 → 7×7. Window size stays fixed at M = 7. Plugging each stage's (h, w, C) into both formulas gives:

Two things fall out of this table that are easy to miss just from staring at the formula. First, the saving is enormous exactly where it matters most: at stage 1, where the token grid is largest, window attention costs 13.81× less than global attention. Second, the saving evaporates by design as the network gets deeper — by stage 4 the feature map has shrunk to exactly 7×7, which is precisely one window, so W-MSA and MSA become mathematically identical (ratio 1.00×, a single window covering the whole map). Windowing isn't a permanent restriction; it's a restriction that matters only while the token grid is still large, which is exactly when the quadratic term would otherwise be unaffordable. There's a clean piece of arithmetic hiding in the table: looking only at the attention term, the ratio (hw)²C ⁄ (M²hwC) = hw/M² is exactly the window count — 64, 16, 4, 1 — and it quarters cleanly at every stage because hw quarters at every merge while M stays fixed at 7. The full ratio (13.81×, 2.70×, 1.18×, 1.00×) falls faster than that, because the 4hwC² projection term is identical in both formulas and doesn't shrink relative to itself — as hw shrinks and C doubles each stage, that shared linear term stops being negligible and increasingly dominates both numerator and denominator, pulling the overall ratio down toward 1 well before the window count itself reaches 1.

Patch merging: how the pyramid gets built

The C-doubling, h-halving pattern in the table above isn't incidental — it's a specific operation called patch merging, applied between stages. Given an h×w×C feature map, patch merging groups each non-overlapping 2×2 block of neighbouring tokens, concatenates their channel vectors (giving h/2 × w/2 × 4C), and applies a single linear layer projecting 4C down to 2C. The result is an (h/2)×(w/2)×2C map — half the spatial resolution, twice the channel width, at every stage boundary. This is architecturally the transformer analogue of a CNN's strided convolution or pooling layer, and it's precisely why a Swin backbone can be dropped into Mask R-CNN or UPerNet in place of a ResNet: it emits the same four-scale feature pyramid (1/4, 1/8, 1/16, 1/32 of the input resolution) that those detection and segmentation heads were designed around. Plain ViT, by contrast, keeps one fixed token resolution throughout — useful for classification, unusable as a drop-in dense-prediction backbone without extra machinery to synthesize multiple scales after the fact.

The problem windowing alone doesn't solve

Non-overlapping windows are cheap, but they're also blind to their own borders — a token in window A has zero attention connection to a token in the adjacent window B, no matter how close they are physically, because the window boundary is absolute. Stack enough plain W-MSA layers and you get something closer to many small independent CNNs bolted side by side than a single coherent transformer: information can travel within a 7×7 patch neighbourhood, but it can never cross into the next one. Swin's second idea directly patches this hole: every other transformer block uses windows that are displaced by half a window relative to the previous block. Displacing the grid line means a window in the shifted layer straddles four windows from the previous, regular layer — so a pair of consecutive blocks (one regular W-MSA, one shifted SW-MSA) lets information that was isolated in layer l reach across a former boundary in layer l+1. Stack several such pairs and the effective receptive field grows layer by layer, the same way stacking 3×3 convolutions grows a CNN's receptive field — except here it's window boundaries dissolving rather than kernel taps overlapping.

The naive way to implement this — literally redraw the window grid lines shifted by M/2 — creates a problem: the windows along the border of the shifted grid are smaller than M×M (some are M/2×M, others M/2×M/2), so you end up batching windows of unequal size, which is awkward and slow on a GPU. The paper's actual implementation avoids this with a cyclic shift: instead of redrawing boundaries, it cyclically rolls the entire feature map by (⌊M/2⌋, ⌊M/2⌋) — content that falls off the bottom-right wraps around to the top-left, like a texture tiling — and then applies the exact same regular M×M partition used in the un-shifted layer. Every window is now uniformly M×M again, so the batched matrix multiply is identical in shape to the regular W-MSA case; the only extra cost is a small residual mask.

The diagram: patch merging pyramid and the shift mechanism

The misconception worth correcting

The natural reading of "the shifted window mixes patches from four different original windows" is that the model now gets free attention between all patches that land in a shifted window — as if the shift simply merges four windows into one bigger, fully-connected window. That's wrong, and the mistake matters because it's exactly the kind of thing you'd assume from the diagram alone. The cyclic shift is a wraparound: the top few rows of a shifted window might contain patches from the true top of the image glued next to patches cyclically wrapped from the true bottom — two regions that are nowhere near each other in the actual picture. Attending between them would let a tree canopy at the top of a tile influence a rooftop at the bottom purely because of how the tensor happened to be rolled in memory. Swin prevents exactly this with an attention mask: before the softmax, every query-key pair that was not spatially adjacent prior to the shift gets its attention score set to a large negative number (−100 in the reference implementation), so after softmax that pair's weight is effectively zero. Only sub-blocks that were genuinely touching in the original image — i.e., across a real former window boundary, not across the wraparound seam — get to exchange information. The shift creates new connections; the mask decides which of the newly-adjacent pairs are real.

Tracing the mask: verified on a toy 8×8 grid

The construction below is the actual masking logic used in the reference implementation, run here on a small 8×8 feature map with window size M = 4 and shift = M//2 = 2 — small enough to inspect by hand, but built from the identical operations Swin-T runs at H = W = 56, M = 7, shift = 3.

import torch

def window_partition(x, window_size):
    # x: (B, H, W, C)
    B, H, W, C = x.shape
    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
    return windows.view(-1, window_size, window_size, C)

def build_shift_mask(H, W, window_size, shift_size):
    img_mask = torch.zeros((1, H, W, 1))
    h_slices = (slice(0, -window_size), slice(-window_size, -shift_size), slice(-shift_size, None))
    w_slices = (slice(0, -window_size), slice(-window_size, -shift_size), slice(-shift_size, None))
    cnt = 0
    for h in h_slices:
        for w in w_slices:
            img_mask[:, h, w, :] = cnt
            cnt += 1
    mask_windows = window_partition(img_mask, window_size).view(-1, window_size * window_size)
    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)
    return attn_mask

H = W = 8
window_size, shift_size = 4, 2
attn_mask = build_shift_mask(H, W, window_size, shift_size)
for w in range(attn_mask.shape[0]):
    blocked = int((attn_mask[w] == -100).sum().item())
    print(f"window {w}: blocked pairs = {blocked} / 256")

Running this produces exactly:

window 0: blocked pairs = 0 / 256
window 1: blocked pairs = 128 / 256
window 2: blocked pairs = 128 / 256
window 3: blocked pairs = 192 / 256

Checking each window's actual patch origins against the mask confirms the story the diagram tells: window 0 — the dashed top-left window in the figure — draws 4 patches from each of the four original windows (A, B, C, D) yet has zero blocked pairs, because that particular shifted window happens to sit entirely inside one contiguous, spatially-adjacent neighbourhood of the real image; the four-way colour mix is cosmetic, not a discontinuity. Windows 1 and 2 wrap around one axis only (128 of 256 pairs blocked — two disjoint groups of 8 mutually-attending tokens each). Window 3, the bottom-right corner, wraps around both the row and column axis simultaneously and ends up as four fully isolated 4-token groups (192 blocked), the maximum fragmentation possible in this toy example. Also worth noting: the mask uses −100 rather than −∞. exp(−100) is close enough to zero that softmax treats it as fully masked, but unlike −∞ it stays a finite, representable fp16 number, avoiding NaNs that can appear from operations on true infinities during mixed-precision training.

Why this backbone shape matters beyond the FLOP count

Coming back to the satellite-tile problem: a Swin backbone processes the 1024×1024 tile through exactly the four-stage pyramid this chapter has been building, ending with feature maps at 1/4, 1/8, 1/16 and 1/32 resolution — the same multi-scale bundle a segmentation head like UPerNet expects, because it was designed for ResNet-style CNN pyramids in the first place. Swap the backbone, keep the head. On ImageNet-1K classification, the original paper reports Swin-T reaching 81.3% top-1 accuracy training from scratch at 224×224 — competitive with the ResNet and EfficientNet backbones it was designed to replace, while additionally supporting the pyramid structure those CNNs provided natively and plain ViT does not. That combination — CNN-like multi-scale output, linear-in-resolution attention cost, transformer-style global modelling capacity within each window — is the specific reason Swin (and its descendants) became a standard detection and segmentation backbone rather than staying a classification curiosity.

Active recall

Attempt each question before reading its answer.

1. For a hypothetical Swin stage with h = w = 32, C = 128, M = 8, compute Ω(MSA) and Ω(W-MSA) and their ratio.
2. Stage 1 of Swin-T uses M = 7 (window count 64, ratio ≈13.81×). If you instead set M = 14 for that same stage (h = w = 56, C = 96), trace the full ripple: new Ω(W-MSA), new ratio to Ω(MSA), new window count, tokens per window, new shift size, and any constraint on M this change bumps into.
3. Why must regular (W-MSA) and shifted (SW-MSA) blocks alternate in pairs — why wouldn't two consecutive SW-MSA blocks with the same shift offset achieve the same cross-window mixing?
4. In the toy 8×8, M = 4, shift = 2 example traced above, how many of the 4 windows are fully unmasked, and how many attention-score entries in total (summed across all 4 windows) are set to −100?
5. Why does the reference implementation mask with −100 instead of −∞, and what specifically could go wrong numerically with −∞ under fp16 training?
6. A 1024×1024 satellite tile at patch size 4 gives a stage-1 grid of h = w = 256 (C = 96, M = 7 unchanged). Compute Ω(MSA):Ω(W-MSA) for this input, and separately estimate the memory (in GB, fp16) needed just to store one global attention score matrix for one head at this resolution. What does that number tell you about deploying plain ViT-style global attention at this scale?

Answers

1. Ω(MSA) = 4(32)(32)(128²) + 2(32·32)²(128) = 4·1024·16384 + 2·1,048,576·128 = 67,108,864 + 268,435,456 = 335,544,320. Ω(W-MSA) = 4·1024·16384 + 2(8²)(1024)(128) = 67,108,864 + 2·64·1024·128 = 67,108,864 + 16,777,216 = 83,886,080. Ratio = 335,544,320 / 83,886,080 = exactly 4.0×.

2. New Ω(W-MSA) = 4(56)(56)(96²) + 2(14²)(56·56)(96) = 115,605,504 + 2·196·3,136·96 = 115,605,504 + 118,013,952 = 233,619,456. Ω(MSA) is unchanged at 2,003,828,736 (it doesn't depend on M), so the new ratio is 2,003,828,736 / 233,619,456 ≈ 8.58× — smaller than the M = 7 ratio of 13.81×, because a bigger window claws back some of the quadratic cost inside each window. Window count drops from (56/7)² = 64 to (56/14)² = 16, while tokens per window rise from 49 to 196. Shift size, defined as ⌊M/2⌋, becomes 7 instead of 3, so the mask's region-splitting slices also change shape. The constraint this bumps into: M must divide h and w exactly at every stage where it's used unless the feature map is padded; 56/14 = 4 still works, but the real danger is stage 4, where h = w = 7 and M would have to be ≤7 and divide 7, so M ∈ {1, 7} only. A window size chosen to help stage 1 can silently become invalid at a deeper stage unless the implementation pads the feature map to the next multiple of M.

3. Two SW-MSA blocks with an identical shift offset partition the grid the same way both times — the second shift doesn't move the window lines any further, so it reconnects exactly the same pairs of previously-separated tokens the first shift already connected, adding no new boundary crossings. It's the alternation between two different partition phases (regular, then shifted) that keeps introducing new adjacent pairs each time the phase changes, the same way a CNN's receptive field only keeps growing because each new layer's kernel is centred differently relative to the previous layer's, not because you stacked two identical kernels back to back.

4. Exactly 1 of the 4 windows (window 0) is fully unmasked (0 blocked pairs). Total blocked entries across all four windows: 0 + 128 + 128 + 192 = 448 out of 4×256 = 1,024 total attention-score entries.

5. exp(−100) ≈ 3.7×10⁻⁴⁴, which softmax treats as effectively zero, achieving the masking goal. −∞ would do the same in exact arithmetic, but fp16 cannot represent finite arithmetic safely once true infinities enter a computation graph: subtracting the row max during a stable softmax can produce −∞ − (−∞) = NaN if an entire row happens to be masked, and NaNs silently propagate through the rest of the forward and backward pass, corrupting gradients well beyond the masked window. −100 sidesteps this because it's an ordinary finite fp16-representable number throughout.

6. Ω(MSA) = 4(256²)(96²) + 2(256²)²(96) = 4·65,536·9,216 + 2·4,294,967,296·96 ≈ 2.42×10⁹ + 8.246×10¹¹ ≈ 8.27×10¹¹. Ω(W-MSA) = 4·65,536·9,216 + 2(7²)(65,536)(96) ≈ 2.42×10⁹ + 6.16×10⁸ ≈ 3.03×10⁹. Ratio ≈ 273×. Separately, one global attention score matrix at this resolution has 256² × 256² = 65,536² ≈ 4.29 billion entries; at 2 bytes (fp16) each, that's ≈8.6 GB — for a single attention head, in a single layer, before counting the query/key/value projections, the other heads, or any of the deeper layers. That single number is the practical argument against plain ViT-style global attention for tile-scale dense prediction: no realistic single-GPU budget survives more than one or two such matrices simultaneously, which is exactly the wall windowed attention is built to avoid.

Think About It

Think about this: How would you explain vision transformers: applying transformer architecture to computer vision to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where vision transformers: applying transformer architecture to computer vision is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting vision transformers: applying transformer architecture to computer vision to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind vision transformers: applying transformer architecture to computer vision, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.