🧩 Vision Transformer (ViT)

An Image is Worth 16×16 Words: applying the transformer architecture directly to sequences of image patches

Table of Contents

• 🏛️ Background: From CNNs to Attention
• 🔲 Image Tokenization: Patches and Projections
• 📌 The CLS Token and Positional Encoding
• 🔍 Multi-Head Self-Attention on Patches
• 🧱 The ViT Encoder
• ⚖️ Inductive Bias: ViT vs. CNN
• 📈 Training Regime and Empirical Results
• 🔗 Connections: DeiT, Swin, and SSL
• 📚 References

🏛️ Background: From CNNs to Attention

For over a decade, convolutional neural networks were the uncontested default for vision. CNNs exploit two strong inductive biases:

1. Locality: Each neuron in layer \(\ell\) depends only on a small spatial neighborhood in layer \(\ell - 1\). Features are local.
2. Translation equivariance: The same filter is applied at every spatial location. If a cat shifts right by 3 pixels, the feature map shifts identically.

These biases make CNNs highly sample-efficient — they need fewer images to learn that a cat’s ear looks the same regardless of position. But they also constrain what the network can represent: long-range dependencies (e.g., matching a face to the body below it) are expensive, requiring many stacked layers to grow the effective receptive field.

Transformers, introduced for NLP by Vaswani et al. (2017), have no built-in locality or translation equivariance. Every token attends to every other token in a single layer. This makes them maximally expressive for long-range structure — but data-hungry, since spatial relationships must be learned entirely from data.

The ViT question: Can you apply a standard transformer encoder without modification to images, by treating image patches as tokens? Dosovitskiy et al. (2020) answer: yes, and at sufficient data scale it outperforms CNNs by a wide margin.

🔲 Image Tokenization: Patches and Projections

📐 From Pixels to Tokens

Let \(x \in \mathbb{R}^{H \times W \times C}\) be an image (\(H\) height, \(W\) width, \(C\) channels). Choose a patch size \(P \times P\). Reshape \(x\) into a sequence of \(N\) flattened patches:

\[x^{(i)}_p = \mathrm{flatten}(\mathrm{patch}_i(x)) \in \mathbb{R}^{P^2 C}, \quad i = 1, \ldots, N, \quad N = \frac{HW}{P^2}.\]

Each patch is a \(P^2 C\)-dimensional vector. For a standard 224×224 RGB image with \(P = 16\):

\[N = \frac{224 \times 224}{16^2} = 196 \text{ tokens}.\]

Linear projection to embedding space. A learned linear map \(E \in \mathbb{R}^{D \times P^2C}\) projects each patch to the model’s hidden dimension \(D\):

\[z^{(i)}_0 = E\, x^{(i)}_p \in \mathbb{R}^D.\]

This is the patch embedding — the ViT analogue of the word embedding in NLP transformers. The projection \(E\) is learned end-to-end; its columns visualize as Fourier-like basis functions resembling oriented Gabor filters after training.

📌 The CLS Token and Positional Encoding

🏷️ Classification Token

Before feeding the patch sequence to the transformer, ViT prepends a learnable classification token \(x_{\mathrm{cls}} \in \mathbb{R}^D\):

\[\mathbf{z}_0 = [x_{\mathrm{cls}};\; z^{(1)}_0;\; z^{(2)}_0;\; \ldots;\; z^{(N)}_0] \in \mathbb{R}^{(N+1) \times D}.\]

After \(L\) transformer layers, the representation at position 0 (the CLS token’s slot) is used as the image representation for classification: \(y = \mathrm{LN}(\mathbf{z}_L^{[0]})\).

The CLS token has no spatial meaning — it contains no image content at initialization. Through training, it learns to aggregate global information from all patch tokens via attention. It acts as a learned readout vector: a query that learns what to ask from the image.

📍 Positional Encoding

A pure set of patch embeddings \(\{z^{(i)}_0\}\) contains no information about where each patch came from — the transformer is permutation-equivariant by default. Positional encodings fix this:

\[\mathbf{z}_0 \leftarrow \mathbf{z}_0 + E_{\mathrm{pos}}, \quad E_{\mathrm{pos}} \in \mathbb{R}^{(N+1) \times D}.\]

ViT uses 1D learnable positional encodings — a learnable matrix where row \(i\) encodes the absolute position of the \(i\)-th token. Ablation studies in the paper show:

Surprisingly, 2D positional encodings (which directly encode the \((r, c)\) grid position) offer almost no benefit over 1D encodings. The model appears capable of learning implicit 2D structure from 1D position tags alone.

🔍 Multi-Head Self-Attention on Patches

The transformer’s key operation is multi-head self-attention (MHSA). Here we derive it in the ViT context.

