Vision transformer

1 Architecture Overview

A Vision Transformer (ViT) adapts the transformer architecture, originally developed for sequence-to-sequence tasks in natural language processing, to image recognition. It treats an image as a sequence of fixed-size patches and processes them through a standard transformer encoder, discarding the convolutional inductive biases inherent in CNNs. The core components are patch embedding, position embeddings, a transformer encoder, and a classification head.

1.1 Patch Embedding

Patch embedding converts raw image pixels into a sequence of feature vectors suitable for transformer input. It consists of two steps: partitioning the image into non-overlapping patches and linearly projecting each flattened patch into a lower-dimensional embedding space.

1.1.1 Image Partitioning into Patches

An input image of size \(H \times W \times C\) is divided into \(N\) patches of size \(P \times P\), where \(N = (H / P) \times (W / P)\). Typical patch sizes are 16×16 or 14×14 pixels. Each patch is flattened into a vector of length \(P^2 \cdot C\).

1.1.2 Linear Projection of Flattened Patches

Each flattened patch is multiplied by a learnable weight matrix (linear projection) to produce a patch embedding of dimension \(D\) (typically 768 or 1024). The resulting sequence of \(N\) embeddings is combined with a special [CLS] token (see Section 1.4.1) to form the input to the transformer encoder.

1.2 Position Embeddings

Because the transformer encoder is permutation-invariant, position embeddings are added to the patch embeddings to retain spatial information. ViT originally uses learnable 1D position embeddings, but variants have explored 2D alternatives.

1.2.1 Learnable 1D Position Embeddings

A learnable vector of dimension \(D\) is assigned to each patch index (0 to \(N-1\)) and to the [CLS] token. These embeddings are added element-wise to the patch embeddings before entering the encoder.

1.2.2 2D Position Embeddings Variants

Some extensions, such as the original ViT paper’s ablation, considered 2D position embeddings that encode both row and column coordinates separately. While 2D embeddings provide a stronger spatial prior, the 1D approach yields comparable performance and is simpler.

1.3 Transformer Encoder

The transformer encoder in ViT is identical to the one used in NLP transformers. It consists of alternating layers of multi-head self-attention (MSA) and feed-forward networks (FFN), each preceded by layer normalization and followed by residual connections.

1.3.1 Multi-Head Self-Attention (MSA)

MSA computes attention scores between every pair of patches. The input sequence of \(N+1\) embeddings (including [CLS]) is linearly projected into queries, keys, and values across \(h\) heads. Attention weights are obtained by softmax over key-query dot products, then used to aggregate values. This allows the model to capture long-range dependencies across the image.

1.3.2 Feed-Forward Network (FFN)

Each encoder layer includes a two-layer MLP with a GeLU activation. The FFN projects the attention output from dimension \(D\) to a hidden dimension (e.g., \(4D\)) and back, providing additional non-linear transformations.

1.3.3 Layer Normalization and Residual Connections

Layer normalization is applied before both MSA and FFN (pre-norm formulation), and residual connections add the block input to its output. This stabilizes training and facilitates gradient flow through many layers.

1.4 Classification Head

The final representation used for classification is derived from the [CLS] token, though global average pooling over patch tokens is an alternative.

1.4.1 [CLS] Token Approach

A learnable [CLS] token is prepended to the patch sequence. After passing through the encoder, the output corresponding to this token serves as an aggregate image representation. It is fed to a linear classifier (or MLP head) to produce class logits.

1.4.2 Global Average Pooling Alternative

Instead of using a [CLS] token, the model can take the mean of all patch token outputs as the image representation. This approach is employed in some variants (e.g., Swin Transformer) and can be more parameter-efficient.

2 Training and Implementation

Training ViT effectively requires large-scale datasets and specific regularization techniques due to its lack of built-in spatial inductive biases. The standard pipeline involves pre-training on massive image collections followed by fine-tuning on target datasets.

2.1 Pre-Training and Fine-Tuning

ViT benefits significantly from pre-training on large datasets, after which it can be fine-tuned to downstream tasks with moderate data.

2.1.1 Large-Scale Pre-Training (e.g., ImageNet-21k, JFT-300M)

