Vision Transformer (ViT)

The Vision Transformer (ViT) is a deep learning architecture that applies the Transformer model — originally designed for natural language processing — directly to sequences of fixed-size image patches, achieving state-of-the-art performance on image classification benchmarks. It matters because it demonstrated that convolutional inductive biases are not necessary for top-tier vision performance, fundamentally reshaping how practitioners build image models. Remarkably, ViT underperforms CNNs on small datasets but dramatically surpasses them when pre-trained on very large datasets (100M+ images), inverting the intuition that more architectural priors always help.

The core insight

The key insight of ViT is that an image can be treated exactly like a sequence of words: slice it into a grid of non-overlapping patches, linearly embed each patch into a token vector, prepend a learnable classification token, and feed the entire sequence through a standard Transformer encoder. Because self-attention operates on all pairs of tokens simultaneously, every patch can directly attend to every other patch from the very first layer — something convolutions can only approximate after many stacked layers. The "aha moment" is realizing that the Transformer's global receptive field, previously exploited for long-range dependencies in text, is equally powerful for long-range spatial relationships in images. You do not need to redesign the Transformer; you only need to redesign how you present the image to it.

How It Works

Patch Embedding

Given an input image $\mathbf{X} \in \mathbb{R}^{H \times W \times C}$ , ViT divides it into a grid of $N$ non-overlapping patches of size $P \times P$ . The total number of patches is:

\quad N = \frac{H \cdot W}{P^2}

Explanation: N is the sequence length fed to the Transformer — the image resolution and patch size together determine how many "tokens" the model sees. Each patch $\mathbf{x}_i \in \mathbb{R}^{P^2 \cdot C}$ is flattened and projected to a $D$ -dimensional embedding space via a learned linear projection matrix $\mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}$ :

\quad \mathbf{z}_0 = [\mathbf{x}_{\text{class}};\, \mathbf{x}_1\mathbf{E};\, \mathbf{x}_2\mathbf{E};\, \ldots;\, \mathbf{x}_N\mathbf{E}] + \mathbf{E}_{\text{pos}}

Explanation: The patch embeddings are concatenated with a learnable [CLS] token and then summed with learnable positional embeddings before being passed to the Transformer. Here $\mathbf{x}_{\text{class}} \in \mathbb{R}^D$ is the learnable classification token, and $\mathbf{E}_{\text{pos}} \in \mathbb{R}^{(N+1) \times D}$ are the positional embeddings — both are free parameters optimized during training.

Transformer Encoder

Each of the $L$ Transformer encoder blocks applies Multi-Head Self-Attention (MSA) followed by a feed-forward MLP, with LayerNorm and residual connections:

\quad \mathbf{z}'_\ell = \text{MSA}(\text{LN}(\mathbf{z}_{\ell-1})) + \mathbf{z}_{\ell-1}, \qquad \mathbf{z}_\ell = \text{MLP}(\text{LN}(\mathbf{z}'_\ell)) + \mathbf{z}'_\ell

Explanation: The residual connections around both MSA and MLP are what make it feasible to stack dozens of Transformer layers without gradient vanishing. In the Multi-Head Self-Attention block, for each head $h$ with projection matrices $\mathbf{W}^Q_h, \mathbf{W}^K_h, \mathbf{W}^V_h \in \mathbb{R}^{D \times d_h}$ where $d_h = D/H$ :

\quad \text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_h}}\right)\mathbf{V}

Explanation: Dividing by $\sqrt{d_h}$ prevents the dot products from growing so large that softmax saturates into near-zero gradients, a critical stabilization trick.

Classification Head

After $L$ encoder layers, only the [CLS] token representation $\mathbf{z}_L^0 \in \mathbb{R}^D$ is used for classification. At pre-training time a two-layer MLP head is used; at fine-tuning time it is replaced by a single linear layer:

\quad \hat{y} = \text{softmax}(\mathbf{W}_{\text{cls}}\,\mathbf{z}_L^0 + \mathbf{b}_{\text{cls}})

Explanation: All global image information must flow into this single [CLS] token through attention during the forward pass — the Transformer's job is to aggregate patch information into it.

PyTorch Implementation Sketch

```python import torch import torch.nn as nn from einops import rearrange # pip install einops class PatchEmbedding(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768): super().__init__() self.n_patches = (img_size // patch_size) ** 2

# A single Conv2d with kernel=patch_size, stride=patch_size

# is algebraically identical to the linear patch projection self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x):

# x: (B, C, H, W) x = self.proj(x) # (B, embed_dim, H/P, W/P) x = rearrange(x, 'b d h w -> b (h w) d') # (B, N, embed_dim) return x

class ViT(nn.Module): def __init__(self, img_size=224, patch_size=16, embed_dim=768, depth=12, n_heads=12, mlp_ratio=4, n_classes=1000): super().__init__() n_patches = (img_size // patch_size) ** 2 self.patch_embed = PatchEmbedding(img_size, patch_size, 3, embed_dim) self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed = nn.Parameter(torch.zeros(1, n_patches + 1, embed_dim)) encoder_layer = nn.TransformerEncoderLayer( d_model=embed_dim, nhead=n_heads, dim_feedforward=embed_dim * mlp_ratio, batch_first=True, norm_first=True # Pre-LN as in ViT paper ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=depth) self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, n_classes) def forward(self, x): B = x.size(0) x = self.patch_embed(x) # (B, N, D) cls = self.cls_token.expand(B, -1, -1) # (B, 1, D) x = torch.cat([cls, x], dim=1) # (B, N+1, D) x = x + self.pos_embed # add positional emb x = self.norm(self.encoder(x)) # (B, N+1, D) return self.head(x[:, 0]) # CLS token -> logits ```

researcher · Alexey Dosovitskiy | b. 1986 | Original ViT Paper (2020)

Dosovitskiy is a research scientist at Google Brain (later Google DeepMind) known for bridging the gap between theoretical deep learning and practical computer vision systems. His landmark 2020 paper "An Image is Worth 16x16 Words" introduced ViT by demonstrating that a pure Transformer applied to image patches, with minimal image-specific modifications, could match or exceed CNN performance when pre-trained on sufficiently large datasets. Surprisingly, the paper was initially rejected from ICLR 2021 by reviewers who doubted the novelty — it was eventually accepted and has since accumulated tens of thousands of citations.

researcher · Ashish Vaswani | b. 1983 | Original Transformer Architecture

Vaswani was a senior research scientist at Google Brain and is the lead author of "Attention Is All You Need" (2017), the paper that introduced the Transformer architecture on which ViT is entirely built. His self-attention mechanism — originally designed to handle variable-length text sequences without recurrence — turned out to be general enough that researchers like Dosovitskiy could apply it verbatim to image patches a mere three years later. Vaswani later co-founded Essential AI, demonstrating that the architecture he helped create would reshape not just NLP but the entire field of deep learning.

Figure

Vision Transformer Architecture

Figure

Patch Tokenization Detail