📐 Single-Head Attention

Given a sequence \(Z \in \mathbb{R}^{(N+1) \times D}\), project to queries, keys, and values:

\[Q = Z W_Q, \quad K = Z W_K, \quad V = Z W_V, \quad W_Q, W_K, W_V \in \mathbb{R}^{D \times d_k}.\]

The attention output is:

\[\mathrm{Attn}(Q, K, V) = A V, \quad A = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) \in \mathbb{R}^{(N+1) \times (N+1)}.\]

The attention matrix \(A\) has entry $A_{ij} = $ (normalized inner product of query \(i\) with key \(j\)), representing how much token \(i\) attends to token \(j\). The \(\sqrt{d_k}\) scaling prevents the dot products from growing too large in magnitude, which would saturate the softmax into near-zero gradients.

📐 Multi-Head Attention

Run \(h\) independent attention heads in parallel, each with projection dimension \(d_k = D/h\), then concatenate and project:

\[\mathrm{head}_i = \mathrm{Attn}(Z W^i_Q, Z W^i_K, Z W^i_V), \quad i = 1, \ldots, h\]

\[\mathrm{MHSA}(Z) = \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_h)\, W^O, \quad W^O \in \mathbb{R}^{D \times D}.\]

With \(h\) heads each of dimension \(d_k = D/h\), the total parameter count is \(4D^2\) (three \(D \times d_k\) projections per head plus \(W^O\)), independent of \(h\).

Why multiple heads? Each head can specialize to attend to different structural relationships simultaneously. In ViT, empirically some heads attend to local structure (nearby patches) while others attend globally (semantically related patches far apart) — even within the same layer.

⏱️ Computational Complexity

For sequence length \(N\) and dimension \(D\), the attention complexity is \(O(N^2 D)\) — quadratic in the sequence length. This is the key bottleneck of standard transformers.

For ViT with \(P=16\) patches: \(N = 196\), making \(N^2 = 38{,}416\) — manageable. For \(P=8\): \(N = 784\), \(N^2 \approx 600{,}000\) — already expensive. This motivates efficient attention variants for dense tasks.

🧱 The ViT Encoder

The full ViT architecture stacks \(L\) identical transformer encoder blocks. Each block applies two sublayers with pre-normalization (LayerNorm before, residual connection after):

