Complete Vision Transformer

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Hard

Problem

Theory

The Vision Transformer (ViT) applies a standard Transformer encoder directly to sequences of image patches for image classification. Rather than using convolutions, it splits an image into fixed-size patches, linearly embeds each one, and processes the resulting sequence with a Transformer. This is the full end-to-end pipeline: patch embedding, CLS token, position embeddings, L encoder blocks, and a classification head.

What It Is

ViT is an image classification architecture that replaces convolutional feature extraction with a Transformer encoder. The input image is divided into non-overlapping patches, each flattened and linearly projected to produce a patch embedding. A learnable [CLS] token is prepended, learnable position embeddings are added, and the full sequence passes through L encoder blocks (pre-LN, bidirectional attention, GELU MLP). After the final block, the [CLS] representation is extracted, LayerNormed, and linearly projected to class logits. No softmax is applied.

The architecture is deliberately minimal. Dosovitskiy et al. made as few modifications to the standard Transformer as possible, proving that a pure Transformer without any convolutional inductive bias can achieve state-of-the-art image classification when trained on sufficient data.

Key Equations

Stage 1 -- Patch embedding. Split the image into N patches of size P \times P. Flatten each patch to length P^2 \cdot C, then project:

z_i^{(0)} = x_{\text{patch}_i} \cdot W_E + b_E, \quad i = 1, \ldots, N

where W_E \in \mathbb{R}^{(P^2 C) \times D} and D is the hidden dimension.

Stage 2 -- Prepend [CLS] and add position embeddings:

z^{(0)} = [z_{\text{cls}};\; z_1^{(0)};\; \ldots;\; z_N^{(0)}] + E_{\text{pos}}

where z_{\text{cls}} \in \mathbb{R}^D is learnable and E_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}.

Stage 3 -- L encoder blocks (pre-LN):

z'^{(\ell)} = \text{MSA}(\text{LN}(z^{(\ell-1)})) + z^{(\ell-1)}

