Vision Transformers

A Vision Transformer (ViT) applies the Transformer encoder architecture, originally developed for sequences of word tokens, directly to sequences of image patch tokens. An input image of resolution $H \times W \times C$ is partitioned into a grid of non-overlapping patches of size $P \times P$, producing $N = HW / P^2$ patches. Each patch is flattened into a vector of length $P^2 C$ and linearly projected into a $D$-dimensional embedding — this projection is called patch tokenization or patch embedding, and it is the entire "vision-specific" part of the architecture. Everything downstream is generic Transformer machinery. Because a Transformer's self-attention is permutation-invariant — it treats its input as an unordered set of tokens — spatial position must be reinjected explicitly. ViT adds a learnable (or fixed sinusoidal) positional embedding to each patch embedding, so the model can, if useful, recover 2D structure such as "this patch is above that one." A special learnable [CLS] token is prepended to the sequence; its output embedding after the final layer is used as the aggregate image representation for classification, analogous to its role in BERT for sentence-level tasks. The heart of the architecture is multi-head self-attention (MHSA). For a sequence of $N{+}1$ token embeddings (patches plus CLS), each layer computes queries, keys, and values via learned linear maps, then lets every token attend to every other token according to a softmax-normalized similarity of queries and keys. This is fundamentally different from a convolution, which only mixes information within a small, fixed spatial neighborhood and only grows its receptive field by stacking layers. Self-attention gives every layer a receptive field covering the entire image, at the cost of quadratic compute in the number of tokens and the loss of built-in locality and translation equivariance — properties CNNs get for free.

Step 1 — Patchify. Split the $H \times W \times C$ image into $N = (H/P)(W/P)$ non-overlapping $P \times P$ patches. Flatten each into a vector of dimension $P^2 C$.

Step 2 — Linear patch embedding. Multiply each flattened patch by a learned matrix $E \in \mathbb{R}^{(P^2C) \times D}$ to obtain a $D$-dimensional token. In practice this is implemented as a single strided convolution with kernel size $P$ and stride $P$ — mathematically identical to patch-then-flatten-then-linear, but faster in standard deep learning frameworks.

Step 3 — Prepend CLS token, add positional embeddings. Prepend a learnable vector $x_{cls}$, then add a learnable positional embedding $E_{pos} \in \mathbb{R}^{(N+1)\times D}$ elementwise to the $(N{+}1)$ tokens.

Step 4 — Transformer encoder. Pass the sequence through $L$ identical blocks, each consisting of LayerNorm, multi-head self-attention with a residual connection, LayerNorm, and a two-layer MLP (with GELU activation) with a residual connection.

Step 5 — Classification head. Take the final-layer CLS token embedding, apply LayerNorm, and pass through a small MLP (or single linear layer during fine-tuning) to produce class logits. ``` function VisionTransformer(image, params):

# image: H x W x C tensor patches = split_into_patches(image, patch_size=P) # N patches, each PPC patch_embeddings = patches @ E # N x D, E is learned P^2C x D tokens = concat(CLS_token, patch_embeddings) # (N+1) x D tokens = tokens + positional_embeddings # (N+1) x D, elementwise add for layer in 1..L:

# --- Multi-head self-attention sub-block --- normed = LayerNorm(tokens) Q, K, V = normed @ Wq, normed @ Wk, normed @ Wv # each (N+1) x D Q, K, V = split_into_heads(Q, K, V, num_heads=h) # h heads of dim D/h for each head: scores = (Q_head @ K_head^T) / sqrt(D/h) # (N+1) x (N+1) attn_weights = softmax(scores, axis=-1) head_out = attn_weights @ V_head # (N+1) x (D/h) attn_out = concat(all head_out) @ Wo # (N+1) x D tokens = tokens + attn_out # residual connection

# --- MLP sub-block --- normed2 = LayerNorm(tokens) mlp_out = GELU(normed2 @ W1 + b1) @ W2 + b2 tokens = tokens + mlp_out # residual connection cls_final = LayerNorm(tokens[0]) # CLS token row logits = cls_final @ W_head + b_head return logits ``` Every step above is differentiable and trained end-to-end with a standard classification loss (cross-entropy against image labels), typically after large-scale pretraining and then fine-tuning at higher resolution on the target task.

Time complexity. Let $N$ be the number of patch tokens (plus one for CLS), $D$ the embedding dimension, and $L$ the number of encoder layers. Each self-attention sub-block computes $Q$, $K$, $V$ projections in $O(ND^2)$, the pairwise attention scores $QK^\top$ in $O(N^2 D)$, and the weighted sum with $V$ in another $O(N^2 D)$. The MLP sub-block, typically expanding to $4D$ hidden units, costs $O(ND^2)$. Summed over $L$ layers, total time is

The $N^2 D$ term dominates when the number of patches is large relative to the embedding width — this is the well-known quadratic bottleneck of self-attention. For a $224\times224$ image with $16\times16$ patches, $N = 196$, small enough that the $ND^2$ term (linear projections) often dominates in practice; but at higher resolutions or with smaller patches, $N$ grows quadratically with image side length ($N \propto (H/P)^2$), so overall attention cost scales as the fourth power of image resolution — a serious practical constraint that motivates windowed and hierarchical attention variants.

Space complexity. Storing the full attention matrix per head requires $O(N^2)$ per layer (before summing over heads it is $O(hN^2)$ split across $h$ heads of size $N \times N$ each, same asymptotic total), plus $O(ND)$ for activations. Memory is dominated by the $N^2$ attention matrices during training when backpropagating through softmax, which is why memory-efficient attention implementations (recomputing rather than storing the full matrix) matter for large images.

