Vision Transformer

The Vision Transformer (ViT) classifies an image by cutting it into a grid of fixed-size patches, treating each patch as a token, and passing the sequence through a standard transformer encoder. Dosovitskiy and eleven Google Research colleagues introduced it in an October 2020 preprint, “An Image is Worth 16x16 Words”, published at ICLR 2021. The model contains no convolutions beyond its patch embedding, which is equivalent to a single strided one. It trails convolutional networks of similar size when trained on ImageNet alone, and overtakes them when pre-trained on tens to hundreds of millions of images. It has since become a general-purpose backbone for self-supervised pre-training, vision–language models such as CLIP, and detectors and segmenters such as ViTDet and SAM.

Problem Addressed

By 2020 the transformer dominated natural language processing, but ResNet-style CNNs remained the strongest large-scale image classification models. Earlier models combined self-attention with convolutions or used local or axial attention patterns that, the ViT authors note, needed complex engineering to run efficiently on accelerators. Global attention over pixels is too costly: it is quadratic in the number of tokens, and a $224 \times 224$ image has 50,176 pixels. ViT applied an unmodified transformer to a short sequence of patches and made up for the missing image structure with data.

Architecture

The model follows the original transformer encoder as closely as possible (Dosovitskiy et al., 2021, §3.1).

Patch Embedding

An image $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ is reshaped into $N = HW/P^2$ non-overlapping patches of $P \times P$ pixels, each flattened to a vector $\mathbf{x}_p^i \in \mathbb{R}^{P^2 C}$. A learned matrix $\mathbf{E} \in \mathbb{R}^{(P^2 C) \times D}$ maps every patch to the model width $D$. Because the patches do not overlap, this linear map is a convolution with $D$ filters of size $P \times P \times C$ and stride $P$, which is how the official JAX code and the timm library implement it. A $224 \times 224$ image cut into $16 \times 16$ patches, the “words” of the title, gives a $14 \times 14$ grid of 196 tokens.

Class Token and Position Embeddings

As in BERT, a learned vector $\mathbf{x}_\text{class}$ is prepended to the sequence, and its output after the last block is the image representation. Learned position embeddings $\mathbf{E}_\text{pos} \in \mathbb{R}^{(N+1) \times D}$, one vector per patch in raster order plus one for the class token, are added to the tokens:

$$\mathbf{z}_0 = [\mathbf{x}_\text{class};\, \mathbf{x}_p^1 \mathbf{E};\, \mathbf{x}_p^2 \mathbf{E};\, \dots;\, \mathbf{x}_p^N \mathbf{E}] + \mathbf{E}_\text{pos}.$$

Without them the encoder would see an unordered set of patches; 2D and relative variants gave nearly the same accuracy (see positional encoding). The class token is a convention: average-pooling the patch outputs worked about as well once its learning rate was tuned (arXiv v2, Appendix D.3).

Transformer Encoder

The encoder stacks $L$ identical pre-norm blocks, each applying layer normalization before multi-head self-attention (MSA) and before a two-layer MLP with a GELU activation, with a residual connection around each:

$$\mathbf{z}'_\ell = \operatorname{MSA}(\operatorname{LN}(\mathbf{z}_{\ell-1})) + \mathbf{z}_{\ell-1}, \qquad \mathbf{z}_\ell = \operatorname{MLP}(\operatorname{LN}(\mathbf{z}'_\ell)) + \mathbf{z}'_\ell, \qquad \ell = 1, \dots, L.$$

Width $D$ and the token grid stay constant throughout; the transformer article covers the block and its parameter count.

Classification Head

The representation is $\mathbf{y} = \operatorname{LN}(\mathbf{z}_L^0)$, the normalized output of the class token. During pre-training the head is an MLP with one tanh hidden layer. For fine-tuning, the whole head is replaced by a single zero-initialized linear layer with one output per target class.

Model Sizes

Base and Large follow BERT; Huge is new. ViT-L/16 denotes the Large model with $16 \times 16$ patches.

