Vision Transformers

A Vision Transformer (ViT) applies a standard Transformers encoder to images by cutting the image into fixed size patches and treating each patch as a token. It shows that convolutions are not necessary for strong image recognition given enough data.

From Image to Tokens

The core trick is patchify: split an image of shape $(C,H,W)$ into a grid of non overlapping $P\times P$ patches, flatten each patch, and linearly project it to the model dimension $D$. This produces $N=\frac{HW}{P^2}$ tokens.

$${\mathbf{z}}_0=[x_{\text{class}};x_p^{1}E;x_p^{2}E;\dots ;x_p^{N}E]+E_{\text{pos}}$$

• $E\in {\mathbb{R}}^{(P^{2}C)\times D}$ is the shared patch embedding projection.
• $x_{\text{class}}$ is a learnable class token prepended to the sequence; its final state is the image representation for classification.
• $E_{\text{pos}}$ are learnable positional embeddings added because the transformer is otherwise permutation invariant.

Patch Embedding in PyTorch

einops.rearrange expresses patchify in one line, exactly as the ViT paper describes it.

import torch
import torch.nn as nn
from einops import rearrange, repeat
 
class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch=16, in_ch=3, dim=768):
        super().__init__()
        self.p = patch
        n_patches = (img_size // patch) ** 2
        self.proj = nn.Linear(patch * patch * in_ch, dim)
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.zeros(1, n_patches + 1, dim))
 
    def forward(self, x):
        # x: (b, c, h, w) -> (b, num_patches, patch_dim)
        patches = rearrange(
            x, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)",
            p1=self.p, p2=self.p,
        )
        tokens = self.proj(patches)                     # (b, n, dim)
        cls = repeat(self.cls, "1 1 d -> b 1 d", b=x.shape[0])
        tokens = torch.cat([cls, tokens], dim=1)        # prepend class token
        return tokens + self.pos                        # add positions

The rearrange pattern

b c (h p1) (w p2) -> b (h w) (p1 p2 c) reads the height axis as h patch-rows each of size p1, the width as w patch-cols of size p2, then flattens the spatial grid into a token axis (h w) and each patch into a feature vector (p1 p2 c). A Conv2d with stride = patch is an equivalent, common alternative.

Transformer Encoder

Each encoder block is pre-norm (LayerNorm) then Self-Attention then MLP, with residuals:

$$\begin{array}{rl}{\mathbf{z}}_{\ell }^{\prime }& =\text{MSA}(\text{LN}({\mathbf{z}}_{\ell -1}))+{\mathbf{z}}_{\ell -1}\\ {\mathbf{z}}_{\ell }& =\text{MLP}(\text{LN}({\mathbf{z}}_{\ell }^{\prime }))+{\mathbf{z}}_{\ell }^{\prime }\end{array}$$

Attention is global from the first layer: every patch can attend to every other patch, unlike a CNN’s local kernel. The class token collects information from all patches; a final linear head on its output produces logits.

Practical Notes

Data hungry

ViTs lack the built in locality and translation equivariance of CNNs, so trained from scratch on small datasets they underperform. They need large scale pretraining (for example JFT-300M) or heavy augmentation and regularization (DeiT) to match or beat CNNs.

• Positional embeddings must be interpolated when changing input resolution, since the number of patches changes.
• Patch size trades resolution for compute: smaller patches mean more tokens and quadratic attention cost.
• Hybrid variants feed CNN feature maps as the patch source, and later models (Swin) reintroduce locality with windowed attention.