Patch Embeddings and Positional Encoding in ViT
Patch embedding turns a 2D image into a token sequence with one strided convolution, and position embeddings restore the spatial layout the transformer would otherwise ignore.
The patch embedding operation
The first layer of a Vision Transformer converts raw pixels into tokens. Given an image of height H, width W, and C channels, and a patch size P, the image contains (H/P) x (W/P) patches. Each patch has P*P*C values, which a learned linear map projects to the model dimension D. Implementations usually realize this as a single convolution with kernel size P and stride P, which is mathematically identical to flattening each patch and multiplying by a weight matrix.
import torch, torch.nn as nn
# image: (B, C, H, W); patch P; model dim D
proj = nn.Conv2d(in_channels=3, out_channels=768, kernel_size=16, stride=16)
x = torch.randn(2, 3, 224, 224)
tokens = proj(x) # (2, 768, 14, 14)
tokens = tokens.flatten(2).transpose(1, 2) # (2, 196, 768)Number of tokens
For a 224x224 image with P=16, the grid is 14x14 = 196 patch tokens. Adding the class token gives a sequence of length 197. Because self-attention cost grows with the square of sequence length, halving the patch size roughly quadruples the number of tokens and increases attention cost by about sixteen times, which is why patch size is a central efficiency knob.
Position embeddings
Self-attention has no notion of order, so a learned position embedding is added to each token. ViT uses a table of 1D learned vectors, one per position, added elementwise. When fine-tuning at a higher resolution, the grid has more patches than the pretrained table, so the position embeddings are reshaped to the original grid and interpolated to the new grid before use.
- Learned 1D embeddings work about as well as 2D-aware ones in practice
- The class-token position also has its own embedding
- Interpolation lets one checkpoint serve multiple input resolutions
Design consequences
Patch embedding is where nearly all spatial inductive bias enters a ViT, and it is minimal: a linear projection plus additive positions. Everything else is learned by attention. This is why ViT is data-hungry but flexible, and why so many variants innovate primarily on the tokenizer and the position scheme rather than the encoder blocks themselves.