Model	Layers $L$	Width $D$	MLP size	Heads	Parameters
ViT-Base	12	768	3,072	12	86M
ViT-Large	24	1,024	4,096	16	307M
ViT-Huge	32	1,280	5,120	16	632M
ViT-22B (Dehghani et al., 2023)	48	6,144	24,576	48	21,743M

The first three rows are Table 1 of Dosovitskiy et al. (2021). Halving the patch size quadruples the tokens without adding parameters; in the authors’ ablations this cost compute and raised accuracy.

Hybrid Architecture

The tokens can also come from a CNN feature map, one token per position ($1 \times 1$ patches). The paper’s hybrids feed ResNet-50 features, taken after stage 4 or an extended stage 3, into a ViT.

Training

ViT is pre-trained by supervised classification on a large dataset and then fine-tuned on the target task, the transfer learning recipe standard in natural language processing; for single-label data the loss is the cross-entropy.

Data. Pre-training used ImageNet (1.3 million images, 1,000 classes), its superset ImageNet-21k (14 million, 21,000 classes) or Google’s in-house JFT-300M (303 million, 18,000 classes).

Pre-training. All models, ResNet baselines included, used Adam ($\beta_1 = 0.9$, $\beta_2 = 0.999$), batch size 4,096, a 10,000-step linear warmup, weight decay 0.1 on JFT-300M and $224 \times 224$ inputs. Training from scratch on ImageNet needed strong regularization: weight decay 0.3, dropout 0.1 and gradient clipping (Appendix B.1).

Fine-tuning at higher resolution. Fine-tuning uses SGD with momentum 0.9, batch size 512 and, by default, $384 \times 384$ inputs: a $24 \times 24$ grid of 576 tokens at the same patch size. The encoder accepts any length, but the position embeddings must be resampled by 2D interpolation (§3.2; see positional encoding).

Inductive bias. A CNN builds locality, 2D neighborhood structure and translation equivariance into every layer; in ViT only the MLPs are local, and self-attention is global. Patch extraction and position-embedding interpolation are, in the authors’ words, “the only points” where 2D structure is injected by hand; all spatial relations between patches must be learned.

Making ViT data-efficient. Augmentation and regularization can replace much of the data, as DeiT showed on ImageNet-1k. In a systematic study of data, augmentation and regularization (AugReg), model size and compute, Steiner et al. (2022) found that more compute plus AugReg can match an order of magnitude more data: their ImageNet-21k models matched or beat counterparts trained on JFT-300M.

Inference

The output is a class distribution or, for a backbone, the grid of $N$ patch features. For $n = N + 1$ tokens of width $D$, a block spends about $12nD^2$ multiply-adds in its linear layers and $2n^2D$ in the attention products (see transformer), so attention cost grows with the square of $HW/P^2$. In the authors’ TPU timings the quadratic growth only begins to show for the largest models at the largest resolutions (arXiv v2, Appendix D.5).

Patch size as a cost dial. The NumPy code below checks on a synthetic $64 \times 64$ RGB image that the linear patch embedding equals a stride-16 convolution, then counts tokens and multiply-adds per block at $224 \times 224$ for three patch sizes with $D = 768$.

import numpy as np
from scipy.signal import correlate2d               # deep-learning "convolution" is cross-correlation

rng = np.random.default_rng(0)

# A 64x64 RGB synthetic image: a red square and a blue disc on a gray gradient.
H = W = 64
yy, xx = np.mgrid[0:H, 0:W]
img = np.repeat((0.3 + 0.4 * xx / (W - 1))[..., None], 3, axis=2)
img[8:28, 36:56] = [0.9, 0.1, 0.1]                                   # red square
img[(yy - 44) ** 2 + (xx - 20) ** 2 < 120] = [0.1, 0.2, 0.9]         # blue disc

