Part of Language AI Handbook
Explains how Vision Transformers work: from image patching and embeddings to self-attention. Topics include ViT architecture, inductive biases.
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Vision Transformer
For years, computer vision and natural language processing followed parallel but distinct trajectories, each developing specialized architectures suited to their respective modalities. While NLP underwent a revolutionary transformation with the transformer architecture, computer vision remained dominated by convolutional neural networks (CNNs). Architectures including ResNet, DenseNet, EfficientNet achieved impressive results by exploiting the inductive biases inherent in images: locality, which assumes that nearby pixels matter more than distant ones, and translation equivariance, which ensures that a cat remains a cat regardless of where it appears in the frame. These architectural priors allowed CNNs to learn efficiently from limited data, as the network structure itself encoded fundamental assumptions about how visual information is organized. Yet these same inductive biases that made CNNs sample-efficient also constrained their ability to model long-range dependencies and to scale gracefully with increasing amounts of data.
The Vision Transformer (ViT), introduced by Dosovitskiy et al. in 2020, posed a counterintuitive question: What if we set aside the convolutional inductive biases entirely and treat an image exactly like we treat text in a transformer? Instead of sliding convolutional filters across pixels to detect local features hierarchically, we could slice the image into uniform patches, linearly embed each patch into a vector representation, and process the resulting sequence with standard self-attention mechanisms. This approach seemed almost naive at first glance. CNNs had been engineered over decades to exploit the structure of images; removing those carefully designed priors should theoretically hurt performance, especially when working with limited training data where architectural constraints provide necessary regularization.
Yet when scaled appropriately with sufficient data and compute, ViT matched the performance of state-of-the-art CNNs and then surpassed them across numerous benchmarks. More importantly, ViT established a unified architectural foundation that enabled multimodal models to emerge. By representing both images and text as sequences of embeddings processed by identical transformer blocks, we can build systems that reason across modalities natively, using the same computational primitives for vision and language. As we discussed in Part XIII: Self-Attention, the transformer architecture's power derives from its ability to model arbitrary dependencies between sequence positions, regardless of distance. ViT demonstrated that this power extends naturally beyond discrete text tokens to continuous visual data, challenging the assumption that different modalities require fundamentally different architectural approaches.
Understanding ViT requires appreciating it as a conceptual shift rather than a technical curiosity. CNNs encode prior beliefs about images into their very structure. ViT instead asks: can we give the model enough data to discover those patterns on its own? The answer turns out to be yes, provided the dataset is large enough. This insight changes how we think about architecture design in deep learning more broadly: the more data you have, the less you need to bake prior knowledge into the architecture, and the more you can rely on the model to learn structure from scratch.
In this chapter, we explore in detail how ViT converts images into sequences through the patching and embedding process, examine the specific architectural details that differ from standard NLP transformers, and understand why ViT requires specific pre-training strategies to overcome its lack of built-in visual inductive biases. We trace the flow from raw pixels to classification logits, understanding both the mechanics and the theoretical underpinnings of this paradigm shift in computer vision.
Image Patching: From Pixels to Sequences
The fundamental challenge in applying transformers to images is a dimensional mismatch between what the architecture expects and how visual data is structured. A standard transformer encoder expects input as a sequence of vectors , where is the sequence length and is the embedding dimension. A typical image, however, arrives as a 3D tensor representing height and width plus color channels. This structural difference poses an immediate problem: transformers process sequences, while images are inherently two-dimensional grids with channel information.
The solution ViT adopts is elegant in its simplicity: reshape the image into a sequence by dividing it into a regular grid of fixed-size square patches. Each patch is flattened into a single vector by concatenating all pixel values across spatial locations and color channels. These flattened vectors then serve as the input sequence to the transformer, with each patch playing the role that a word token plays in a language model.
The Patching Operation in Detail
Consider a standard input image of size pixels with 3 color channels (RGB). If we choose a patch size of pixels per side, the patching proceeds as follows. First, we divide the image into a grid of non-overlapping patches, since patches fit along each spatial dimension. The total number of patches is:
Each individual patch is a sub-image containing values. We flatten this 3D sub-image into a 1D vector by concatenating all pixel values in row-major order across height and width for every channel. More precisely, for the patch at grid position where , we extract:
where:
- : the flattened patch vector at grid row , grid column
- : the slice of the image tensor corresponding to the patch at position , with all channels included
- : the patch size (number of pixels per side, e.g., 16)
- : the number of color channels (3 for RGB)
- : the operation that reshapes a tensor into a vector of length
We then flatten the 2D grid of patches into a 1D sequence by linearizing the grid positions. Conventionally, this is done in row-major order: patches in the first row of the grid appear first in the sequence, followed by patches in the second row, and so on. Patch gets sequence index . We now have a sequence of vectors, each of dimension 768. This is a sequence length and dimensionality comparable to what standard language transformers process, making the same architecture applicable with minimal modification.
It is worth pausing to appreciate what has been lost and what has been retained in this transformation. We have retained the pixel values of every region of the image; no information has been discarded or aggregated. What we have lost is the explicit 2D spatial structure. The sequence representation no longer encodes the fact that patch 1 is adjacent to patch 2 horizontally, or that patch 15 lies directly below patch 1. The transformer receives an unordered bag of patch vectors with positional indices, and it must learn from those indices and from training data which patches are spatially related. This is the inductive bias gap that we will discuss in depth later.
The choice of patch size involves fundamental trade-offs that affect both computational efficiency and model capacity:
- Smaller patches (e.g., ): Produce longer sequences ( patches for a image), capturing finer spatial resolution and preserving more detail. However, self-attention has complexity, so quadrupling patch count leads to roughly 16 times more attention computation.
- Larger patches (e.g., ): Produce shorter sequences ( patches), reducing computational cost significantly, but at the expense of fine-grained spatial detail. Individual patches encompass more image area, potentially obscuring small objects or fine textures within a single token.
Standard ViT models use patches as a balanced compromise. ViT variants optimized for high-resolution tasks (such as detection and segmentation) often use patches and pair them with efficient attention mechanisms to manage the quadratic cost.
The patching approach also imposes an implicit tokenization granularity. A word token in NLP typically carries semantic content: the word "cat" has a meaning that the model can use directly. A patch of pixels, by contrast, often contains no clear semantic unit. It might span the boundary between a cat's ear and the background, or contain only a uniform textured region with no object at all. This means patch-level representations start at a much lower semantic level than word tokens, requiring the model to learn compositional semantics through attention across patches, rather than beginning with semantically meaningful units. This is another reason why ViT needs more data: it must learn both low-level visual features and their high-level composition, whereas NLP models begin with tokens that already carry semantic weight encoded by the vocabulary.
Patch Embeddings and the CLS Token
Once we have flattened the image patches into vectors, we need to project them into the transformer's embedding dimension . This process is analogous to word embeddings in NLP, where we map discrete vocabulary indices to dense continuous vectors. In ViT, we apply a trainable linear projection to each patch vector.
The embedding for the -th patch is computed as:
where:
- : the embedding vector for the -th patch after projection
- : the flattened -th patch (input vector), where is the patch size and is the number of channels
- : the trainable linear projection matrix mapping from patch dimension to embedding dimension
- : the patch index, ranging from to (total number of patches)
This linear projection is a learned weight matrix, initialized randomly and trained end-to-end along with the rest of the model. Conceptually it serves the same purpose as an embedding table in NLP: it maps the raw input representation to a space where meaningful relationships can be captured and transformed by the downstream attention layers. Unlike a discrete embedding table that has one learned vector per vocabulary entry, the patch embedding matrix is a general linear transformation applied uniformly to every patch, regardless of its content. The model must therefore learn a transformation that works for all possible patches it might encounter, a more challenging requirement that again explains why more pre-training data is beneficial.
An important practical equivalence exists between this explicit flatten-then-project operation and a convolutional layer applied with stride and no padding. Both operations extract non-overlapping windows of size , apply the same linear map to each window, and produce a spatial grid of feature vectors. ViT implementations therefore sometimes use nn.Conv2d(in_channels=3, out_channels=D, kernel_size=P, stride=P) as the patch embedding layer, which is computationally equivalent to the explicit formulation but benefits from highly optimized convolution routines on modern hardware.
The Classification Token
Following the convention established by BERT (which we covered in BERT Architecture), ViT prepends a special learnable embedding to the beginning of the sequence. This classification token (CLS token) serves a unique role in the architecture. It begins as a learnable parameter initialized randomly, with no inherent content. As information flows through the transformer layers via self-attention, the CLS token can attend to every patch in the sequence, aggregating global information from all spatial locations.
Why use a separate CLS token for classification rather than, say, averaging all patch representations? The averaging approach, called global average pooling, is also used in practice and works well. The CLS token approach is conceptually cleaner for sequence classification because the model can learn to selectively aggregate the information that is most relevant for the classification decision, rather than uniformly combining all spatial information. Some patches might be background clutter that contributes little to the class label; a learnable aggregation via the CLS token can in principle learn to downweight such patches. In practice, both approaches yield similar accuracy, and the choice is sometimes made for architectural consistency with BERT-style models.
Position Embeddings and the Full Input Sequence
After prepending the CLS token, we add learned position embeddings to inject spatial information. The complete input sequence to the transformer is:
where:
- : the complete input sequence to the transformer, with position embeddings added element-wise
- : the learnable CLS token embedding prepended to the sequence at position 0
- : the linearly projected -th patch embedding
- : the matrix of position embeddings, with one row per sequence position including the CLS token
- : vertical concatenation (stacking vectors row-by-row to form a matrix)
- : total number of patches ( in the running example)
The here denotes element-wise addition: each row of the concatenated patch matrix is added to the corresponding row of the position embedding matrix. This means every token in the sequence, including the CLS token, receives a position-specific bias vector that encodes its index in the sequence. After this addition, the model's self-attention layers can use these position signals to infer spatial relationships.
Position Embeddings in Vision
Recall from Part XIV: Positional Encoding that transformers process sequences in parallel and have no inherent notion of token order or sequence position. In NLP, we explored sinusoidal encodings, learned embeddings, and relative position representations to inject ordering information. For vision, the question is more complex: images have inherently 2D spatial structure, but our patching operation has flattened them into a 1D sequence. How do we encode 2D spatial relationships when the sequence is 1D?
The original ViT used standard 1D learned position embeddings, treating the sequence of patches exactly like a sequence of word tokens in a language model. Each position index has an associated learnable vector that is added to the token at that position. These position embeddings are initialized randomly and trained alongside the rest of the model. Remarkably, this simple approach works well in practice despite its apparent disregard for the 2D nature of images.
Why does 1D learned position encoding work for 2D images? The key insight is that the 2D spatial arrangement of patches maps directly to a specific 1D ordering in row-major format. Patch always appears at position . The model can learn position embeddings such that patches at similar 2D positions receive similar position embeddings. If patch at position (row 1, column 0) and patch at position (row 0, column 1) both have position embeddings that share useful features, the model can learn spatial adjacency implicitly. Research has confirmed that after training, the learned 1D position embeddings indeed exhibit 2D spatial structure: when you compute cosine similarities between position embeddings, spatially adjacent patches have more similar embeddings than distant ones.
Alternative position encoding strategies that have been explored in ViT variants include:
-
2D position embeddings: Use separate row and column embeddings, summed to produce a composite spatial representation:
where is the learned embedding for row and is the learned embedding for column . This approach explicitly factors spatial position into orthogonal row and column components, which may make it easier for the model to generalize to images of different resolutions.
-
Relative position encodings: Instead of encoding absolute positions, encode the relative offset between any two patches. When computing attention between patch and patch , a bias term is added that depends on and . This approach generalizes naturally to images of different sizes and can better capture local spatial relationships.
-
2D sinusoidal encodings: Extend the sinusoidal position encoding from Sinusoidal Position Encoding to two dimensions by creating sinusoidal patterns across both height and width axes independently, then concatenating or summing them.
Empirically, the differences between these approaches are relatively small for classification tasks on fixed-resolution images, suggesting that the transformer is remarkably effective at learning spatial relationships from 1D position embeddings given sufficient training. For tasks requiring fine-grained spatial reasoning or adaptation to variable resolutions, 2D or relative position encodings offer modest improvements. The flexibility of attention-based position encoding, compared to the fixed positional assumptions embedded in CNN architectures, is itself an advantage of the ViT approach.

