Vision Transformer (ViT)

A Vision Transformer (ViT) is a neural network that applies the transformer architecture from language modeling to images. It cuts an image into fixed-size patches, turns each patch into a token embedding, and processes the sequence with self-attention, so every patch can draw on every other patch from the first layer onward. Google researchers introduced the architecture in "An Image is Worth 16x16 Words", which showed that a pure transformer with no convolution layers matches state-of-the-art convolutional neural networks (CNNs) on image classification when pretrained on enough data. ViTs now serve as the image encoders of vision-language models, segmentation foundation models, and self-supervised backbones.

How Vision Transformers Work#

A ViT turns an image into a sequence and then processes it like text. The steps below use ViT-Base/16 on a 224×224 image:

1. Patch embedding: The image is cut into 16×16-pixel patches, a 14×14 grid of 196 patches. Each patch is flattened and multiplied by a learned matrix to become a 768-dimensional embedding, which is usually implemented as a 16×16 convolution with stride 16.
2. Class token and position embeddings: A learnable class token is prepended to the sequence, and learnable position embeddings are added to every token. Attention by itself ignores order, so these embeddings are the only source of information about where each patch came from.
3. Transformer encoder: A stack of identical blocks, 12 in ViT-Base, alternates multi-head self-attention, which mixes information across all tokens, with a two-layer MLP using GELU activations, which transforms each token on its own. Layer normalization precedes each sublayer, and a residual connection wraps it.
4. Head: The class token's final output feeds a linear classifier. Dense tasks such as segmentation use the patch tokens instead, reshaped back into a 14×14 grid.

The PyTorch code below builds the token sequence for one image and passes it through a single encoder block.

import torch
from torch import nn

# ViT-Base/16 settings: 16x16 patches, 768-dim tokens, 12 attention heads
image = torch.randn(1, 3, 224, 224)
patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16)  # one linear projection per patch

# Cut the image into a 14x14 grid of patches and flatten it into 196 tokens
tokens = patch_embed(image).flatten(2).transpose(1, 2)

# Prepend a learnable class token and add learnable position embeddings
cls_token = nn.Parameter(torch.zeros(1, 1, 768))
pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
tokens = torch.cat([cls_token, tokens], dim=1) + pos_embed

# One pre-norm encoder block in which every token attends to every other token
block = nn.TransformerEncoderLayer(
    d_model=768, nhead=12, dim_feedforward=3072, activation="gelu", batch_first=True, norm_first=True
)
print(block(tokens).shape)

It prints torch.Size([1, 197, 768]): the 196 patch tokens plus the class token, each still 768-dimensional, ready for the next of ViT-Base's 12 blocks.

Model Sizes and Variants#

ViT names give the model size and patch size: ViT-B/16 is the 86-million-parameter Base model with 16×16 patches, and the Large and Huge models have 307 million and 632 million parameters. Smaller patches produce more tokens, which improves detail at a higher compute cost. Several variants changed how ViTs are trained or structured:

• DeiT: A training recipe with strong augmentation and a distillation token that trains competitive ViTs on ImageNet-1k alone, without extra pretraining data.
• Swin Transformer: Computes attention inside shifted local windows and merges patches between stages, producing a CNN-like hierarchy that suits detection and segmentation.
• Masked autoencoders (MAE) and DINOv2: Self-supervised methods that pretrain ViTs on unlabeled images, by reconstructing masked patches in MAE and by matching a teacher network's outputs in DINOv2.
• CLIP and SigLIP: ViT image encoders trained to match images with their captions, widely reused as the vision tower of multimodal models.

Applications#

• Vision-language models: A ViT encoder converts an image into a sequence of tokens that a large language model reads alongside text, enabling captioning and visual question answering.
• Dense prediction: Patch tokens from ViT backbones drive image segmentation and depth estimation, where each output pixel needs context from far across the image.
• Video understanding: Video ViTs split clips into spatiotemporal patches, and self-supervised video models such as V-JEPA learn motion and appearance features from unlabeled footage for tasks like action recognition.

Vision Transformers vs. CNNs#

CNNs build two assumptions into every layer: locality, because each filter sees a small neighborhood, and translation equivariance, because the same filter is applied at every position. Apart from the patch grid, a ViT has to learn spatial structure from data. In the original paper, ViTs trained on ImageNet alone (1.3 million images) scored a few percentage points below ResNets of comparable size, and pretraining on the 300-million-image JFT-300M dataset lifted ViT-H/14 to 88.55% ImageNet top-1 accuracy. Cost is the other difference. Self-attention compares every token with every other token, so doubling the image's side length quadruples the token count and multiplies the attention cost by 16. Flash attention reduces the memory traffic of that computation, and windowed attention limits the comparisons themselves. Ultralytics runs Meta's Segment Anything Model, whose sam_b.pt checkpoint encodes each 1024×1024 image with a ViT-B backbone: the 16×16 patches form a 64×64 grid of 4,096 tokens, so 8 of its 12 blocks attend only within 14×14 windows and 4 use global attention.