# Patchify: (H, W, C) -> (N, P*P*C), patches in raster order, each flattened row by row.
P, C, D = 16, 3, 32
def patchify(x, P):
    h, w, c = x.shape
    return x.reshape(h // P, P, w // P, P, c).transpose(0, 2, 1, 3, 4).reshape(-1, P * P * c)

xp = patchify(img, P)
print("patches:", xp.shape, f"(N = HW/P^2 = {H * W // P**2}, P^2 C = {P * P * C})")

# Linear patch embedding E (P^2 C x D) plus bias, as in Eq. (1) of the ViT paper.
E = rng.normal(0, 0.02, (P * P * C, D))
b = rng.normal(0, 0.02, D)
tokens_linear = xp @ E + b                                            # (16, D)

# The same map as a convolution: D filters of size P x P x C applied with stride P (no padding).
K = E.reshape(P, P, C, D)                           # filter d, tap (dy, dx), channel c
out = np.stack([sum(correlate2d(img[..., c], K[..., c, d], mode="valid") for c in range(C))[::P, ::P]
                for d in range(D)], axis=-1) + b    # dense valid output, keep every P-th row and column
tokens_conv = out.reshape(-1, D)
print(f"max |linear embedding - stride-{P} convolution| = {np.abs(tokens_linear - tokens_conv).max():.1e}")

# Sequence length and attention cost at 224 x 224 for three patch sizes (ViT-B width D = 768).
# Per layer: 4 N D^2 multiply-adds for the Q, K, V and output projections, 8 N D^2 for the MLP,
# and 2 N^2 D for the scores QK^T and the weighted sum of values.
print("\n224 x 224 input, D = 768")
print(" P   patches  tokens  score entries/head  attention MACs  linear-layer MACs  attention share")
for P in (8, 16, 32):
    n = (224 // P) ** 2 + 1                                           # + class token
    d = 768
    attn, lin = 2 * n * n * d, 12 * n * d * d
    print(f"{P:2d}  {n - 1:7d}  {n:6d}  {n * n:18,d}  {attn / 1e9:12.3f} G  {lin / 1e9:15.3f} G  {attn / (attn + lin):14.1%}")

patches: (16, 768) (N = HW/P^2 = 16, P^2 C = 768)
max |linear embedding - stride-16 convolution| = 6.7e-16

224 x 224 input, D = 768
 P   patches  tokens  score entries/head  attention MACs  linear-layer MACs  attention share
 8      784     785             616,225         0.947 G            5.556 G           14.6%
16      196     197              38,809         0.060 G            1.394 G            4.1%
32       49      50               2,500         0.004 G            0.354 G            1.1%

The two embeddings agree to rounding error. From $32 \times 32$ to $8 \times 8$ patches the tokens grow 16-fold and the attention products about 250-fold, yet at $224 \times 224$ the linear layers still dominate. In our arithmetic the attention term catches up only at detection and segmentation resolutions (4,097 tokens for a $1024 \times 1024$ input), where ViTDet and SAM restrict most blocks to windows.

Performance

All figures are top-1 accuracy as reported by Dosovitskiy et al. (2021).

Against the state of the art. Fine-tuned on ImageNet at $518 \times 518$ with Polyak averaging, ViT-H/14 pre-trained on JFT-300M reached 88.55%, against 87.54% for BiT-L (a ResNet-152x4 pre-trained on the same data) and 88.4/88.5% for Noisy Student (EfficientNet-L2); it also reached 77.63% on the 19-task VTAB suite (Table 2). Pre-training took 2.5k TPUv3-core-days, against 9.9k for BiT-L and 12.3k for Noisy Student. ViT-L/16 reached 87.76% with 0.68k core-days, and 85.30% pre-trained on the public ImageNet-21k (about 30 days on one 8-core TPUv3).

Data scale. Fine-tuned at $384 \times 384$ (Table 5), ViT-B/16 and ViT-L/16 reached 77.91% and 76.53% when pre-trained on ImageNet itself, 83.97% and 85.15% with ImageNet-21k, and 84.15% and 87.12% with JFT-300M. The larger model lost on ImageNet, gained about a point on ImageNet-21k and three on JFT-300M; on ImageNet alone both trailed the BiT ResNets (Figure 3). The transformer article covers the JFT-subset experiment behind the authors’ conclusion that convolutional inductive bias helps on smaller datasets.

Compute. On JFT-300M, ViT needed about 2–4× less pre-training compute than ResNets (with group instead of batch normalization) for the same transfer accuracy, averaged over five datasets. Hybrids were slightly better at small budgets, the difference vanishing for larger models, and ViT did not saturate within the range tried (§4.4).

What the model learns. The position embeddings learn the 2D layout: nearby patches, and patches in the same row or column, have similar embeddings (Figure 7). The authors’ attention distance, the attention-weighted image-space distance between query and attended pixels, is an analogue of receptive field size. In the lowest layers some heads already attend across most of the image, so global context is used from the start, while others stay local; hybrids show less local attention, which the authors read as those heads doing the job of early convolutions. Attention distance grows with depth.

A toy version of the data effect. The script below (PyTorch 2.14.1, about 30 seconds on one CPU thread, identical output on every run) trains a tiny ViT ($4 \times 4$ patches on $32 \times 32$ images, $D = 64$, two blocks) and a four-layer CNN of nearly the same size to classify synthetic squares, discs, triangles and crosses of random size, brightness and position on a noisy background. Both get 600 AdamW steps of 32 images, drawn from 200 or 4,000 training images, and share a fixed 1,000-image test set.

import math
import numpy as np
import torch
from torch import nn

torch.set_num_threads(1)                     # single-threaded CPU run: identical output on every run
print(f"torch {torch.__version__}")

def make_shapes(n, rng, size=32):
    """n grayscale 32x32 images of four classes (square, disc, triangle, cross) with random
    size, brightness and position anywhere in the image, on a noisy shaded background."""
    yy, xx = np.mgrid[0:size, 0:size].astype(np.float32)
    X = np.empty((n, 1, size, size), np.float32)
    y = rng.integers(0, 4, n)
    for i in range(n):
        r = rng.uniform(4, 7)                                      # half-width of the shape
        cy, cx = rng.uniform(r + 1, size - r - 1, 2)               # position jitter
        dy, dx = yy - cy, xx - cx
        if y[i] == 0:   m = (abs(dy) < r) & (abs(dx) < r)                                     # square
        elif y[i] == 1: m = dy ** 2 + dx ** 2 < r ** 2                                         # disc
        elif y[i] == 2: m = (abs(dy) < r) & (abs(dx) < (dy + r) / 2)                           # triangle
        else:           m = ((abs(dy) < r) & (abs(dx) < r / 3)) | ((abs(dx) < r) & (abs(dy) < r / 3))  # cross
        g = rng.uniform(-0.2, 0.2, 2)                              # random shading of the background
        bg = 0.4 + g[0] * (xx / size - 0.5) + g[1] * (yy / size - 0.5)
        X[i, 0] = bg + m * rng.uniform(0.3, 0.6) + rng.normal(0, 0.08, (size, size))
    return torch.from_numpy(X), torch.from_numpy(y)

class Block(nn.Module):
    """Pre-norm transformer block: z + MSA(LN(z)), then z + MLP(LN(z)) with a GELU."""
    def __init__(self, d, heads, mlp):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(d, mlp), nn.GELU(), nn.Linear(mlp, d))
    def forward(self, z):
        h = self.ln1(z)
        z = z + self.attn(h, h, h, need_weights=False)[0]
        return z + self.mlp(self.ln2(z))

class TinyViT(nn.Module):
    """4x4 patches of a 32x32 image -> 64 tokens + [class] token, learned 1D position embeddings."""
    def __init__(self, img=32, patch=4, d=64, depth=2, heads=4, mlp=128, classes=4):
        super().__init__()
        n = (img // patch) ** 2
        self.embed = nn.Conv2d(1, d, patch, stride=patch)          # patch embedding = stride-P convolution
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.pos = nn.Parameter(0.02 * torch.randn(1, n + 1, d))
        self.blocks = nn.Sequential(*[Block(d, heads, mlp) for _ in range(depth)])
        self.norm, self.head = nn.LayerNorm(d), nn.Linear(d, classes)
    def forward(self, x):
        t = self.embed(x).flatten(2).transpose(1, 2)               # (B, 64, d), patches in raster order
        t = torch.cat([self.cls.expand(len(t), -1, -1), t], 1) + self.pos
        return self.head(self.norm(self.blocks(t))[:, 0])          # classify from the [class] token

def tiny_cnn(c=28, classes=4):
    return nn.Sequential(
        nn.Conv2d(1, c, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),              # 16x16
        nn.Conv2d(c, 2 * c, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),          # 8x8
        nn.Conv2d(2 * c, 2 * c, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),      # 4x4
        nn.Conv2d(2 * c, 2 * c, 3, padding=1), nn.ReLU(),
        nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(2 * c, classes))

def n_params(m):
    return sum(p.numel() for p in m.parameters())

def train_eval(model, X, y, Xte, yte, steps=600, bs=32, lr=2e-3, warmup=30):
    """Same budget for every model and training-set size: 600 AdamW steps of 32 images."""
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.05)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min((s + 1) / warmup,
                0.5 * (1 + math.cos(math.pi * (s - warmup) / (steps - warmup)))))
    for _ in range(steps):
        idx = torch.randint(0, len(y), (bs,))
        loss = nn.functional.cross_entropy(model(X[idx]), y[idx])
        opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    with torch.no_grad():
        return (model(Xte).argmax(1) == yte).float().mean().item()