\[\mathbf{z}'_\ell = \mathrm{MHSA}(\mathrm{LN}(\mathbf{z}_{\ell-1})) + \mathbf{z}_{\ell-1} \tag{attention sublayer}\]

\[\mathbf{z}_\ell = \mathrm{MLP}(\mathrm{LN}(\mathbf{z}'_\ell)) + \mathbf{z}'_\ell \tag{MLP sublayer}\]

The MLP has two linear layers with a GELU activation and hidden dimension \(4D\):

\[\mathrm{MLP}(x) = W_2\, \mathrm{GELU}(W_1 x + b_1) + b_2, \quad W_1 \in \mathbb{R}^{4D \times D},\; W_2 \in \mathbb{R}^{D \times 4D}.\]

After \(L\) blocks, extract the CLS token and apply a final LayerNorm:

\[y = \mathrm{LN}(\mathbf{z}_L^{[0]}) \in \mathbb{R}^D.\]

For classification, \(y\) is projected by a linear head: \(\hat{p} = \mathrm{softmax}(Wy)\).

🏗️ Architecture Diagram

flowchart TD
    img["Input Image
H x W x C"]
    patch["Patch Embedding
N patches, each P^2C dim"]
    proj["Linear Projection E
N x D embeddings"]
    cls["Prepend CLS token
(N+1) x D"]
    pos["Add Position Encoding E_pos"]
    block["Transformer Block x L
LN -> MHSA -> residual
LN -> MLP -> residual"]
    clsout["Extract z_L^0
CLS token output"]
    ln["LayerNorm"]
    head["Linear Head
D -> num_classes"]
    out["Class Logits"]

    img --> patch --> proj --> cls --> pos --> block --> clsout --> ln --> head --> out

📊 Model Variants

The /P suffix denotes the patch size. ViT-B/16 is the workhorse; ViT-H/14 (14×14 patches for finer granularity) is the largest and best-performing variant.

⚖️ Inductive Bias: ViT vs. CNN

This comparison is the conceptual core of the ViT paper.

🔑 What CNNs Bake In

A convolutional layer is equivariant to translation by construction: \[f(T_\delta x) = T_\delta f(x)\] where \(T_\delta\) denotes a spatial shift by \(\delta\). Additionally, the limited receptive field enforces locality — features at position \((i, j)\) only depend on a small neighborhood.

These inductive biases are correct priors for natural images (cats look the same regardless of position), which makes CNNs extremely sample-efficient: ResNets learn strong features from ImageNet-1k (1.2M images) without any pretraining.

🔑 What ViT Lacks

ViT has no built-in locality or translation equivariance. The patch embedding \(E\) applies the same projection to all patches (a weak form of equivariance), but the positional encoding \(E_\text{pos}\) breaks it immediately — position-\(i\) and position-\(j\) patches are treated differently even if their content is identical.

Consequence: ViT must learn that nearby patches tend to be correlated, that the same feature can appear at different positions, and that spatial relationships are meaningful. With limited data, it may not learn these properties reliably.

📊 The Data Requirement Gap

Without large-scale pretraining, ViT-L/16 underperforms a comparably sized ResNet. This is the inductive bias gap: CNNs encode the right priors for free; ViT must learn them from data.

The payoff at scale: With 300M pretraining images, ViT matches and eventually exceeds the best CNN baselines — achieving better accuracy with \(2\text{–}4\times\) less compute at inference than comparably accurate ResNets. The learned representations are also more transferable: ViT pretrained on JFT shows stronger few-shot transfer to diverse downstream tasks.

📈 Training Regime and Empirical Results

🏋️ Pretraining and Fine-Tuning

ViT is trained in two stages:

1. Pretraining on a large dataset (ImageNet-21k or JFT-300M) with a classification head, using AdamW optimizer, learning rate warmup, and weight decay regularization.

2. Fine-tuning on the target dataset at higher resolution. The sequence length grows (e.g., \(N = 784\) at \(448 \times 448\) vs. \(N = 196\) at \(224 \times 224\)), and position embeddings are 2D-interpolated. The pretraining classification head is replaced by a zero-initialized linear head for the target classes.

📊 What Attention Learns

A striking empirical finding is the attention distance analysis. For each attention head, compute the average distance (in pixels) between a token and the tokens it attends to most heavily. Plots from the paper reveal:

• Lower layers: Some heads attend very locally (distance \(< 20\) pixels), others globally (full image distance). High variance across heads.
• Higher layers: Attention distance uniformly increases — all heads become global. The model integrates local features (lower layers) into global semantics (upper layers), analogously to the expanding receptive field in CNNs.

The key difference from CNNs: ViT develops global heads from the very first layer, allowing early layers to directly reason about cross-image relationships. In CNNs, global receptive field requires stacking many layers.

🔗 Connections: DeiT, Swin, and SSL

🎓 DeiT: Data-Efficient Training

ViT’s large data requirement is a practical barrier. Touvron et al. (DeiT, 2021) show that ViT can be trained competitively on ImageNet-1k alone via: - Strong augmentation: Mixup, CutMix, RandAugment, random erasing - Distillation token: A second special token (alongside CLS) that is trained to match the output of a CNN teacher, enabling knowledge distillation without large-scale data

DeiT-B/16 matches ViT-B/16 (JFT-pretrained) on ImageNet using only ImageNet-1k. The distillation token learns to attend to features useful for the CNN teacher’s predictions — complementing the class token.

🪟 Swin Transformer: Restoring Locality

Liu et al. (Swin, 2021) reintroduce locality and hierarchy into ViT by computing attention within local windows (e.g., \(7 \times 7\) patch windows) and using a shifted window scheme across layers to allow cross-window interactions:

• Window attention: \(O(N)\) cost instead of \(O(N^2)\) — attention within a \(7 \times 7\) window is \(O(49^2)\) per window, with \(N/49\) windows total.
• Hierarchical feature maps: Like CNNs, Swin progressively merges patches to produce multi-scale features, making it directly compatible with dense prediction heads (FPN, Mask R-CNN).

Swin became the dominant backbone for detection and segmentation (replacing ResNets) while ViT remained dominant for classification.

🔗 ViT and Self-Supervised Learning

ViT architecture interacts particularly well with SSL objectives:

DINO (Caron et al., 2021): Applies BYOL-like self-distillation to ViT. The CLS token of the teacher (EMA network) provides the target; the student’s CLS token must match it. Remarkably, the attention maps of DINO-pretrained ViTs cleanly segment foreground objects with no segmentation supervision — a property unique to ViT + SSL (CNNs trained with DINO do not exhibit this).

MAE (He et al., 2022): Masks 75% of patches (random) and trains the encoder to reconstruct the masked patches via a lightweight decoder. The high masking ratio forces the encoder to develop a global understanding — a single visible patch must predict masked content across the image. MAE uses ViT exclusively; the masked-input framework is incompatible with CNNs (which cannot skip spatial locations).

See Self-Supervised Vision and SSL Theory for the theoretical foundations of these SSL methods.

📚 References