The block-diagonal structure in the heatmap reveals a key property: patches from the same row of the spatial grid (separated by 7 positions in the 1D sequence) show higher similarity to each other than patches from distant rows. This demonstrates that learned 1D position embeddings spontaneously develop 2D spatial awareness during training, even without an explicit 2D structure in the encoding scheme.
ViT Architecture
With patches converted to embeddings and positional information injected, the remainder of the architecture follows the standard transformer encoder we detailed in Encoder Architecture. The embeddings pass through identical transformer blocks, each consisting of several key components working in concert.
The four core components of each transformer block are:
-
Multi-Head Self-Attention (MSA): As covered in Multi-Head Attention, this allows each patch to attend to all other patches in the sequence, modeling global relationships from the very first layer. This is in stark contrast to CNNs, where receptive fields grow gradually with depth and only achieve global context in the final layers. In ViT, a patch in the upper-left corner of an image can directly attend to a patch in the lower-right corner in layer 1, allowing the model to capture long-range visual relationships that would require many convolutional layers to achieve.
-
Feed-Forward Network (MLP): A two-layer MLP with GELU activation and hidden dimension typically set to , following the transformer convention from Feed-Forward Networks. Each position in the sequence passes through this MLP independently. This provides non-linear transformation capacity beyond what the attention mechanism alone can achieve. The MLP can be understood as transforming each patch representation in place, while attention mixes information across patches.
-
Layer Normalization: Applied before each sublayer (Pre-Norm configuration), as we discussed in Pre-Norm vs Post-Norm. Pre-Norm stabilizes training by ensuring that the inputs to each sublayer have controlled scale, which is particularly important for deep networks where vanishing or exploding gradients can derail training.
-
Residual Connections: Applied around both the attention and MLP blocks, creating gradient highways essential for training deep networks, as explored in Residual Connections. These connections allow gradients to flow directly back through the network to early layers without passing through the non-linearities of intermediate layers.
The Layer Equations
The equations for layer process the input through self-attention and feed-forward sublayers with residual connections:
where:
- : input to layer , the output from layer (or for the first layer)
- : intermediate representation after the Multi-Head Self-Attention sublayer and its residual connection
- : output of layer after the MLP sublayer and its residual connection
- : layer index, ranging from to (total number of layers)
- : Multi-Head Self-Attention function applied to all tokens in parallel
- : Feed-Forward Network applied position-wise with GELU activation and expansion ratio 4
- : Layer Normalization applied before each sublayer (Pre-Norm configuration)
The symmetry between the NLP transformer and ViT is striking: after the patch embedding step, the equations are identical. The same PyTorch TransformerEncoderLayer module used for BERT or GPT can be used for ViT without modification. This architectural equivalence is what makes unified multimodal models possible.
From CLS Token to Classification Logits
After layers, the final hidden state of the CLS token is extracted and passed to a classification head. This head is typically a single linear layer during fine-tuning, and a two-layer MLP during pre-training:
where maps from the embedding dimension to class logits. The CLS token is used here because it has no direct input content (it begins as a learned null token), which means its final representation is entirely determined by what it has aggregated through attention over the other tokens. It acts as a summary of the entire image, specialized through training to carry information relevant to the classification objective.
The computational cost of self-attention is , where is the number of patches. For a image with patches, , which is manageable on modern hardware. However, if we increase image resolution to , the number of patches grows to , and the attention computation becomes roughly 27 times more expensive due to the quadratic scaling. This limits ViT's ability to handle high-resolution images directly, motivating sparse attention patterns and hierarchical architectures such as the Swin Transformer, which we explore in the efficient attention chapters referenced from Part XVII.
The Inductive Bias Problem
The central trade-off of Vision Transformers is that they lack the visual inductive biases that make CNNs sample-efficient. To understand why this matters, it helps to examine exactly what these biases are and how they are embedded into the CNN architecture.
A CNN exploits two key architectural priors through the structure of convolution:
Locality: Nearby pixels form coherent patterns such as edges and textures alongside local color gradients. CNNs assume that meaningful features can be detected within small spatial neighborhoods. A convolutional filter explicitly processes only a region of the input at each position, ignoring all other regions. This forced locality constraint is a strong prior: it tells the network that relevant visual features are local. This is approximately true for low-level vision (edge detectors, blob detectors) and approximately true for mid-level features (corners, junctions) but less true for high-level semantics (object category), where global context often matters.
Translation equivariance: The same visual features should be detected regardless of their absolute position in the image. CNNs implement this through weight sharing: the same filter weights are applied at every spatial location. A horizontal edge detector learned for the top-left region automatically works for the bottom-right region, because the same weights slide across the entire image. This prior directly addresses the fact that the visual world is shift-invariant: a cat in the corner of an image looks the same as a cat in the center.
ViT has no such architectural assumptions. It must learn from data alone that adjacent patches correspond to adjacent image regions, and that visual features that appear at one spatial location also appear at others. Without weight sharing, there is no mechanism to enforce that what the model learns about patch 1 (top-left) applies to patch 100 (lower-right). The same cat, appearing at different positions, would produce different patch vectors, and the model must learn to handle this variation through exposure to many examples at many positions.
This is why the original ViT paper found that when pre-trained on ImageNet-21k (14 million images), ViT matched or exceeded ResNet performance on downstream tasks, but when trained only on ImageNet-1k (1.3 million images) from scratch, ResNets significantly outperformed ViT. The transformer architecture, without built-in visual priors, requires far more data to independently discover the locality and translation invariance that CNNs receive for free by architectural design.
The situation is analogous to learning grammar from scratch versus learning it from a teaching grammar. A child acquiring language receives structural guidance from the humans around them that constrains the hypothesis space. A model learning from raw text must infer grammatical structure from distributional statistics alone. Given enough data, both approaches produce equivalent linguistic competence, but the data-guided approach requires far more exposure to raw data to compensate for the absent structural guidance.