The original ViT was pre-trained on ImageNet-21k (14 million images, 21k classes) or JFT-300M (300 million images, 18k classes). This large-scale supervised pre-training learns strong generic visual features that transfer well to smaller datasets.

2.1.2 Fine-Tuning on Downstream Datasets

After pre-training, the classification head is replaced and the entire model is fine-tuned on a target dataset (e.g., ImageNet-1k, CIFAR-100). Fine-tuning often uses a lower learning rate and fewer epochs than pre-training.

2.2 Data Augmentation and Regularization

To compensate for weak inductive biases and prevent overfitting, ViT relies on extensive data augmentation and regularization.

2.2.1 Mixup, CutMix, and RandAugment

Mixup creates convex combinations of pairs of images and labels, CutMix replaces rectangular regions from one image with another, and RandAugment applies a random selection of augmentation operations (e.g., rotation, color jitter) with controlled magnitude. These techniques are crucial for ViT’s generalization.

2.2.2 Dropout and Stochastic Depth

Dropout is applied to the attention weights or the output of the MSA and FFN modules. Stochastic depth (randomly dropping entire layers during training) is also used to improve regularization and training efficiency.

2.3 Optimization Hyperparameters

ViT training demands careful tuning of learning rate schedules and regularization parameters.

2.3.1 Learning Rate Schedules (Warmup, Cosine Decay)

A linear warmup over the first few thousand steps increases the learning rate from zero to a peak value, followed by cosine decay to zero. This schedule stabilizes early training and leads to better convergence.

2.3.2 Weight Decay and Optimizer Choice (AdamW)

AdamW is the optimizer of choice for ViT, applying decoupled weight decay. A typical weight decay value is 0.1 for pre-training and moderate decay for fine-tuning. The weight decay acts as an additional regularizer.

3 Variants and Extensions

Since the introduction of ViT, numerous variants have been proposed to address its limitations (e.g., data efficiency, computational cost) or extend its capabilities.

3.1 DeiT (Data-efficient Image Transformers)

DeiT introduces training strategies that enable ViT to perform well on medium-sized datasets (e.g., ImageNet-1k) without external pre-training on gigantic corpora. The key innovation is knowledge distillation from a CNN teacher.

3.1.1 Knowledge Distillation with a CNN Teacher

A CNN (e.g., RegNet) serves as a teacher model, and the student ViT is trained to mimic the teacher’s soft labels in addition to the ground-truth labels. A special distillation token is added to the transformer input to facilitate this.

3.1.2 Hard Distillation vs. Soft Distillation

Soft distillation uses the teacher’s probability distribution as training signal, while hard distillation uses the teacher’s predicted class label as a hard target. DeiT shows that hard distillation works better in practice.

3.2 Swin Transformer

Swin Transformer introduces a hierarchical architecture and shifted window attention to overcome ViT’s quadratic complexity and lack of multi-scale features.

3.2.1 Shifted Window Multi-Head Self-Attention

Self-attention is computed within local windows (e.g., 7×7 patches), and windows are shifted between consecutive layers to allow cross-window connections. This reduces complexity from \(O(N^2)\) to \(O(N \cdot M)\) where \(M\) is the window size.

3.2.2 Hierarchical Feature Maps

Unlike ViT’s single-scale output, Swin Transformer progressively downsamples the spatial resolution (like a CNN), producing feature maps at multiple scales. This makes it more suitable for dense prediction tasks like object detection and segmentation.

3.3 Other Notable Variants

Several other ViT-inspired models have been developed, each addressing specific aspects.

3.3.1 ViT with Convolutional Stem (CvT)

CvT replaces the linear patch embedding with a small convolutional stem that captures low-level features and introduces a convolution-based token mixing mechanism, improving performance and efficiency.

3.3.2 Pyramid Vision Transformer (PVT)

PVT adopts a pyramid structure similar to Swin but uses spatial reduction attention to reduce complexity. It also includes overlapping patch embedding to retain more local information.

3.3.3 Cross-Attention Vision Transformer (CrossViT)

CrossViT employs a dual-branch architecture with different patch sizes, cross-attention between branches to fuse multi-scale information, and a token fusion mechanism, enhancing performance on classification and segmentation.

4 Performance and Applications

