Vision Transformers

Vision transformers treat image patches as tokens and process them with transformer blocks. They replace the hard-coded locality of CNN architectures with learned attention over patch sequences, using the same core mechanism as attention.

CNN locality versus learned attention

A vision transformer keeps the transformer machinery but swaps the CNN’s built-in locality for attention that any patch can direct anywhere:

	CNN	Vision transformer
Inductive bias	local, translation-equivariant	learned global attention
Information mixing	neighboring convolution windows	any patch to any patch
Data efficiency	works with less data	usually needs large-scale pretraining
Cost driver	scales with resolution	attention is $O(N^2)$ in patch count

Patches to tokens

An image $x\in {\mathbb{R}}^{C\times H\times W}$ is split into $N=HW/P^2$ patches, flattened, projected, and combined with positional embeddings:

$$Z_0=[x_{1}E;\dots ;x_{N}E]+E_{\mathrm{pos}}.$$

Here $P$ is the patch width and height, $x_j$ is the flattened vector for patch $j$, $E$ is the learned projection into token space, and $E_{\mathrm{pos}}$ adds position information so the transformer can distinguish where each patch came from. The initial token sequence $Z_0$ is therefore an ordered grid of visual patch embeddings.

Each attention head computes

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d}}\right)V.$$

The query, key, and value matrices $Q$, $K$, and $V$ are linear projections of the patch tokens, and $d$ is the key dimension used to scale dot products. Attention turns patch-to-patch similarity into weights, then mixes value vectors from all patches.

Patch size controls spatial granularity: a large patch lowers compute but makes small details harder to represent.

Worked example

This snippet splits an image tensor into patches, projects them into tokens, and computes one row of token attention weights.

import torch
 
torch.manual_seed(8)
img = torch.arange(1*1*8*8, dtype=torch.float32).reshape(1,1,8,8)
patches = img.unfold(2,4,4).unfold(3,4,4).contiguous().view(1,1,4,4,4)
patches = patches.permute(0,2,1,3,4).reshape(1,4,16)
W = torch.randn(16,6) * 0.01
tokens = patches @ W
attn = torch.softmax((tokens @ tokens.transpose(-1,-2)) / (6**0.5), dim=-1)
print("patches_shape", tuple(patches.shape), "tokens_shape", tuple(tokens.shape))
print("first_attention_row", torch.round(attn[0,0], decimals=3).tolist())

Observed output:

patches_shape (1, 4, 16) tokens_shape (1, 4, 6)
first_attention_row [0.03200000151991844, 0.04399999976158142, 0.38999998569488525, 0.5339999794960022]

The 8-by-8 image becomes four patch tokens. Attention then mixes information globally across all patches rather than only through neighboring convolution windows.

Caveats

Attention cost grows as $O(N^2)$ in the number of patches, so high-resolution semantic segmentation needs architectural compromises. ViTs also tend to rely on pretraining and strong data augmentation; weak data regimes can favor CNN inductive bias.

References

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• PyTorch documentation: scaled dot product attention

Nav

Section — Computer Vision

← CNN Architectures Self Supervised Visual Learning →

Learning path — Computer vision

← Detection and Segmentation Metrics