Vision Transformers: Image Understanding with Attention

During major flood events, ISRO's National Remote Sensing Centre (NRSC) processes satellite imagery — optical frames from Cartosat and radar frames from RISAT — into rapid flood-extent maps that district control rooms use to route relief boats and evacuate low-lying wards. The hard part of that job is rarely spotting an obvious lake of standing water in the middle of a frame. It is a patch of grey pixels near a river bend that could be floodwater, could be the shadow of a cloud, or could be a wet paddy field that always looks like this in August. A human analyst resolves that ambiguity by looking elsewhere in the image: does the grey patch connect to the river's known course? Do nearby elevation contours suggest a natural floodplain? Is the same tone visible in a dozen other patches that are unambiguously river?

A convolutional neural network, the workhorse you have already studied for image tasks, is structurally bad at this. A 3×3 convolution kernel centred on the ambiguous patch can only see pixels a few steps away in that one layer. To let that patch "see" evidence from a river bend forty patches away, you need to stack enough convolution and pooling layers that the receptive field grows wide enough to include it — and even then, distant evidence gets compressed and diluted as it passes through many layers. The Vision Transformer (ViT), introduced by Dosovitskiy et al. in the 2021 paper "An Image is Worth 16×16 Words," removes this constraint entirely. It treats an image as a sequence of patches and applies the same self-attention mechanism you met in the NLP transformer chapter, where every token can directly attend to every other token in a single layer, regardless of distance. A patch in the top-left corner can attend to a patch in the bottom-right corner just as easily as it attends to its immediate neighbour. This chapter builds that mechanism from first principles and traces it through actual numbers.

From pixels to a sequence: patchify and embed

A transformer encoder, as you have already seen for text, consumes a sequence of vectors. Text gives you that sequence for free — words, or sub-word tokens, in order. An image is a 2-D grid of pixels, not a sequence, so ViT's first job is to manufacture one. Given an image of height H, width W, and C channels, ViT cuts it into a grid of non-overlapping square patches of side P. The number of patches is

N = (H / P) × (W / P)