The visualization makes the contrast concrete. In a CNN, early layers have tiny receptive fields (the small inner circle) and global context only emerges after many stacked layers. In ViT, every patch can directly reach every other patch in the very first layer, with no sequential accumulation of context required. This global connectivity is simultaneously ViT's greatest architectural strength and its greatest data challenge.
Pre-training Strategies
To overcome the data inefficiency inherent in the absence of inductive biases, ViT relies on large-scale pre-training. The original paper explored two main approaches, and subsequent research has expanded these significantly.
Supervised Pre-training
The most direct approach involves pre-training ViT on large supervised image classification datasets. The model is trained with standard cross-entropy loss to predict image labels from the CLS token representation. After pre-training on a large dataset, the classification head is replaced or fine-tuned for specific downstream tasks.
The two primary supervised pre-training datasets used in the original ViT work were:
- ImageNet-21k: 14 million images across 21,000 classes. This provides broad visual diversity. Pre-training on this dataset was enough for ViT to match or exceed well-tuned ResNets on ImageNet-1k classification.
- JFT-300M: 300 million images across 18,000 classes, an internal Google dataset. Training at this scale demonstrated ViT's strongest performance and clearest advantages over CNNs in terms of accuracy and scaling behavior.
The key practical concern with supervised pre-training is the requirement for labeled data at scale. While ImageNet-21k is publicly available, collecting and labeling 300 million images requires substantial resources. This limitation motivated the development of self-supervised pre-training approaches that can learn from unlabeled data.
An important pre-training detail for supervised ViT is that the classification head used during pre-training differs from the one used during fine-tuning. During pre-training on JFT-300M, ViT uses an MLP head with one hidden layer and tanh activation. For fine-tuning on downstream datasets, a single linear layer is used instead, initialized to zero. The zero initialization of the fine-tuning head encourages the model to preserve the pre-trained representations at the start of fine-tuning, avoiding the disruption that a randomly initialized large head might cause.
Self-Supervised Pre-training: Masked Autoencoders
More recent approaches eliminate the need for manual labels by learning representations directly from unlabeled image data. The Masked Autoencoder (MAE) approach, introduced by He et al. in 2021, draws direct inspiration from masked language modeling in BERT (which we covered in Masked Language Modeling).
MAE works by randomly masking a high proportion of input patches (typically 75%) and training the model to reconstruct the pixel values of the masked patches from the visible ones. The high masking ratio is intentional: it creates a challenging task that prevents the model from solving the reconstruction problem through simple linear interpolation or nearest-neighbor copying. To recover a heavily masked image, the model must develop an internal representation that captures the semantic and structural regularities of natural images.
The architecture uses an asymmetric encoder-decoder design that provides an efficiency advantage:
The encoder is a standard ViT that processes only the visible (unmasked) patches. Because 75% of patches are masked, the encoder handles only 25% of the sequence. This reduces the computational cost of the encoder by roughly compared to processing all patches, enabling larger batch sizes and faster training iterations.
The decoder is a much smaller transformer that takes both the encoded visible patches and learned mask tokens (placeholders for the missing patches) as input. The mask tokens are shared across all masked positions (the same vector is used for every masked patch). The decoder reconstructs all patches in their original spatial positions, with the loss computed only on the masked patches using mean squared error between predicted and actual pixel values.
The asymmetry is strategically important: making the decoder lightweight means the encoder bears the full burden of learning good visual representations. The decoder, being simple and narrow, cannot compensate for a weak encoder through its own capacity. This forces the encoder to develop rich semantic representations rather than relying on decoder complexity to smooth over impoverished features.