z^{(\ell)} = \text{MLP}(\text{LN}(z'^{(\ell)})) + z'^{(\ell)}

Stage 4 -- Classification head:

y = \text{LN}(z_0^{(L)}) \cdot W_{\text{head}} + b_{\text{head}}

where z_0^{(L)} is the CLS token after block L, and y \in \mathbb{R}^K is raw logits.

Stage 1: Patch Embedding

The image has shape (C, H, W). It is divided into a grid of non-overlapping P \times P patches. The number of patches is N = HW / P^2. For a 224x224 image with $ = 16$: N = 196. Each patch is flattened to length P^2 C = 768 (for RGB), then projected through W_E \in \mathbb{R}^{(P^2 C) \times D} with bias.

Why patches instead of pixels? A 224x224 image has 50,176 pixels. Self-attention is O(N^2), making per-pixel tokenization prohibitive. Patch size 16 reduces this to 196 tokens. Smaller patches give finer resolution but quadratically increase cost. The ViT naming convention reflects this: ViT-B/16 means Base model with patch size 16.

In practice, this linear projection is equivalent to a 2D convolution with kernel size P and stride P. All weights are initialized from \mathcal{N}(0, 0.02^2).

Stage 2: CLS Token and Position Embeddings

The [CLS] token. A learnable vector z_{\text{cls}} \in \mathbb{R}^D is prepended, increasing sequence length from N to N + 1. It has no spatial content and serves as a global aggregation point. Because attention is bidirectional, CLS can attend to all patches and vice versa. After L blocks, the CLS representation encodes the entire image. This follows BERT's design.

Position embeddings. Learnable E_{\text{pos}} \in \mathbb{R}^{(N+1) \times D} are added element-wise. Each position gets its own D-dimensional vector, encoding spatial layout: position 1 is the top-left patch, position 2 the next, and so on. Without these, the Transformer treats patches as an unordered set and loses all spatial information.

The paper found that learned 1D embeddings perform as well as more complex 2D-aware schemes. The authors tested fixed sinusoidal, learned 1D, and learned 2D embeddings, finding no significant accuracy difference. The Transformer learns to infer 2D spatial relationships from 1D indices.

Stage 3: Transformer Encoder

The sequence z^{(0)} \in \mathbb{R}^{(N+1) \times D} passes through L identical encoder blocks. Each has two sub-layers: multi-head self-attention (MSA) and a feed-forward MLP, both with pre-LayerNorm and residual connections.

Multi-head self-attention. Input is normalized: \hat{z} = \text{LN}(z^{(\ell-1)}). Queries, keys, values are computed, split into h heads of dimension d_h = D/h. Each head computes:

\text{head}_j = \text{softmax}\!\left(\frac{Q_j K_j^T}{\sqrt{d_h}}\right) V_j

Heads are concatenated and projected: \text{MSA} = [\text{head}_1; \ldots; \text{head}_h] W_O. There is no causal mask. Every token attends to every other token bidirectionally. Image patches have no sequential ordering, so masking future positions would create arbitrary asymmetry.

Feed-forward MLP. After the attention residual, the output is normalized and passed through:

\text{MLP}(x) = \text{GELU}(x W_1 + b_1) W_2 + b_2

where W_1 \in \mathbb{R}^{D \times 4D} and W_2 \in \mathbb{R}^{4D \times D}. The activation is GELU, not ReLU.

Pre-LN ordering means LayerNorm is applied before each sub-layer, not after. This stabilizes training for deeper models because the residual path carries unnormalized gradients directly.

Stage 4: Classification Head

Extract the [CLS] token from position 0 of the final block output: z_0^{(L)} \in \mathbb{R}^D. Apply LayerNorm, then project:

y = \text{LN}(z_0^{(L)}) \cdot W_{\text{head}} + b_{\text{head}}

where W_{\text{head}} \in \mathbb{R}^{D \times K} and K is the number of classes. Output is raw logits (no softmax). During training, logits go to cross-entropy loss which applies log-softmax internally.

During pre-training on large datasets, the paper uses an MLP head with tanh activation. During fine-tuning, this is replaced with a single linear layer. This problem uses the fine-tuning variant: LayerNorm then linear. All head weights are initialized as \mathcal{N}(0, 0.02^2).

The Data Requirement

The central empirical finding: pure Transformers need large-scale pre-training to match CNNs. On ImageNet-1K alone (1.3M images), ViT-B/16 underperforms a comparably sized ResNet. The Transformer lacks convolutional inductive biases: translation equivariance and locality. It must learn spatial structure entirely from data.

As the paper states: "When trained on mid-sized datasets, these models yield modest accuracies. However, the picture changes if trained on larger datasets." On JFT-300M (303M images), ViT-L/16 surpasses the best CNNs on every benchmark.

• Small data: CNN inductive biases act as regularizers, reducing the hypothesis space. ViT must learn these patterns from scratch.
• Large data: Inductive biases become constraints. Self-attention can learn arbitrary spatial relationships that CNNs only capture through many stacked layers. With enough data, flexibility wins.
• Compute: ViT maps directly onto hardware optimized for matrix multiplications. ViT-H/14 reached 88.55% top-1 on ImageNet using fewer TPU-days than competing models.

Paper Context

"An Image is Worth 16x16 Words" by Dosovitskiy et al. (2020, Google Brain) showed that a pure Transformer applied directly to image patches can match or exceed convolutional networks at scale.

Three model configurations:

• ViT-Base: L = 12, D = 768, h = 12 heads, MLP = 3072, ~86M params.
• ViT-Large: L = 24, D = 1024, h = 16 heads, MLP = 4096, ~307M params.
• ViT-Huge: L = 32, D = 1280, h = 16 heads, MLP = 5120, ~632M params.

The design philosophy was deliberate minimalism. Prior work combined CNN feature extractors with Transformers. Dosovitskiy et al. asked: can a standard Transformer with no convolutions work for vision? The patch embedding is the only vision-specific component. Weight initialization throughout is \mathcal{N}(0, 0.02^2).

Numerical Example

Trace a tiny image through all stages. Image: C = 3, H = 4, W = 4, P = 2, D = 4, L = 1 block, h = 2 heads (d_h = 2), K = 3 classes.

Stage 1: Patch embedding. N = 16 / 4 = 4 patches. Each flattened to length 4 \cdot 3 = 12, projected through W_E \in \mathbb{R}^{12 \times 4}:

• z_1^{(0)} = [0.42, -0.18, 0.31, 0.05] (top-left)
• z_2^{(0)} = [-0.11, 0.53, 0.27, -0.34] (top-right)
• z_3^{(0)} = [0.29, 0.14, -0.22, 0.61] (bottom-left)
• z_4^{(0)} = [0.55, -0.07, 0.48, 0.19] (bottom-right)

Stage 2: CLS + position. Prepend z_{\text{cls}} = [0.10, 0.10, 0.10, 0.10]. Sequence: 5 tokens. Add $ \in \mathbb{R}^{5 \times 4}$:

• CLS: [0.10 + 0.01, 0.10 - 0.02, 0.10 + 0.03, 0.10 - 0.01] = [0.11, 0.08, 0.13, 0.09]
• Patch 1: [0.44, -0.17, 0.30, 0.08], patches 2-4 similar.

Stage 3: Encoder block. Focus on CLS token.

LayerNorm CLS [0.11, 0.08, 0.13, 0.09]: mean = 0.1025, std = 0.0192. Normalized: [0.39, -1.17, 1.43, -0.65].

MSA: CLS query attends to all 5 positions (no mask). After 2-head attention and W_O projection, output = [0.06, -0.03, 0.08, -0.02]. Residual: [0.17, 0.05, 0.21, 0.07].

MLP: LayerNorm, expand to 8 dims with GELU, compress back to 4. Output = [0.04, -0.01, 0.05, 0.02]. Residual: [0.21, 0.04, 0.26, 0.09].

Stage 4: Classification head. Extract CLS: z_0^{(1)} = [0.21, 0.04, 0.26, 0.09]. LayerNorm: h_f = [0.64, -1.17, 1.17, -0.64]. Project through W_{\text{head}} \in \mathbb{R}^{4 \times 3}:

• Class 0: 0.64(0.3) + (-1.17)(-0.1) + 1.17(0.2) + (-0.64)(0.4) = 0.192 + 0.117 + 0.234 - 0.256 = 0.287
• Class 1: 0.64(-0.2) + (-1.17)(0.5) + 1.17(0.1) + (-0.64)(-0.3) = -0.128 - 0.585 + 0.117 + 0.192 = -0.404
• Class 2: 0.64(0.1) + (-1.17)(0.3) + 1.17(-0.4) + (-0.64)(0.2) = 0.064 - 0.351 - 0.468 - 0.128 = -0.883

Logits: [0.287, -0.404, -0.883]. Predicted class: 0. No softmax. The CLS token started with no image content but aggregated all patch information through bidirectional attention.

ViT's Impact

ViT proved that Transformers work for vision without any convolutional layers. This sparked a transformation of the field:

• DeiT (2021): Knowledge distillation from a CNN teacher made ViT data-efficient, achieving competitive results with only ImageNet-1K.
• Swin Transformer (2021): Shifted windows for local attention created hierarchical features, enabling dense prediction tasks like detection and segmentation.
• MAE (2022): Masked 75% of patches and trained reconstruction, enabling powerful self-supervised pre-training for ViT.
• CLIP (2021): Paired a ViT image encoder with a text encoder for zero-shot classification from natural language supervision.

Pitfalls

1. Adding a causal mask to self-attention.

ViT uses bidirectional attention with no causal mask. Adding one would prevent patches from attending to later positions in the flattened sequence. Since "later" means "further right or down" in the image, a causal mask creates arbitrary asymmetry where top-left patches see less context than bottom-right patches.

2. Forgetting position embeddings.

Without position embeddings, the Transformer treats patches as an unordered set. Permuting patches would produce identical output. The model cannot learn spatial relationships without positional information.

3. Wrong patch size or miscounting patches.

P must evenly divide H and W. The position embedding matrix has N + 1 entries (the +1 is for CLS). Using $ = 14$ on a 224x224 image gives 256 patches instead of 196, requiring a different position embedding size.

4. Applying softmax to the output.

The forward pass outputs raw logits. Cross-entropy loss applies log-softmax internally. Adding softmax in the model creates double-softmax, compressing gradients toward zero and stopping learning.

5. Not extracting the CLS token.

The classification head operates only on position 0 (CLS), not the full sequence. Passing all N + 1 tokens through the linear head produces a (N+1) \times K output instead of a K-dimensional logit vector. Average pooling is a valid alternative but architecturally different from what ViT specifies.

Examples

Hints

1. Crop to complete patches before reshaping the image.
2. Use len(encoder_weights) as the number of encoder blocks.
3. Infer embedding and class widths from the supplied matrices.

Requirements

• Use NumPy.
• Extract complete patches in row-major order and ignore incomplete boundary patches.
• Apply W_patch and patch_bias to every flattened patch.
• Prepend cls_token and add pos_embed.
• Apply every encoder_weights dictionary in order.
• Return float64 class logits with shape (B, C).

Constraints

• image has shape (B, H, W, C_in) and dtype float64.
• patch_size is positive and no larger than H or W.
• W_patch has shape (patch_size × patch_size × C_in, D), and patch_bias has shape (D,).
• cls_token has shape (1, 1, D), and pos_embed has shape (1, N + 1, D).
• Every encoder matrix has compatible float64 dimensions, and D is divisible by num_heads.
• W_head has shape (D, C), and every supplied numeric array has dtype float64.

Starter Code

import numpy as np

def vit_forward(image: np.ndarray, patch_size: int, num_heads: int,
                W_patch: np.ndarray, patch_bias: np.ndarray,
                cls_token: np.ndarray, pos_embed: np.ndarray,
                encoder_weights: list, W_head: np.ndarray) -> np.ndarray:
    """
    Returns float64 Vision Transformer logits with shape (B, C).
    """
    pass

Test Cases