For the ViT-Base configuration in the original paper, H = W = 224 and P = 16, giving N = 14 × 14 = 196 patches. Each patch, before anything else happens to it, is just a small block of pixels: P × P × C numbers. For a 16×16 RGB patch that is 16 × 16 × 3 = 768 raw values. That block is flattened into a single vector of length 768 and passed through a learned linear projection E (a 768 × D weight matrix, where D is the transformer's embedding width — also 768 in ViT-Base, a coincidence of that particular configuration, not a rule) to produce one patch embedding vector. Do this for every patch and you have converted a 224×224×3 image into a sequence of 196 vectors, each of length 768 — precisely the shape a transformer encoder expects.

Two more pieces complete the input sequence, both borrowed directly from BERT-style NLP transformers. First, a single learnable vector called the class token, written [CLS], is prepended to the sequence. It carries no image content of its own at the start; over the course of the encoder's layers, it accumulates information from every patch via attention, and its final-layer representation is what gets read off for classification. Second, because attention has no built-in sense of order or position (more on this in the misconception below), a learned position embedding is added to every token, patch embeddings and the class token alike, so the model can tell "top-left patch" apart from "bottom-right patch" even when their content is identical.

Self-attention, patch by patch: a fully worked example

To see exactly what happens inside one layer, work through a deliberately tiny case by hand — small enough to trace every number, large enough to show the real mechanism. Take a 4×4 grayscale image with pixel intensities normalised to [0, 1]:

0.1  0.1  0.9  0.9
0.1  0.1  0.9  0.9
0.2  0.2  0.8  0.8
0.2  0.2  0.8  0.8

Use patch size P = 2, giving N = (4/2) × (4/2) = 4 patches. Flattening row-major:

Patch A (top-left)     = [0.1, 0.1, 0.1, 0.1]
Patch B (top-right)    = [0.9, 0.9, 0.9, 0.9]
Patch C (bottom-left)  = [0.2, 0.2, 0.2, 0.2]
Patch D (bottom-right) = [0.8, 0.8, 0.8, 0.8]

Project each 4-value patch down to a 2-dimensional embedding using a fixed linear map E with weight matrix rows [[1,0],[1,0],[0,1],[0,1]] and no bias, so that for a patch [p1,p2,p3,p4], the embedding is e = [p1+p2, p3+p4]:

e_A = [0.2, 0.2]   e_B = [1.8, 1.8]   e_C = [0.4, 0.4]   e_D = [1.6, 1.6]

Notice something important: every one of these embeddings lies exactly on the line y = x. Content alone cannot tell the model where a patch sits in the image; A and C both look "small and diagonal," B and D both look "large and diagonal." This is exactly why position embeddings are added next. Take four small, distinct position vectors — one per grid location — and add them:

P_A = [0.0, 0.1]   P_B = [0.1, 0.0]   P_C = [-0.1, 0.0]   P_D = [0.0, -0.1]

z_A = e_A + P_A = [0.20, 0.30]
z_B = e_B + P_B = [1.90, 1.80]
z_C = e_C + P_C = [0.30, 0.40]
z_D = e_D + P_D = [1.60, 1.50]

These four vectors, z_A through z_D, are the tokens that enter self-attention (a real ViT would also carry a [CLS] token through this; it is omitted here only to keep the arithmetic small). Self-attention with a single head, dimension d_k = 2, and — purely to isolate the attention computation itself — query, key and value projection matrices set to the identity, so Q = K = V = Z:

import numpy as np

Z = np.array([[0.2, 0.3],   # z_A
              [1.9, 1.8],   # z_B
              [0.3, 0.4],   # z_C
              [1.6, 1.5]])  # z_D

Q = K = V = Z
d_k = 2
scores = (Q @ K.T) / np.sqrt(d_k)
weights = np.exp(scores) / np.exp(scores).sum(axis=1, keepdims=True)
output = weights @ V

print(scores[0])    # [0.091924 0.650538 0.127279 0.544472]
print(weights[0])   # [0.186687 0.326376 0.193406 0.293531]
print(output[0])    # [1.185123 1.161142]

Trace the query row for patch A by hand to confirm it. The raw dot products Q_A · K_j are 0.13 (with itself), 0.92 (with B), 0.18 (with C), and 0.77 (with D) — B stands out because its coordinates are large and point in nearly the same direction as A's. Dividing by √2 ≈ 1.4142 and applying softmax turns these into the attention weights [0.187, 0.326, 0.193, 0.294]. The output for patch A is the weights-times-values sum: 1.185×0.2 + ... , working out to [1.185, 1.161].

Read what that number means physically. Patch A started as [0.20, 0.30] — a small, dark patch. After one layer of self-attention, its representation has moved to [1.185, 1.161], dragged more than halfway toward the bright cluster because it put its single largest attention weight (0.326) on patch B, a patch that shares no border with it. Running the same computation with patch B as the query (you can verify by swapping which row of weights you read) gives attention weights [0.010, 0.671, 0.013, 0.306] — B attends overwhelmingly to itself and to D, the other bright patch, and almost ignores the dark patches. Bright patches reinforce each other; the dark patch reaches out and borrows context from the bright ones. That asymmetric, content-driven, distance-independent mixing is the entire point of the architecture, and it is exactly what the flood-mapping scenario needs: an ambiguous patch pulling in evidence from patches that are relevant by content, not merely by proximity.

Real ViT layers do not use a single head or an identity projection. They split the D-dimensional embedding into h heads of size d_k = D / h, run the scaled dot-product attention above independently in each head with its own learned W_Q, W_K, W_V, concatenate the h outputs, and pass the result through one more learned linear layer. ViT-Base uses D = 768 split across h = 12 heads, so each head operates on 64-dimensional queries and keys — the toy example's d_k = 2 stands in for that same 64 in miniature. Each attention sub-layer sits inside a full transformer encoder block, identical in structure to the NLP transformer block you already know: layer normalisation, multi-head self-attention, a residual (skip) connection, a second layer normalisation, a two-layer MLP with a GELU activation, and a second residual connection. ViT-Base stacks twelve of these blocks; the [CLS] token's representation after the final block is fed to a small MLP head that outputs class scores.

How the pieces connect

Correcting a common misconception

A student who has just learned self-attention for text often assumes that because attention lets every patch "see" every other patch, ViT must automatically know the spatial layout of the image — which patch is above which, which is to the left of which — as an inherent property of the mechanism. It does not. Self-attention, on its own, is permutation-equivariant: it treats its input as an unordered set of tokens. If you shuffled the order of the patch embeddings e_A, e_B, e_C, e_D before feeding them in, and shuffled the output the same way afterward, you would get back the exact same set of output vectors. Attention has no concept of "next to" or "above" baked into its formula; the dot product Q_A · K_B in the worked example above is computed exactly the same way whether B is A's immediate neighbour or on the opposite corner of the image.

This is precisely why the worked example added distinct position vectors P_A, P_B, P_C, P_D before running attention, rather than feeding the raw patch embeddings e_A, e_B, e_C, e_D directly. Recall that those raw embeddings all lay on the line y = x: a patch with intensity 0.2 produces embedding [0.4, 0.4] regardless of whether it sits at the top-left or the bottom-right of the image. Without position embeddings, a ViT literally cannot distinguish "dark patch in the top-left corner" from "dark patch in the bottom-right corner" — it only knows patch content, never patch location. A CNN never has this problem, because convolution is applied at fixed pixel coordinates by construction; spatial position is implicit in where the kernel slides. ViT has to be told position explicitly, as an extra additive signal, precisely because its core operation is otherwise blind to it. This is also the underlying reason ViT is described as having "weaker inductive bias" than a CNN: locality and spatial order are architectural assumptions in a CNN, but in a ViT they are, at best, patterns the position embeddings and attention weights must learn from data.

What ViT trades away, and what it buys

• Receptive field. A CNN's receptive field grows layer by layer, from a few pixels near the input to the whole image only after many layers. Every ViT layer already has a global receptive field — any patch can attend to any other patch from layer one — which is exactly what resolves the ambiguous grey patch in the flood-mapping example without waiting for depth to accumulate.
• Inductive bias versus data. A CNN's convolution kernel hard-codes the assumption that nearby pixels are related and that a pattern useful in one location is useful in another (translation equivariance). ViT has neither assumption built in; it must discover useful spatial structure purely from training examples. This is why the original ViT paper found it underperforms CNNs of similar size when trained only on ImageNet-scale data (roughly 1.3 million images), but overtakes them once pretrained on much larger datasets (ImageNet-21k, 14 million images, or the internal JFT-300M, 300 million images) — with enough data, the model learns spatial structure that a CNN gets for free architecturally.
• Compute scaling with patch count. Self-attention over N tokens costs O(N² · D) per layer, since it forms an N × N score matrix. Halving the patch size quadruples N (patch 16 on a 224×224 image gives N = 196; patch 8 would give N = 784), so attention cost rises by roughly 16×, not 4×. This quadratic term is why ViT patch sizes are chosen coarser than a CNN's finest convolution kernels — 16×16 or 32×32 patches, not 3×3 — trading fine-grained spatial resolution for tractable compute.

Patch embedding as a single convolution, in code

In practice, no implementation flattens patches and multiplies by a weight matrix as two separate steps — that computation is exactly equivalent to a single strided convolution, and frameworks implement it that way for speed. A convolution with kernel size equal to the patch size and stride equal to the patch size slides across the image without overlap, and each kernel position computes precisely the flatten-then-project operation from the worked example, once per patch, in parallel:

import torch
import torch.nn as nn

class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        self.n_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_chans, embed_dim,
                               kernel_size=patch_size, stride=patch_size)

    def forward(self, x):
        # x: (batch, in_chans, img_size, img_size)
        x = self.proj(x)              # (batch, embed_dim, 14, 14)
        x = x.flatten(2)              # (batch, embed_dim, 196)
        x = x.transpose(1, 2)         # (batch, 196, embed_dim)
        return x                      # ready to prepend [CLS] and add position embeddings

