Paper Explained 1: Vision Transformer
Extending Transformer to Computer Vision Tasks
Paper link: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
The Vision Transformer model, also called ViT, has revolutionized computer vision (CV) by introducing transformer-based architectures originally proposed for NLP to image processing tasks. ViT demonstrated that self-attention mechanisms could effectively model global dependencies in images, outperforming convolutional neural networks (CNNs) in many tasks with sufficient data and computational resources.
It revolutionized areas like image classification, object detection, and segmentation, establishing new benchmarks. Furthermore, ViT has also been widely applied in multi-modality modeling tasks, such as visual-language models such as CLIP, DALL·E and more recent large vision-language models such as LLAVA, by seamlessly integrating textual and visual data and enabling robust cross-modal understanding and generation.
This article will present a comprehensive introduction to the Vision Transformer model. Below are the topics we will cover:
- Background: revisit CNN architectures and their limitations, Transformer model in NLP and why it is needed for CV tasks, etc.
- Vision Transformer Architecture.
- Training of ViT.
- Advantages and Limitations.
Background
CNN Architectures in Computer Vision
The ViT model was proposed in 2021, and at that time convolutional architectures are still dominant in the CV community, especially the ResNet-like architectures.
And there is a good reason for that. To understand that, let’s briefly revisit two important operations in convolutional networks: convolution and pooling.
Below is an illustration of how a convolution operation works.
In convolution, a filter (also known as a kernel) is applied to the source input. A kernel is typically a square matrix with common size 3x3 or 5x5, where larger kernel size can capture larger area within an image but come at the cost of reduced spatial resolution and increased computational cost.
Another two core parameters in convolution are stride and padding.
Strideis the number of pixels by which the kernel moves as it slides over the image. For example, a stride of 1 simply means the kernel moves one pixel at a time. Increasing stride will reduces the output dimensions, which can help decrease computational cost and control overfitting but at the loss of some image detail.Paddingmeans adding an appropriate number of rows and columns (typically of zeros) to the input image borders. This ensures that the convolution kernel fits perfectly at the borders, allowing the output image to retain the same size as the input image, which is crucial for deep networks to allow the stacking of multiple layers.
During the convolution process, this kernel will slide over the source layer, filling the pixels in the destination layer one by one. This operation can capture spatial features like edges, corners, and textures.
To see how convolution helps capture edges, consider the example below: The input is an 6x6 image with a vertical edge, with light on the left and dark on the right. After convolution with the vertical edge detection kernel, it will result in detection of the vertical edge, with the detected edge shown in the middle of the output image.
Pooling is the operation to reduce the spatial size of the representation from the results of convolution. A pooling layer in CNN down samples the spatial dimensions of the input feature maps and reduces their size while preserving important information.
Common pooling methods include max pooling and average pooling. Below is an illustration of a max pooling operation with 2 × 2 filter and stride of 2:
Now it is easier to understand why Convolutional networks are so popular in CV applications. Basically, it has several key benefits that especially suitable for processing images:
Spatial Awareness: CNNs use convolutional layers with kernels that slide across the image to capture spatial features like edges, corners, and textures. This process captures local spatial relationships between pixels, allowing the network to “understand” the structure of the image.Translation Invariance: During the sliding process, the same features (e.g., an edge or a pattern) will be detected regardless of their location, and the following pooling layers also makes feature detection more robust to small translations by down-sampling feature maps. This ensures that objects can be recognized even if they appear in different parts of the image, making the model more generalizable.Parameter Sharing: Instead of having unique weights for each pixel, the same kernel is applied across the entire image, which drastically reduces the number of parameters compared to fully connected networks, lowering memory requirements and computational complexity.Hierarchical Feature Learning: CNNs usually employ a stack of convolutional layers, with early convolutional layers detect simple features like edges and textures, and then as data passes through deeper layers, these simple features are combined to form more complex patterns like shapes or parts of objects. This mimics how humans perceive images, starting from basic shapes to full object recognition, making CNNs highly effective for vision tasks.
The figure below shows the model architecture for ResNet-34, where you can see how multiple convolutional blocks are stacked together to extract feature in a hierarchical manner.
So far we have explained why CNNs are so popular for vision tasks before ViT was proposed. Next we will talk about why and how the CV community try to adopt Transformer.
Why We Need Transformer?
While CNNs are popular choices for computer vision tasks, it is important to acknowledge their limitations as well, many of which stem from the fundamental characteristics of convolution:
Global Context Understanding: CNNs primarily focus on local patterns due to the size of their kernels and rely on stacking layers to gradually expand their receptive field. This can make it challenging to capture global context, especially in complex scenes.Scalability: As CNNs grow increasingly deeper (for example, ResNet can have as many as 152 layers!), computational and memory requirements increase significantly. They also require careful tuning of architectures (e.g., kernel sizes, strides).Dynamic Feature Representation: Convolutional kernels are fixed during inference, and their size limits the patterns they can detect. This can lead to rigid feature extraction.
These limitations have driven the computer vision community to explore Transformers as an alternative.
The Transformer model was proposed in 2018, and has revolutionized many NLP tasks since its introduction. For a detailed introduction to Transformer models, you may refer to my previous articles listed below:
- Transformer Part 1 : Understanding The Attention Mechanism in Transformers with Code
- Transformer Part 2: Understanding Positional Encoding in Transformers and Beyond with Code
- Transformer Part 3: Understanding Transformer Models with Code
- BERT: A Comprehensive Introduction
The success of Transformer models in NLP also inspired a lot of efforts in the CV community by trying to apply it to handle visual tasks, motivated mainly by its computational efficiency and scalability. To some extent, this aligns with the trend of increasing dataset sizes in CV area at that time, such as ImageNet and custom datasets.
Before ViT
There are many works trying to extend Transformer to CV area. For example, several early works tried to combine self-attention mechanism with CNN architectures, by treating pixels as tokens and then let each pixel attend to every other pixel. However, this design clearly cannot scale to realistic input sizes due to the quadratic cost in the number of pixels.
Following that idea, several approximations have been explored to overcome the above issue, for example applying the self attention only in local neighborhoods instead of globally, leveraging Sparse Transformers, or applying self-attention in blocks of varying sizes. However, many of these specialized attention architectures require complex engineering to be implemented efficiently on hardware accelerators.
This motivates the authors of ViT to explore an alternative approach, by applying a standard Transformer directly to images, with the fewest possible motivations.
To achieve this, we have to handle two major challenges:
How to generate the input sequence: images are not naturally represented as sequences. Instead, they have a spatial structure and are typically represented as high-dimensional tensors (e.g., width × height × channels). A naive solution would be flattening an entire image into a sequence of pixels, but as we have seen before this will create extremely long sequences, which are computationally expensive and make it impossible to scale to realistic resolutions due to the quadratic complexity of self-attention.How to handle the loss of local spatial relationships: in CNNs, locality, two-dimensional neighborhood structure, and translation equivariance are naturally baked into each layer through the hierarchical feature extraction, but the vanilla Transformer model treats every token equivalently, making it difficult to capture the local spatial relationships.
In the following sections we will see how ViT mitigate both issues and finally brings ViT to state-of-the-art.
Vision Transformer (ViT) Architecture
Image Pre-Processing: Convert Images to a Sequence of Patches
Given an input image, ViT firstly resize it to a predefined resolution, and then go through the following steps:
Image Splitting: The input image (e.g., H×W) is divided into non-overlapping patches of size P×P. This results in (H/P)(W/P) patches.Flattening: Each patch is flattened into a 1D vector of size PP*C, where C is the number of channels (e.g., 3 for RGB images).Linear Embedding: Each flattened vector is passed through a learnable linear projection to create a fixed-length (D-dimensional) embedding.Positional Encoding: Learnable positional encodings are added to the patch embeddings to retain spatial information.Add CLS token: A [CLS] token is prepended to the sequence of embedded patches, and its state at the output of the Transformer encoder serves as the representation for the entire image.
Note that the image splitting step can also be implemented by a 2D convolution, where
Patch Size as Kernel Size: The patch size P is treated as the kernel size in the convolution operation.Stride Equals Patch Size: The stride of the convolution is set to P, ensuring that the patches are non-overlapping.Number of Output Channels: The convolution’s output channels are set to D, the desired embedding dimension for each patch.
Pre-LayerNorm Transformer in ViT
Having converted the an input image into a sequence of patches or visual tokens, the next step is to feed these visual tokens into Transformer encoder.
The figure below compares the Transformer layer in the original Attention is All You Need paper vs. the Transformer layer used in ViT, which is essentially the differences between a Post-LayerNorm vs. a Pre-LayerNorm design.
More specifically, in the pre-LayerNorm architecture, LayerNorm is applied before the attention and feed-forward operations, which can help control the scale of activations and improve convergence during training. This helps stabilize training, especially in deeper models like ViT.
Training of ViT
Data Requirements
When training a ViT model, one thing we need to keep in mind is that ViT has much less image-specific inductive bias than CNNs, as a result it requires learning from much more data to achieve the same level of capability in terms of image understanding.
What does that mean?
Inductive bias refers to the assumptions or prior knowledge that a machine learning model makes about the data when learning from it. These assumptions can guide the model’s learning process and help it generalize to unseen data.
In the case of CNNs, they assume that images have local structures (like edges, textures, and patterns) that repeat across the image, and this assumption allows CNNs to learn efficiently from image data by detecting these repeating patterns through small filters, rather than learning from scratch for every part of the image.
In other words, CNNs can learn faster and perform pretty well even with less data, because they are already assuming certain patterns exist in the data. On the other hand, for models do not have inductive bias, such as ViT, they have to learn everything from scratch, which requires much more data and computing power.
Pre-Training and Finetuning
ViTs are pre-trained either via self-supervised or supervised learning on large datasets. The goal is to learn useful patch representations through the Transformer model.
After pre-training, the model is fine-tuned on specific tasks using labeled data. Fine-tuning adapts the pre-trained model to the target task, improving performance.
Advantages and Limitations
We start with advantages of ViTs first:
Global Attention and Context: Transformers can capture long-range dependencies and global relationships between image patches, making it better at understanding overall image composition and spatial relationships compared to CNNs’ local receptive fields . This is particularly useful for tasks requiring global context like scene understanding.Flexibility and scalability: ViT scales well with more data and larger models and can also handle variable input sizes. Meanwhile, it can handle a variety of vision tasks such as classification, detection, segmentation etc., with relatively small architectural changes.Parallel Processing: Self-attention operations can be computed in parallel, making ViT potentially faster than sequential CNN.
Then comes the limitations of ViT models:
Data Hungry: ViT requires much larger datasets for training from scratch compared to CNNs, and may not perform well with limited data unless pre-trained, and thus presenting a higher risk of overfitting on smaller datasets.Computational Complexity: due to the adoption of Transformer, it has quadratic complexity O(n²) with respect to input size and is memory intensive. This can be prohibitively expensive for high-resolution images.Loss of Inductive Biases: Unlike CNNs that have built-in assumptions about locality and translation invariance, ViTs need to learn these properties from data and may struggle with tasks where local features are crucial.
Summary
In this article we present a comprehensive introduction to the Vision Transformer models, covering aspects including a revisit of CNNs, why and how Transformers are extended to CV tasks, model architecture of ViT models, training procedures, as well as its advantages and limitations.
Since its publication, Vision Transformers have revolutionized computer vision in tasks like image classification, object detection, and segmentation, particularly when trained on large-scale datasets.
Meanwhile, ViTs have also opened new opportunities in multi-modal modeling, enabling the integration of vision and language tasks such as image captioning, visual question answering, and image-text alignment.