ViT achieves state-of-the-art results on several vision benchmarks, especially when scaled with large pre-training datasets. It has also been adapted to a wide range of tasks beyond classification.

4.1 Benchmark Results

ViT has been evaluated extensively on ImageNet and other standard datasets, often surpassing traditional CNNs.

4.1.1 ImageNet Classification Accuracy

On ImageNet-1k, a large ViT (ViT-L/16) pre-trained on JFT-300M achieves 88.55% top-1 accuracy, comparable to the best CNNs at the time. Smaller variants, such as ViT-B/16 with moderate pre-training, achieve around 77–79% top-1 accuracy.

4.1.2 Comparison with State-of-the-Art CNNs (ResNet, EfficientNet)

ViT outperforms ResNet and EfficientNet of similar model size when pre-trained on large datasets, but on ImageNet-1k alone, CNNs often remain competitive or superior with fewer parameters. The gap is larger on large-scale pre-training regimes.

4.2 Applications Beyond Image Classification

ViT’s architecture has been adapted to various vision tasks through specialized heads or hybrid designs.

4.2.1 Object Detection and Instance Segmentation (DETR, ViTDet)

DETR (Detection Transformer) uses a transformer encoder-decoder based on ViT for end-to-end object detection. ViTDet replaces the CNN backbone in detection pipelines with a ViT, achieving competitive results on COCO.

4.2.2 Semantic Segmentation (SETR, SegFormer)

SETR (SEgmentation TRansformer) treats segmentation as a sequence-to-sequence task using a ViT encoder. SegFormer combines a hierarchical transformer encoder with a lightweight MLP decoder for efficient and accurate segmentation.

4.2.3 Video Understanding (Video Vision Transformer)

Video ViT extends ViT to video by treating spatiotemporal patches (cubes) as tokens. It processes a sequence of frames, achieving strong results on action recognition and video classification.

4.3 Computational Efficiency and Scalability

ViT’s scaling behavior and computational footprint are important considerations for real-world deployment.

4.3.1 Complexity Analysis (Quadratic vs. Linear Attention)

Standard self-attention has quadratic complexity \(O(N^2)\) in the number of patches, which becomes prohibitive for high-resolution images or long video sequences. Variants like Swin and Performer reduce this to linear or near-linear complexity.

4.3.2 Hybrid CNN-Transformer Architectures

Hybrid models combine a CNN-based early stem (to extract low-level features) with a transformer encoder (for global context). This reduces patch count while retaining efficiency, offering a good balance between performance and computation.

5 Limitations and Future Directions

Despite its successes, ViT faces several challenges that motivate ongoing research.

5.1 Data Hungriness and Lack of Inductive Biases

ViT does not encode translation equivariance or locality priors, making it reliant on large-scale pre-training to learn these properties from data. On small datasets, it often underperforms CNNs without heavy augmentation.

5.2 Scalability with Limited Data

When pre-training data is scarce, ViT’s performance degrades significantly. Methods like DeiT and self-supervised pre-training (see Section 5.4.2) aim to mitigate this issue.

5.3 Computational Cost of Self-Attention

Quadratic attention complexity remains a bottleneck for high-resolution images and dense prediction tasks. Efficient attention mechanisms and hierarchical designs address this but often at the cost of some accuracy.

5.4 Ongoing Research Areas

Current research focuses on improving efficiency, reducing data requirements, and expanding ViT’s role in multimodal systems.

5.4.1 Efficient Attention Mechanisms (Linformer, Performer)

Linformer approximates self-attention with linear projections, while Performer uses random Fourier features to achieve linear complexity. These methods enable ViT to process longer sequences with lower memory.

5.4.2 Self-Supervised Pre-Training for Vision Transformers (MAE, SimMIM)

Masked Autoencoders (MAE) and SimMIM train ViT by masking random patches and reconstructing the missing pixels. These self-supervised methods achieve strong performance without labeled data, reducing ViT’s data hungriness.

5.4.3 Integration with Multimodal Models (CLIP, Flamingo)

ViT is used as the visual backbone in multimodal models such as CLIP (contrastive language-image pre-training) and Flamingo (few-shot vision-language). These models link images and text, enabling tasks like zero-shot classification and image captioning.