SEED = 4
Xte, yte = make_shapes(1000, np.random.default_rng(1000))       # one fixed test set of 1,000 images
print(f"TinyViT {n_params(TinyViT()):,} parameters, CNN {n_params(tiny_cnn()):,} parameters")
for n_train in (200, 4000):
    X, y = make_shapes(n_train, np.random.default_rng(SEED))
    accs = []
    for make in (TinyViT, tiny_cnn):
        torch.manual_seed(SEED)
        accs.append(train_eval(make(), X, y, Xte, yte))
    print(f"{n_train:5d} training images, seed {SEED}: test accuracy ViT {accs[0]:.3f}, CNN {accs[1]:.3f}")

torch 2.14.1
TinyViT 72,644 parameters, CNN 71,236 parameters
  200 training images, seed 4: test accuracy ViT 0.491, CNN 0.872
 4000 training images, seed 4: test accuracy ViT 0.793, CNN 0.903

Over seeds 0–9 the ViT scored 0.409–0.522 with 200 images and 0.765–0.832 with 4,000; the CNN scored 0.798–0.963 and 0.740–0.984. Seed 4 is near the mean in every column. With 200 images the ViT fits its training set perfectly (training accuracy 1.000 for seed 4) but generalizes poorly, while the CNN, whose weight sharing and pooling tolerate shifts (our reading), does well. With 4,000 images the ViT gains about 0.34 on average and the mean gap shrinks from 0.44 to 0.10. The CNN’s mean barely moves because two of its runs ended near 0.75 on training and test images alike, unconverged after 600 steps. Without built-in locality, the transformer needs more examples than a CNN on a task where locality matters; the toy says nothing about the millions of images at which ViT overtakes CNNs.