Correctness / why it learns the right thing. There is no combinatorial correctness proof for a learned model the way there is for a sorting algorithm; the relevant argument is representational and statistical. Representational sufficiency: self-attention with a residual stream can represent, in principle, any function a convolution can — a convolution is a special case of a sparse, weight-shared attention pattern restricted to a local window with fixed weights. A sufficiently expressive attention layer can therefore learn to emulate local, translation-equivariant filtering if that is what the data demands, and the model additionally has access to genuinely global interactions from layer one. This is a strict superset of representational power, at the cost of a much larger hypothesis space to search. Why more data is needed: because CNNs bake in the assumption "nearby pixels matter most, and the same filter should work everywhere," they need fewer examples to find good filters — the inductive bias substitutes for data. ViT drops that assumption, so early layers must learn from data that nearby pixels are usually more relevant. On small datasets (e.g. ImageNet-1k trained from scratch), ViT underperforms comparable CNNs because it has not seen enough examples to discover these regularities reliably. On large datasets (JFT-300M-scale pretraining, or heavy augmentation/distillation as in DeiT), ViT matches or exceeds CNNs, because the extra capacity that was previously "wasted" relearning locality can instead be spent on patterns convolution could never represent efficiently, such as long-range part relationships. Position embeddings encode 2D structure empirically: since self-attention is permutation-invariant, the only signal telling the model which patch is where is the positional embedding. Empirically, learned positional embeddings in trained ViTs exhibit smooth 2D similarity structure — nearby patches acquire similar embeddings — even though this structure is never explicitly imposed, confirming that the model recovers spatial locality as an emergent, learned property rather than a hard-coded one.

Tradeoffs, stated plainly. ViT trades built-in locality and translation equivariance for flexibility and scalability; it trades better small-data sample efficiency for better large-data ceiling performance; and it trades linear-in-pixels convolutional cost for quadratic-in-tokens attention cost, which becomes the binding constraint at high resolution. None of these are free lunches — the right choice between CNN and ViT backbones still depends on data scale, resolution, and latency budget.

Trace a single forward pass on a small, concrete configuration to see the shapes and one attention computation explicitly. Setup: image $32\times32\times3$, patch size $P=16$, so $N = (32/16)\times(32/16) = 4$ patches, embedding dimension $D = 6$ (tiny, for hand tracing), 1 attention head, 1 encoder layer.

Concretely, suppose after LayerNorm the query vector for the CLS token is $q_{cls} = [0.1, -0.2, 0.4, 0.0, 0.3, -0.1]$ and the four patch key vectors, after the learned $W_k$, happen to be

Computing dot products $q_{cls}\cdot k_i$ and dividing by $\sqrt{6}\approx 2.449$ gives four raw scores, say $\{0.06, -0.05, 0.09, 0.14\}$ (illustrative values consistent with the vectors above). Softmax over these four numbers (plus the CLS-to-CLS score) produces attention weights that are nearly uniform but slightly favor patch 4 (the largest score), e.g. weights approximately $\{0.198, 0.191, 0.198, 0.207, 0.206\}$ across CLS-self, patch 1–4. The CLS token's new representation is then this weighted average of all five value vectors — meaning even at a single layer, the classification token has already blended information from every patch in the image, something a single convolutional layer with a small kernel could never do. Stacking more layers lets this mixing become increasingly selective, as the query/key projections are trained to sharpen attention toward semantically relevant patches (e.g., patches containing the object rather than background). This trace exposes the mechanism's essence: attention weights are not fixed like a convolution kernel, they are data-dependent, recomputed per image from the learned query/key projections, and they connect arbitrarily distant tokens in one hop.

The Transformer architecture itself originated in sequence modeling for machine translation, replacing recurrence with pure attention. Applying it to raw image patches without convolution required the empirical demonstration that self-attention could match or beat CNNs given sufficient scale — a result that surprised much of the vision community, which had spent a decade optimizing convolutional inductive biases.

ViT backbones now sit inside a wide range of production and research systems. Image classification leaderboards are dominated by ViT and its hierarchical descendants, often pretrained at massive scale and fine-tuned per task. Object detection and segmentation frameworks that once assumed a convolutional backbone (building on the same feature-pyramid ideas discussed in Single-Shot Detectors) now commonly swap in ViT-style backbones, sometimes adapting them to produce multi-scale feature maps for dense prediction. Multimodal systems such as CLIP use a ViT as the image encoder half of a joint image-text embedding space, aligning patch-derived image representations with text embeddings for zero-shot classification and retrieval. Vision-language assistants and generative image models frequently reuse ViT-style patch tokenization as the bridge between pixel grids and Transformer-based reasoning or diffusion architectures. Because ViT's patch-token formulation requires no labels to define — only the raw pixels and their spatial layout — it is also the natural backbone for Self-Supervised Vision methods that learn representations via masked patch prediction or contrastive objectives, extending the same architecture built here far beyond supervised classification.

:::quiz[What is the primary role of the [CLS] token in a Vision Transformer?]

• It marks the boundary between image patches during tokenization
• Its final-layer embedding aggregates information for classification ✓
• It stores the positional embedding for patch zero
• It is discarded before the encoder and only used at input time

Ask the tutor: How does windowed (hierarchical) attention in models like Swin Transformer reduce the quadratic cost while still capturing global context?