Other self-supervised approaches applied to ViT include:
DINO (Self-Distillation with No Labels): Uses knowledge distillation between a student network and a teacher network that is a momentum-averaged version of the student. The student and teacher both receive different augmented views of the same image, and the student is trained to predict the teacher's output distribution. DINO-trained ViTs develop particularly strong feature representations that exhibit striking emergent properties: the self-attention maps of DINO-trained models often segment objects from background without any segmentation supervision, revealing that the model has learned object boundaries as a side effect of self-supervised visual understanding.
BEiT (BERT Pre-training of Image Transformers): Instead of reconstructing raw pixel values, BEiT converts each patch into a discrete token from a visual vocabulary learned by a separate vector quantization model (a visual tokenizer). The model then performs masked patch prediction in this discrete token space, much like BERT predicts masked word tokens. This approach provides a cleaner target signal than raw pixel reconstruction, which may focus the model on learning semantically meaningful features rather than low-level texture details.
Hybrid Approaches
Some architectures combine CNNs and transformers to capture the benefits of both paradigms. A CNN stem, typically using the early layers of a ResNet, processes the image first to extract low-level features while preserving spatial inductive biases. The resulting feature maps are then patched and fed into the transformer for global modeling. This hybrid approach offers a useful engineering compromise: the CNN stem provides strong inductive bias that enables efficient learning from limited data, while the transformer backbone enables global reasoning and flexible scaling with more data. Hybrid ViT models trained on ImageNet-1k alone can outperform pure ViT models trained on the same data, since the CNN stem compensates for ViT's lack of visual priors in the data-scarce regime.
Scaling Properties
ViT exhibits fundamentally different scaling characteristics compared to CNNs. While ResNet performance improves predictably but with diminishing returns as model size increases, ViT demonstrates steeper scaling curves when data is sufficient. With large datasets, larger ViT models continue to improve significantly, while CNN performance plateaus at a lower ceiling. This aligns with the scaling laws we explored in Part XXIII: Scaling Laws: transformer-based models follow power-law improvement with scale when data is abundant.
The original ViT paper introduced a family of standardized model sizes using naming conventions borrowed from BERT:
- ViT-Tiny: embed_dim=192, depth=12, num_heads=3 (~5.7M parameters)
- ViT-Small: embed_dim=384, depth=12, num_heads=6 (~22M parameters)
- ViT-Base (ViT-B/16): embed_dim=768, depth=12, num_heads=12 (~86M parameters)
- ViT-Large (ViT-L/16): embed_dim=1024, depth=24, num_heads=16 (~307M parameters)
- ViT-Huge (ViT-H/14): embed_dim=1280, depth=32, num_heads=16, patch_size=14 (~632M parameters)
The notation ViT-B/16 means Base model with patches; ViT-H/14 uses patches instead. Moving to smaller patches for the Huge model increases the sequence length and allows the model's greater capacity to capture finer spatial detail.