embedder = PatchEmbed()
print(embedder.n_patches)   # 196, matches N = (224/16)^2 computed above

Trace the shapes: a 224×224×3 image enters proj, a convolution with a 16×16 kernel and stride 16, so it slides to 224/16 = 14 positions along each axis, producing a (batch, 768, 14, 14) tensor — one 768-dimensional vector per patch, arranged on a 14×14 grid. flatten(2) collapses the two spatial axes into one of size 14 × 14 = 196, and transpose(1, 2) puts the sequence dimension before the feature dimension, giving (batch, 196, 768) — the shape the transformer encoder expects, and exactly the N = 196 derived by hand earlier in this chapter.

Active recall

1. An image is 384×384 with 3 channels, patch size P = 16. Compute the number of patches N and the length of each patch's raw flattened vector before projection.
2. In the worked example, patch A's attention weight on itself (0.187) was lower than its weight on patch D (0.294). Using the raw dot-product scores computed in the chapter, explain why, referencing what makes a key vector "align" with a query vector under the dot product.
3. If you removed the position embeddings from the worked example entirely (used z = e for every token), would the attention weights for query A change? Justify using the actual embedding values e_A, e_B, e_C, e_D.
4. Doubling the patch size from 16 to 32 on a 224×224 image changes N from 196 to 49. By what factor does the per-layer self-attention compute roughly change, and why is that factor not simply 4?
5. Why does a Vision Transformer typically need a larger pretraining dataset than a comparably-sized CNN to reach the same accuracy, in terms of what each architecture assumes about images before training even starts?
6. In the ISRO flood-mapping scenario, state in one sentence why a global receptive field resolves the ambiguous grey-patch problem faster (in terms of network depth) than a stack of small convolutions would.

