A Vision Transformer (ViT) applies the same attention machinery as a language model, but the inputs are images. Walk through how the per-token sequence is built from an image, what attention mask is used, and how the model knows where each patch came from.
Free · 2 AI evals / day
ViT applies the same transformer attention machinery as a language model, but the inputs are images. The trick is converting the image into a sequence of tokens.
Start with an image of shape (H, W, 3), say 224 x 224 x 3 RGB. Cut it into a grid of fixed-size patches, typically 16 x 16 pixels each. A 224 x 224 image gives a 14 x 14 grid, which is 196 patches.
Flatten each patch into a vector of length 16 x 16 x 3 = 768 numbers, then pass it through one linear projection to get a d_model-dimensional embedding. The sequence of 196 patch embeddings is what attention runs over.
For classification, ViT prepends a learnable [CLS] token to the sequence (so the input is 197 tokens total: 196 patches plus 1 CLS). After all transformer layers, the final hidden state of the CLS token gets fed to a classification head.
The attention mask is FULL bidirectional. Every patch attends to every other patch in both directions. There is no causal mask because there's no notion of 'future' in a static image, the patches all exist simultaneously and any patch can in principle inform any other.
Position is injected via learned 2D positional embeddings added to the patch embeddings. Without PE, ViT would be permutation-invariant over patches, which would destroy the spatial structure the model needs.
**The pipeline.** An image of shape (H, W, 3) is split into a grid of fixed-size patches (typically 16 x 16 pixels). For a 224 x 224 image that's a 14 x 14 grid, 196 patches. Each patch is flattened (16 * 16 * 3 = 768 values for RGB) and passed through ONE linear projection that produces a d_model-dimensional token embedding. The output is a sequence of patch tokens.
**The CLS token.** For classification ViT prepends a learnable [CLS] token, making the sequence 197 tokens long. After all the transformer layers run, the final hidden state of the CLS token is the pooled image representation that feeds the classification head. The CLS token is not strictly required (some variants use mean-pooling over patch tokens instead), but it's the convention from the original ViT paper.
**The attention mask: FULL bidirectional.** Every patch can attend to every other patch in both directions. There is no causal mask because an image has no temporal direction, patches exist simultaneously and any pair can be structurally related (a corner can inform the center, the center can inform a corner). Applying a causal mask would arbitrarily block half the attention based on the raster order of patches, which would destroy structural reasoning.
**Position encoding.** Self-attention is permutation-invariant over input tokens. Without position information, ViT couldn't distinguish patch (0, 0) from patch (10, 10), they'd be the same token. The original ViT adds learned 2D positional embeddings to the patch embeddings before the transformer. Swin and MaxViT use relative position biases inside attention instead. Either way, position has to be injected, otherwise the spatial structure is lost.
The rest of the architecture (transformer encoder, LayerNorm, FFN, output projection) is identical to a BERT-style encoder.
Quick reference
TL;DR
Split the image into fixed-size patches, linearly project each into a token, prepend a CLS token, run full bidirectional attention, and add learned 2D positional embeddings.
Imagine cutting a photograph into a grid of small square tiles, like a jigsaw puzzle. You hand the model the tiles in a stack along with a sticky note on each one saying 'I came from row 3, column 7'. The model looks at every tile and lets each one see every other tile, no rules about which tiles came first because all tiles exist at once. After all that looking, one special blank tile (the CLS token) collects everything the model figured out about the whole picture and hands the summary off for classification.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.
Very small images or low-resolution inputs (e.g., MNIST at 28 x 28) Patch-based tokenization produces too few patches to benefit from the attention machinery; classical CNN architectures often outperform ViT on small inputs.
Tasks requiring dense per-pixel output (semantic segmentation, dense prediction) without dedicated decoder architecture ViT outputs are per-patch; getting back to per-pixel resolution requires a decoder (UPerNet, Mask2Former) or a hybrid hierarchical backbone like Swin.
Very high-resolution inputs without windowed or hierarchical attention Full bidirectional attention is O(T^2) in sequence length; a 1024 x 1024 image with 16 x 16 patches has T = 4096, which makes vanilla ViT prohibitively expensive. Use Swin or hierarchical ViT variants instead.
Settings with very small datasets Vanilla ViT lacks the inductive biases (translation equivariance, locality) of CNNs and tends to need larger pretraining datasets to match CNN performance. CNN backbones or hybrid architectures may work better with limited data.
Quick scan · 2–4 min
Everything important, quickly.
6m: patch tokenization pipeline, CLS token role and alternatives, why bidirectional attention is the right choice for images, position encoding variants, hierarchical Swin alternative, and downstream applications (CLIP, DINOv2, SAM).
Code example
python
import torch
import torch.nn as nn
class PatchEmbed(nn.Module):
def __init__(self, img_size=224, patch_size=16, in_chans=3, d_model=768):
super().__init__()
self.proj = nn.Conv2d(in_chans, d_model,
kernel_size=patch_size, stride=patch_size)
num_patches = (img_size // patch_size) ** 2
self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, d_model))
def forward(self, x): # x: (B, 3, 224, 224)
x = self.proj(x) # (B, d_model, 14, 14)
x = x.flatten(2).transpose(1, 2) # (B, 196, d_model)
cls = self.cls_token.expand(x.size(0), -1, -1)
x = torch.cat([cls, x], dim=1) # (B, 197, d_model)
x = x + self.pos_embed
return x
Real-world examples
Real products, models, and research that use this idea.
1OpenAI CLIP uses a ViT image encoder paired with a text encoder for joint image-text contrastive learning.
2Meta DINOv2 (released 2024) trains ViT self-supervised on internet-scale image data, producing the standard pretrained image backbone for downstream tasks.
3Segment Anything Model (SAM) uses a ViT-H backbone for the image encoder; the foundation of zero-shot segmentation.
4DETR (DEtection TRansformer) uses ViT-style attention for end to end object detection.
5Llama 4 Maverick and Gemini 3.1 Pro use ViT-style patch tokenization to encode images for multimodal input alongside text tokens.
Interviewer follow-ups
What an interviewer would ask next. Try answering before peeking at the approach.
QWhat happens when you run a ViT trained at 224 x 224 on a 384 x 384 image?
Sign in or go Pro to see the approach for this follow-up.
QWhy doesn't ViT use a CNN stem for the patch projection?
Sign in or go Pro to see the approach for this follow-up.
QHow does Swin's hierarchical attention differ from ViT's flat attention?
Sign in or go Pro to see the approach for this follow-up.
Don't say this
Red flags and common mistakes that signal junior thinking. Click to expand.
Most common mistake
Using a causal mask for ViT. Images have no temporal direction, all patches exist simultaneously, so attention is fully bidirectional. Causal masking would arbitrarily block patches from seeing later patches in the (artificial) raster order.