The crossover point in the scaling curve reflects the core trade-off: CNNs start with strong inductive priors that provide a head start in the low-data regime, but these same priors impose a ceiling on what the model can learn. ViT starts with no visual priors, requiring more data to learn the basics, but its more flexible representational capacity allows it to keep improving at larger scales.
Worked Example: Tracing a Patch Through ViT
To solidify the mechanics, let's trace a single patch from raw pixels all the way to the transformer encoder output, working through the dimensions explicitly.
Consider an RGB image of size and a model configuration of patch size , embedding dimension , 12 transformer layers, and 12 attention heads. We focus on the patch at grid position , the second patch in the first row.
Step 1: Extract the patch. The patch at corresponds to pixels in the image tensor. Slicing this region gives a sub-tensor of shape .
Step 2: Flatten. We flatten this sub-tensor into a 1D vector: . In Python: patch_vec = patch.reshape(-1) yields a vector . (We index from 1, so is patch 2 in the sequence.)
Step 3: Linear projection. The patch embedding matrix transforms the raw patch vector to an embedding: .
Step 4: Add position embedding. The sequence index of this patch is in 0-indexed patch coordinates, which becomes position 2 in the full sequence (after the CLS token at position 0). We add the position embedding: .
Step 5: Transformer processing. The full sequence (196 patches plus 1 CLS token) passes through 12 transformer layers. In each layer, self-attention is computed across all 197 tokens. Our patch at position 2 can receive information from the CLS token (position 0) and all other 195 patches. After 12 layers, the representation reflects a complex mixture of information from all spatial locations, weighted by learned attention patterns.
Step 6: Classification. The final representation of the CLS token, , is passed to the linear classification head. Our individual patch representation is not directly used for classification but contributed to the CLS token's representation through attention in all 12 layers.
This trace shows that the path from raw pixels to a classification decision passes through a sequence of linear transformations and attention-based mixing operations, all operating in the same -dimensional space. No special-purpose vision operations are required after the initial patching step.
Code Implementation
Let's implement the core components of ViT to solidify our understanding. We'll build a simplified but complete version that processes images through patching and embedding before the transformer layers.
import torch
torch.manual_seed(42)We start with the patch embedding layer. This module takes an image tensor and produces a sequence of patch embeddings.
import torch.nn as nn
class PatchEmbed(nn.Module):
"""
Image to Patch Embedding.
Divides input image into non-overlapping patches and linearly projects
each flattened patch to the embedding dimension.
"""
def __init__(
self, img_size=224, patch_size=16, in_channels=3, embed_dim=768
):
super().__init__()
self.img_size = img_size
self.patch_size = patch_size
self.n_patches = (img_size // patch_size) ** 2
# Linear projection: maps P*P*C flattened patch to embed_dim
self.proj = nn.Linear(patch_size * patch_size * in_channels, embed_dim)
def forward(self, x):
# x shape: (batch_size, channels, height, width)
B, C, H, W = x.shape
# Extract non-overlapping patches using unfold
# unfold(dim, size, step) slides a window of 'size' along 'dim' with 'step'
x = x.unfold(2, self.patch_size, self.patch_size) # (B, C, H//P, P, W)
x = x.unfold(
3, self.patch_size, self.patch_size
) # (B, C, H//P, W//P, P, P)
# Permute to (B, H//P, W//P, C, P, P) then flatten spatial+channel dims
x = x.permute(0, 2, 3, 1, 4, 5)
x = x.reshape(B, self.n_patches, -1) # (B, N_patches, C*P*P)
# Project to embedding dimension: (B, N_patches, embed_dim)
x = self.proj(x)
return xNow let's test the patch embedding and inspect the output dimensions:
dummy_img = torch.randn(2, 3, 224, 224) # Batch of 2 images
patch_embed = PatchEmbed(img_size=224, patch_size=16, embed_dim=768)
patches = patch_embed(dummy_img)Input image shape: [2, 3, 224, 224] Patch embedding output: [2, 196, 768] Number of patches: 196 Patch dimensionality before projection: 768 Patch dimensionality after projection: 768
The output confirms our images are converted to sequences of 196 patch embeddings, each of dimension 768. Now let's implement the transformer building blocks. We implement multi-head attention first, returning attention weights for later visualization:
class MultiHeadAttention(nn.Module):
"""Scaled dot-product multi-head self-attention with optional attention weight output."""
def __init__(
self, dim, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
):
super().__init__()
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.scale = self.head_dim**-0.5
# Single linear maps all tokens to Q, K, V in one pass
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, N, C = x.shape
# Compute Q, K, V and reshape for multi-head processing
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
qkv = qkv.permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2] # Each: (B, heads, N, head_dim)
# Scaled dot-product attention
attn = (q @ k.transpose(-2, -1)) * self.scale # (B, heads, N, N)
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
# Aggregate values and project
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x, attn # Return attention weights for visualization
class MLP(nn.Module):
"""Feed-forward network with GELU activation and 4x hidden expansion."""
def __init__(
self, in_features, hidden_features=None, out_features=None, drop=0.0
):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features * 4
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = nn.GELU()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return xThe transformer block combines these two sublayers with Pre-Norm layer normalization and residual connections:
class TransformerBlock(nn.Module):
"""Standard ViT transformer block with Pre-LayerNorm."""
def __init__(
self,
dim,
num_heads,
mlp_ratio=4.0,
qkv_bias=False,
drop=0.0,
attn_drop=0.0,
):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = MultiHeadAttention(
dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
attn_drop=attn_drop,
proj_drop=drop,
)
self.norm2 = nn.LayerNorm(dim)
self.mlp = MLP(
in_features=dim, hidden_features=int(dim * mlp_ratio), drop=drop
)
def forward(self, x, return_attention=False):
# Pre-norm attention with residual: z'_l = MSA(LN(z_{l-1})) + z_{l-1}
attn_out, attn_weights = self.attn(self.norm1(x))
x = x + attn_out
# Pre-norm MLP with residual: z_l = MLP(LN(z'_l)) + z'_l
x = x + self.mlp(self.norm2(x))
if return_attention:
return x, attn_weights
return xNow the complete Vision Transformer, which assembles these components with the CLS token, position embeddings, and classification head:
class VisionTransformer(nn.Module):
"""Vision Transformer for image classification."""
def __init__(
self,
img_size=224,
patch_size=16,
in_channels=3,
num_classes=1000,
embed_dim=768,
depth=12,
num_heads=12,
mlp_ratio=4.0,
drop_rate=0.0,
):
super().__init__()
self.embed_dim = embed_dim
self.patch_embed = PatchEmbed(
img_size, patch_size, in_channels, embed_dim
)
num_patches = self.patch_embed.n_patches
# Learnable CLS token: starts as noise, trained to aggregate image info
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
# Learned 1D position embeddings: one per patch + one for the CLS token
self.pos_embed = nn.Parameter(
torch.zeros(1, num_patches + 1, embed_dim)
)
self.pos_drop = nn.Dropout(p=drop_rate)
# L identical transformer blocks
self.blocks = nn.ModuleList(
[
TransformerBlock(
dim=embed_dim,
num_heads=num_heads,
mlp_ratio=mlp_ratio,
drop=drop_rate,
)
for _ in range(depth)
]
)
# Final layer norm and classification head
self.norm = nn.LayerNorm(embed_dim)
self.head = nn.Linear(embed_dim, num_classes)
# Initialize with small normal noise to break symmetry
nn.init.normal_(self.cls_token, std=0.02)
nn.init.normal_(self.pos_embed, std=0.02)
def forward(self, x, return_attention=False):
B = x.shape[0]
# Step 1: Patch embedding
x = self.patch_embed(x) # (B, N, D)
# Step 2: Prepend CLS token
cls_tokens = self.cls_token.expand(B, -1, -1) # (B, 1, D)
x = torch.cat((cls_tokens, x), dim=1) # (B, N+1, D)
# Step 3: Add position embeddings
x = x + self.pos_embed
x = self.pos_drop(x)
# Step 4: Process through transformer layers
attn_weights = []
for block in self.blocks:
if return_attention:
x, attn = block(x, return_attention=True)
attn_weights.append(attn)
else:
x = block(x)
# Step 5: Extract CLS token and classify
x = self.norm(x)
cls_output = x[:, 0] # CLS token only: (B, D)
logits = self.head(cls_output) # (B, num_classes)
if return_attention:
return logits, attn_weights
return logitsLet's instantiate a small model suitable for demonstration, verify the output dimensions, and check the parameter count:
# Small ViT configuration for demonstration
model = VisionTransformer(
img_size=224,
patch_size=16,
embed_dim=192, # Tiny model: ViT-Tiny width
depth=4, # Shallow for fast demo
num_heads=3,
num_classes=10,
)
total_params = sum(p.numel() for p in model.parameters())
test_input = torch.randn(1, 3, 224, 224)
output = model(test_input)
sample_logits = output.squeeze()[:5]Model parameters: 1,965,130 Input shape: [1, 3, 224, 224] Output shape: [1, 10] Sample logits (first 5): ['-0.094', '0.791', '0.309', '-0.657', '-0.169']
The compact model configuration yields a manageable parameter count while maintaining all the architectural properties of full-scale Vision Transformers. Now let's visualize the core patching operation that turns a 2D image into a 1D token sequence:


The conversion from the structured 2D image on the left to the flat sequence on the right captures the core challenge: the sequence representation loses the two-dimensional adjacency that is natural in the image. Patches that were side-by-side are now adjacent in the sequence, but patches that were vertically adjacent (14 positions apart in row-major order) appear far apart. The model must learn from position embeddings and training data to reconstruct these spatial relationships.
Finally, let's visualize how self-attention distributes weights across patches:

The attention heatmap illustrates two distinct learned behaviors. The top row (CLS token) exhibits diffuse, roughly uniform attention across all patches. This reflects its role as a global aggregator that must summarize the entire image for classification. The remaining rows show concentration along the diagonal. This reflects locality: patches attend most strongly to spatially adjacent patches. No convolutional architecture enforced this locality; the model discovered it from data because local correlations are a measurable statistical property of natural images. ViT learns the same visual priors as CNNs, but through gradient descent rather than architectural design.
Key Parameters and Model Variants
The key parameters for Vision Transformer configurations, along with their practical implications, are:
- img_size: Size of the input image in pixels (assumes square). Controls how many patches are produced; larger images create more patches and increase attention computation quadratically.
- patch_size: Side length of each square patch in pixels. The standard choice is 16. This provides 196 patches for a image. ViT-H uses 14 to give the largest model finer spatial resolution.
- embed_dim: Embedding dimension . Common values are 192 (Tiny), 384 (Small), 768 (Base), 1024 (Large), and 1280 (Huge). This is the primary driver of model width and parameter count.
- depth: Number of transformer blocks . Ranges from 4 in lightweight configurations to 32 in ViT-Huge. Deeper models have more capacity for compositional reasoning but require more memory and computation.
- num_heads: Number of attention heads. Must evenly divide embed_dim. More heads allow the model to attend to different representation subspaces simultaneously, with each head specializing in different spatial relationships or feature types.
- mlp_ratio: Ratio of the MLP hidden dimension to embed_dim. Standard is 4.0, giving a hidden dimension of in the feed-forward layers. Higher ratios increase model capacity and parameter count.
- drop_rate: Dropout probability applied to position embeddings, MLP activations, and attention weights. Provides regularization during fine-tuning on smaller datasets.
Limitations and Impact
Vision Transformers achieved breakthrough performance on image classification, but their deployment introduces specific constraints that shape practical use and motivate ongoing research.
Data Efficiency
The most significant limitation is data efficiency. CNNs can achieve strong performance on ImageNet-1k (1.3 million images) when trained from scratch because their inductive biases provide powerful implicit regularization. ViT trained on ImageNet-1k alone performs substantially below ResNet baselines. Only with datasets like ImageNet-21k or JFT-300M does ViT match and exceed CNN performance, and even then, the gap is largest in the low-data regime. This means that practitioners working with domain-specific datasets of moderate size (tens of thousands to a few hundred thousand images) often find that fine-tuning a CNN baseline or using a hybrid architecture yields better results than fine-tuning a pure ViT, unless a large ViT pre-trained on a general-purpose dataset is available as the starting point.
Recent work has substantially reduced this gap. DeiT (Data-efficient Image Transformers) demonstrated that ViT can match ResNet performance on ImageNet-1k alone using improved training recipes: aggressive data augmentation, CutMix, Mixup, random erasing, repeated augmentation, and a knowledge distillation token that learns from a CNN teacher. These techniques collectively act as data multipliers, allowing the model to see more variation and partially compensate for the absent inductive biases. However, even with these improvements, ViT's sample efficiency remains below that of well-tuned CNNs in truly data-scarce scenarios.
Computational Scaling with Resolution
Computational scaling with image resolution presents another significant challenge. CNNs scale approximately linearly with pixel count for fixed filter sizes. ViT's attention scales quadratically with patch count. For a image with patches, , and the attention computation is approximately 27 times more expensive than for a image. This makes high-resolution vision tasks such as semantic segmentation, object detection, and medical image analysis computationally demanding.
This limitation motivated the development of hierarchical vision transformers. The Swin Transformer, for example, computes attention within local windows rather than globally, and uses shifted windows across layers to enable cross-window information exchange. This design achieves linear computational complexity with respect to image size and produces hierarchical feature maps at multiple resolutions, matching the feature pyramid structure of CNN-based detectors. Swin and its descendants have largely supplanted pure ViT in tasks requiring fine-grained spatial output, while pure ViT remains dominant for classification and representation learning where global attention is beneficial.
Fixed Patch Granularity
The fixed patch size creates a rigidity that CNNs with multi-scale architectures avoid. A patch might contain an entire small object, only the center of a large object, or purely background. The model must handle this semantic variability within a single sequence position, whereas CNN feature pyramids process the image at multiple scales simultaneously. This can limit performance on detection tasks where object sizes vary greatly, unless the model is adapted with multi-scale patch processing or additional architectural modifications.
Impact on Multimodal AI
Despite these limitations, ViT's influence extends far beyond classification accuracy benchmarks. By demonstrating that the transformer can process images without specialized architectural components, ViT enabled the unified multimodal models that now define the frontier of AI research. CLIP (Contrastive Language-Image Pre-training), which we will explore in the next chapter, uses ViT as its vision encoder alongside a standard text transformer. Because both modalities are processed by architecturally similar transformers producing similarly structured embeddings, CLIP can align visual and linguistic representations in a shared semantic space using contrastive training. This architectural symmetry would be far harder to achieve if vision required fundamentally different operations from language.
Beyond CLIP, ViT's architecture enables models like Flamingo, which processes interleaved sequences of images and text; CoCa, which uses a ViT encoder and transformer decoder for captioning and retrieval; and large vision-language models like GPT-4V and Gemini, which accept images and text as unified input sequences. The common thread is ViT's reduction of images to sequences of dense vectors, the same basic data type that language models have learned to process at scale.
ViT also validated the generality of transformer scaling laws beyond language. As we explored in Part XXIII: Scaling Laws, language models follow power-law relationships between performance and the product of model size, data size, and compute. ViT exhibited the same behavior: doubling model size or data consistently produced predictable accuracy improvements, without the saturation effects observed in CNN scaling. This has motivated investment in large-scale vision foundation models analogous to large language models, with capabilities that generalize across visual tasks through fine-tuning.
Summary
The Vision Transformer demonstrated that the transformer architecture, originally designed for discrete text sequences, could process continuous visual data with remarkable effectiveness through a single key insight: treating image patches as sequence tokens. By dividing images into fixed-size patches, linearly embedding them into vectors, and processing them with standard transformer encoders, ViT matched and surpassed state-of-the-art CNNs while eliminating the convolutional inductive biases that had defined computer vision for decades.
Key takeaways from this chapter include:
- Image patching converts a 2D image into a 1D sequence of patch embeddings, each of dimension , enabling transformer processing of visual data with no architectural modifications.
- Architectural parity: After the patching step, ViT uses identical components to NLP transformers: multi-head self-attention, MLP feed-forward layers, Pre-Norm layer normalization, and residual connections. This architectural equivalence is what makes unified multimodal systems possible.
- The CLS token is a learnable aggregate representation for classification, analogous to BERT's [CLS] token. It accumulates global information from all spatial locations through self-attention, then feeds a linear classification head.
- Position embeddings inject spatial information that the patching operation strips away. Standard 1D learned position embeddings work surprisingly well, with the model spontaneously developing 2D spatial structure in the learned embedding space.
- The inductive bias trade-off is central to understanding ViT: CNNs embed locality and translation equivariance into their architecture, enabling efficient learning from limited data. ViT must discover these properties from data, requiring much larger pre-training datasets to compensate for the absent structural priors.
- Pre-training strategies range from supervised classification on massive labeled datasets (ImageNet-21k, JFT-300M) to self-supervised masked autoencoding (MAE) and self-distillation (DINO), each offering different trade-offs between data requirements, label availability, and representation quality.
- Quadratic attention complexity with patch count limits ViT's application to high-resolution tasks, motivating hierarchical variants like Swin Transformer that reintroduce local processing while retaining the flexibility of attention.
- Scaling laws observed in NLP transfer directly to ViT: larger models trained on more data follow power-law improvement curves, paving the way for large-scale vision foundation models.
ViT's most lasting contribution is not any specific accuracy number but the architectural unification it demonstrated. By showing that a single sequence-to-sequence processing paradigm could handle both language and vision, ViT created the foundation for the multimodal AI systems that now define the frontier of the field. In the next chapter, we will explore CLIP, which uses ViT as its vision encoder alongside a text transformer, learning joint visual and linguistic representations through contrastive training at scale.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about Vision Transformers.
Vision Transformer Fundamentals
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of Language AI Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore Language AI HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.

Comments
No comments yet. Be the first to share your thoughts!