Strengths

• Simplicity and scalability. An unmodified transformer runs on standard accelerator kernels, and scaling depth, width or token count kept improving accuracy in the authors’ experiments.
• Efficiency at scale. Less pre-training compute than ResNets for the same accuracy on JFT-300M, and, for large ViTs, larger per-core batches than ResNets (arXiv v2, Appendix D.5).
• A uniform token interface. Patch tokens can be dropped, masked or mixed with text tokens, as in masked image modeling and vision–language models.

Limitations

• Data hunger. Without large pre-training sets, or heavy augmentation, regularization and distillation, ViTs trail CNNs of similar size.
• Quadratic attention cost in $N = HW/P^2$, which limits resolution and patch size.
• A single coarse scale. Liu et al. (2021) cite ViT’s single low-resolution feature map, with its quadratic cost, as unsuitable for a general-purpose dense-prediction backbone.

Historical Importance

ViT was not the first transformer on image patches: its authors name Cordonnier et al. (2020), with full self-attention on $2 \times 2$ patches, as the closest prior work. ViT showed that with large-scale pre-training a plain transformer could match or beat the best CNNs at lower pre-training cost, on medium-resolution images. In our reading, it turned the field’s question from how to add attention to CNNs into how to give transformers the useful properties of CNNs (hierarchy, locality, data efficiency), which DeiT, Swin and their successors pursued. A small masked-patch-prediction experiment in the paper lifted ViT-B/16 to 79.9% on ImageNet, 2 points above training from scratch but 4 below supervised pre-training, an early step toward masked image modeling.

