Vision Transformers: Revolutionizing Image Recognition with Self-Attention
Introduction
In recent years, the field of computer vision has been dominated by convolutional neural networks (CNNs), which have achieved remarkable success on tasks such as image classification, object detection, and semantic segmentation. However, a new challenger has emerged in the form of the Vision Transformer (ViT), introduced by Dosovitskiy et al. in their seminal paper "An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale" [1]. This work demonstrates that a pure transformer architecture, adapted from natural language processing (NLP), can outperform state-of-the-art CNNs on image recognition tasks when trained on sufficiently large datasets. In this article, we will delve into the key ideas behind ViT, with a particular focus on the patch embedding process that enables the application of transformers to images.
The Transformer Architecture
Before diving into Vision Transformers, let‘s briefly review the transformer architecture as originally proposed by Vaswani et al. for machine translation [2]. The core idea of the transformer is the self-attention mechanism, which allows the model to weigh the importance of different parts of the input sequence when computing a representation for each element.
In the encoder-decoder structure of the original transformer, the encoder consists of a stack of identical layers, each containing two sub-layers: a multi-head self-attention mechanism and a position-wise fully connected feed-forward network. The decoder has a similar structure but includes an additional sub-layer that performs multi-head attention over the encoder output.
The multi-head attention mechanism is the key component that enables the transformer to capture dependencies between elements of the input sequence. It works by linearly projecting the input embeddings into multiple lower-dimensional spaces (heads) and then computing a weighted sum of the values, where the weights are determined by the dot product of the queries and keys. This allows each head to focus on different aspects of the input, and the outputs of all heads are concatenated and linearly projected again to produce the final attention output.
Position-wise feed-forward networks are applied to each position separately and identically, consisting of two linear transformations with a ReLU activation in between. Layer normalization is applied before each sub-layer, and residual connections are used to facilitate gradient flow.
For a more detailed explanation of the transformer architecture, see [3].
Adapting Transformers to Vision
To apply transformers to image recognition tasks, Dosovitskiy et al. introduced several key modifications to the original architecture:
- Splitting an image into patches and linearly embedding them to form a sequence of input tokens.
- Adding learnable position embeddings to the patch embeddings to retain spatial information.
- Appending a classification token to the sequence to serve as a global representation of the image.
- Feeding the resulting sequence through a stack of transformer encoder layers.
- Using an MLP head on top of the classification token output for the final classification.
The most crucial aspect of this adaptation is the patch embedding process, which we will examine in detail in the next section.
The Patch Embedding Process
In order to apply a transformer to an image, we need to convert the 2D grid of pixels into a sequence of tokens that can be processed by the self-attention layers. ViT achieves this by splitting the image into a grid of fixed-size patches and then linearly embedding each patch to form a sequence of input tokens.
Formally, an image $x \in \mathbb{R}^{H \times W \times C}$ is reshaped into a sequence of flattened 2D patches $x_p \in \mathbb{R}^{N \times P^2 \cdot C}$, where $H$, $W$, and $C$ are the height, width, and number of channels of the original image, $P$ is the height and width of each patch, and $N = HW/P^2$ is the number of patches. The patches are then linearly projected to $D$ dimensions using a trainable matrix $E \in \mathbb{R}^{P^2 \cdot C \times D}$, resulting in a sequence of patch embeddings $z0 = [x{class}; x_p^1E; x_p^2E; …; x_p^NE]$.
To this sequence, a learnable classification token $x{class}$ is prepended, serving as a global representation of the image that can be used for classification. Learnable position embeddings $E{pos} \in \mathbb{R}^{(N+1) \times D}$ are added to the patch embeddings to retain positional information.
The resulting sequence of embeddings $z_0$ is then passed through a stack of $L$ transformer encoder layers, each consisting of multi-head self-attention (MSA) and multi-layer perceptron (MLP) blocks, with layer normalization (LN) applied before every block:
$$zl‘ = MSA(LN(z{l-1})) + z_{l-1}, \forall l \in [1, L]$$
$$z_l = MLP(LN(z_l‘)) + z_l‘, \forall l \in [1, L]$$
The output of the transformer encoder is a sequence of embeddings $zL$, where the embedding corresponding to the classification token $x{class}$ serves as a representation of the entire image. This embedding is passed through an MLP head to produce the final class probabilities:
$$y = MLP(LN(z_L^0))$$
The patch embedding process can be seen as a way to "tokenize" an image, transforming it from a 2D grid of pixels into a sequence of feature vectors that can be processed by a standard transformer encoder. By using a fixed patch size and linear projection, ViT can handle images of arbitrary resolution (up to memory constraints) without requiring any architectural changes.
Comparison to Convolutional Neural Networks
Convolutional neural networks (CNNs) have been the dominant approach in computer vision for the past decade, owing to their ability to capture local patterns and spatial hierarchies in images. CNNs rely on hand-crafted features such as convolutions and pooling operations that bake in certain assumptions about the structure of images, such as translation invariance and local connectivity.
In contrast, Vision Transformers make minimal assumptions about the image structure and instead rely on self-attention to capture both local and global dependencies between image patches. This allows ViT to learn more flexible representations that can adapt to the specific characteristics of the data.
One key advantage of self-attention over convolutions is the ability to capture long-range dependencies between distant parts of an image. In a CNN, the receptive field of a neuron is limited by the size of the convolution kernels and the depth of the network. In a transformer, each self-attention layer allows each patch to attend to all other patches in the image, regardless of their spatial distance. This global connectivity enables ViT to better capture complex relationships and contextual information.
However, the lack of inductive biases in ViT also means that it requires more training data and compute to achieve competitive performance compared to CNNs. In the original paper, Dosovitskiy et al. found that ViT outperformed state-of-the-art CNNs when pre-trained on the large-scale JFT-300M dataset (containing 300 million images) and fine-tuned on smaller downstream tasks. When trained on smaller datasets like ImageNet, ViT achieved lower accuracy than comparable CNNs.
This trade-off between flexibility and sample efficiency is a key consideration when choosing between CNNs and transformers for a given vision task. CNNs may be preferred when training data is limited or when the inductive biases of convolutions are well-suited to the problem at hand. Transformers may be preferred when there is an abundance of training data and when the ability to capture long-range dependencies is critical.
Analyzing Attention Patterns in ViT
One of the benefits of the self-attention mechanism in transformers is the interpretability of the learned attention patterns. By visualizing the attention maps of a trained ViT model, we can gain insights into how the model is using different image patches to make predictions.
In the original ViT paper, the authors analyzed the average distance between the query patch and the patches attended to in each head and layer of the network. They found that in the lower layers, the attention distances were highly variable, with some heads attending to most of the image while others focused on small regions around the query patch. In the higher layers, the attention distances increased, with most heads attending to almost all of the image.
This analysis suggests that the ViT model is able to capture both local and global information at different scales, adapting the receptive field size based on the requirements of the task. In contrast, the receptive field size in CNNs is fixed by the architecture and cannot be adaptively adjusted.
Further studies have used attention rollout [4] to compute the attention maps between the classification token and the input image patches, revealing which parts of the image are most important for the model‘s predictions. These visualizations can be used to debug and interpret the model‘s behavior, as well as to identify potential biases or failure modes.
Performance and Scaling
One of the key findings of the ViT paper is the impressive performance of the model when scaled up to large datasets and model sizes. The authors trained three different ViT variants – ViT-Base (86M parameters), ViT-Large (307M parameters), and ViT-Huge (632M parameters) – on the JFT-300M dataset and evaluated their performance on several benchmark tasks.
On the ImageNet dataset, ViT-Huge achieved a top-1 accuracy of 88.55%, outperforming the previous state-of-the-art CNN (EfficientNet-L2) by a significant margin. ViT-Huge also achieved state-of-the-art results on the CIFAR-100, Oxford-IIIT Pets, and Oxford Flowers-102 datasets when fine-tuned.
Interestingly, the authors observed that the performance of ViT continued to improve with increasing pre-training dataset size, whereas CNNs tended to saturate earlier. This suggests that given enough data and compute, transformers may have a higher performance ceiling than CNNs for image recognition tasks.
The following table compares the performance of ViT and CNNs on the ImageNet dataset:
| Model | Params | ImageNet Top-1 Acc. |
|---|---|---|
| ResNet-152 [5] | 60M | 78.3% |
| ViT-Base/16 | 86M | 77.9% |
| ViT-Large/16 | 307M | 76.5% |
| ViT-Huge/14 | 632M | 88.55% |
| EfficientNet-L2 | 480M | 88.4% |
While ViT-Huge outperforms all other models, it is also much larger in terms of parameter count. The base and large variants of ViT are more comparable to ResNets in size but achieve slightly lower accuracy when trained on ImageNet alone. However, when pre-trained on JFT-300M and fine-tuned on ImageNet, all ViT variants significantly outperform the ResNets.
These results demonstrate the potential of transformers for image recognition, especially when combined with large-scale pre-training. As more efficient transformer architectures and training methods are developed, we can expect to see further improvements in performance and scalability.
Conclusion and Future Directions
Vision Transformers have emerged as a powerful alternative to convolutional neural networks for image recognition tasks, showing impressive performance when scaled up to large datasets and model sizes. By adapting the self-attention mechanism from natural language processing to vision, ViT is able to capture long-range dependencies and learn flexible representations of images.
The key innovation of ViT is the patch embedding process, which enables the application of transformers to grid-like data such as images. By splitting an image into patches and linearly embedding them, ViT can process images of arbitrary resolution using a standard transformer encoder.
While ViT has shown great promise, there are still many open questions and avenues for future research. One direction is to explore more efficient transformer architectures and training methods that can reduce the computational cost and improve the sample efficiency of the model. For example, the DeiT [6] model introduces a distillation token to enable more efficient training of ViT on smaller datasets.
Another direction is to extend ViT to other vision tasks beyond image classification, such as object detection, semantic segmentation, and image generation. Some recent works, such as DETR [7] and ViT-FRCNN [8], have already begun to explore the use of transformers for these tasks, showing promising results.
Finally, there is still much to be learned about the properties and behaviors of self-attention in vision models. Further analysis of the learned attention patterns and representations could lead to new insights and improved architectures.
As the field of computer vision continues to evolve, it is clear that transformers will play an increasingly important role alongside CNNs. With their ability to capture long-range dependencies and learn flexible representations, transformers have the potential to unlock new levels of performance and capabilities in image recognition and beyond.