Answers.

1. N = (384/16) × (384/16) = 24 × 24 = 576 patches. Each raw patch vector has length 16 × 16 × 3 = 768.
2. The raw score is the dot product Q_A · K_j. Q_A = [0.2, 0.3]. K_A = [0.2, 0.3] gives dot product 0.13; K_D = [1.6, 1.5] gives 0.2×1.6 + 0.3×1.5 = 0.32 + 0.45 = 0.77. Even though D is a completely different token, its coordinates are both larger than A's own coordinates in the same direction, so the dot product with D exceeds the dot product with A's own (smaller-magnitude) vector. Self-attention rewards direction and magnitude alignment, not identity — a token is not guaranteed to be its own best match.
3. No — the four raw embeddings are e_A=[0.2,0.2], e_B=[1.8,1.8], e_C=[0.4,0.4], e_D=[1.6,1.6], all distinct in magnitude (0.2, 1.8, 0.4, 1.6 respectively), so attention scores would still differ token to token; the weights would change numerically. What removing position embeddings destroys is not variation but location information: since e_A and e_C both lie on y=x and differ only by being "small" versus "small-ish," and likewise B and D differ only by magnitude, the model would attend based purely on how dark or bright a patch is, with no way to know which corner of the image that patch came from. Two images with the same four intensity values placed in different corners would produce identical attention patterns.
4. Attention cost per layer is O(N² · D). N drops from 196 to 49, a factor of 4, but the cost depends on N², so the compute factor is 4² = 16, not 4. Patch size acts on both the height and width axes of the patch grid simultaneously, so its effect on token count — and therefore on quadratic attention cost — is squared.
5. A CNN's convolution operation already assumes that nearby pixels are related and that useful patterns repeat across locations (translation equivariance); those are correct assumptions for almost all natural images, so the CNN starts training with useful structure already built in. A ViT's self-attention makes no such assumption — every patch pair starts equally "reachable" — so the network must learn from examples which patches are actually worth attending to, which takes more data to get right.
6. Every ViT layer gives each patch a direct, single-hop connection to every other patch in the image, so evidence from a distant confirmed-dry or confirmed-river patch can influence the ambiguous patch's representation in one layer, whereas a CNN needs its receptive field to grow across many stacked convolution and pooling layers before it can combine information from patches that far apart.

Think About It

Think about this: How would you explain vision transformers: image understanding with attention 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.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind vision transformers: image understanding with attention, 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.