Variants

• DeiT (Touvron et al., 2021) trained the unchanged ViT-B on ImageNet-1k alone, on one 8-GPU node in two to three days, reaching 83.1% at $384 \times 384$. A distillation token trained to reproduce a CNN teacher’s predictions raised this to 85.2% after 1,000 epochs (recipe in transformer).
• Swin Transformer (Liu et al., 2021) restricts attention to $7 \times 7$-token windows that shift between consecutive blocks, so cost grows linearly with image size, and merges neighboring patches between stages into a CNN-like feature pyramid. Swin-L pre-trained on ImageNet-22K reached 87.3% top-1 on ImageNet-1K, and 58.7 box AP on COCO test-dev and 53.5 mIoU on ADE20K val as the backbone of HTC++ and UperNet.
• Self-supervised ViTs (see self-supervised learning). Masked autoencoders (MAE; He et al., 2022) encode only the 25% of patches left visible and reconstruct the rest; an MAE ViT-H fine-tuned at $448 \times 448$ reached 87.8% using only ImageNet-1k data. DINO trains ViTs by self-distillation without labels, and its ViT-B/8 reached 80.1% in linear evaluation (Caron et al., 2021); DINOv2’s frozen ViT-g/14 features reached 86.5% (Oquab et al., 2024).
• CLIP (Radford et al., 2021) trained ResNet and ViT image encoders (ViT-B/32, B/16 and L/14, with one extra layer normalization) contrastively against a text encoder on 400 million image–text pairs.
• Plain ViT backbones for detection and segmentation. ViTDet (Li et al., 2022) builds a simple feature pyramid from the last, stride-16 feature map and uses window attention plus four global-attention blocks. With MAE pre-training on ImageNet-1K, its ViT-H reached 61.3 box AP on COCO val2017 with a Cascade Mask R-CNN and multi-scale testing. SAM uses the same recipe for its image encoder: an MAE-pre-trained ViT-H/16 with $14 \times 14$ windowed attention and four global-attention blocks on $1024 \times 1024$ inputs (Kirillov et al., 2023). DETR took the other route, a transformer on top of a CNN backbone.
• Scaling. ViT-22B (Dehghani et al., 2023) uses $14 \times 14$ patches and was trained on a JFT version of about 4 billion images. For efficiency and stability it runs attention and MLP in parallel, normalizes queries and keys, and drops QKV-projection and layer-norm biases. As a frozen feature extractor with a linear probe at $224 \times 224$, it reached 89.51% on ImageNet.

Implementations

The authors’ JAX/Flax code and checkpoints are in the google-research/vision_transformer repository. Steiner et al. (2022) released more than 50,000 AugReg checkpoints through it and through Ross Wightman’s PyTorch library timm (pytorch-image-models), which implements ViT and many variants; their training code is in google-research/big_vision. When loading a checkpoint at a new resolution, check how the library resamples the position embeddings (see positional encoding).