Vision Transformer (ViT)

A vision transformer (ViT) is a transformer model adapted for computer vision, processing images as sequences of patches rather than text tokens. It was introduced in 2020 as an alternative to convolutional neural networks.

A vision transformer (ViT) is a type of Transformer (architecture) designed for computer vision. Unlike convolutional neural networks (CNNs) that process pixels through convolutional filters, a ViT decomposes an input image into a series of square patches, serializes each patch into a vector, and maps it to a lower-dimensional embedding using a single matrix multiplication. These patch embeddings are then processed by a transformer encoder, which applies Multi-Head Attention mechanisms to capture relationships between patches. ViTs were introduced as alternatives to CNNs, offering different inductive biases, training stability, and data efficiency. They are generally less data-efficient than CNNs but have higher capacity, and some of the largest modern computer vision models, such as one with 22 billion parameters, are ViTs.

The ViT architecture was first proposed in 2020 and quickly achieved state-of-the-art results in image classification, overcoming the previous dominance of CNNs. Subsequent research produced many variants, including hybrid architectures that combine features of both ViTs and CNNs. ViTs have found applications in image recognition, image segmentation, weather prediction, and autonomous driving.

History

Transformers were introduced in the 2017 paper "Attention Is All You Need" and became widely used in natural language processing. In 2019, a paper applied transformer ideas to computer vision by starting with a ResNet, a standard CNN, and replacing all convolutional kernels with the self-attention mechanism. This approach yielded superior performance but was not a true vision transformer.

In 2020, an encoder-only transformer was adapted for computer vision, yielding the ViT, which reached state of the art in image classification. The masked autoencoder (2022) extended ViT to work with unsupervised training. These developments also stimulated new advances in convolutional neural networks, leading to cross-fertilization between the two approaches.

In 2021, several important ViT variants were proposed to improve efficiency, accuracy, or domain suitability. Two studies improved efficiency and robustness by adding a CNN as a preprocessor. The Swin Transformer achieved state-of-the-art results on object detection datasets such as COCO by using convolution-like sliding windows of attention and a pyramid structure.

Overview

The basic architecture of the original 2020 ViT is a BERT-like encoder-only transformer. The input image is represented as a tensor of shape H × W × C, where H, W, and C are height, width, and number of channels (typically RGB). The image is split into square patches of size P × P × C. Each patch is flattened and passed through a linear layer to obtain a patch embedding. A positional encoding, which encodes the patch's position in the image, is added to the embedding. The original paper compared no embedding, 1D, 2D, and relative embeddings, and adopted 1D.

The sequence of patch embeddings is then processed by several transformer encoder layers. The attention mechanism in a ViT repeatedly transforms representation vectors of image patches, incorporating semantic relations between patches, analogous to how transformers in NLP capture relations between words.

To use the output for downstream tasks, an additional head is trained. For classification, a shallow MLP (linear-GeLU-linear-softmax) is added on top, which outputs a probability distribution over classes.

Variants

Original ViT

The original ViT was an encoder-only transformer trained with supervision to predict image labels from patches. It used a special [CLS] token in the input, and the corresponding output vector served as the input to the final MLP head. This architectural hack allowed the model to compress all label-relevant information into one vector.

Transformers initially found success in NLP, as seen in models like BERT and GPT-3. In contrast, typical image processing used CNNs, with well-known examples including Xception, ResNet, EfficientNet, DenseNet, and Inception. Transformers measure relationships between pairs of input tokens, with cost quadratic in the number of tokens. For images, computing relationships for every pixel pair is prohibitive, so ViT computes relationships among pixels in small sections (e.g., 16×16 pixels), drastically reducing cost. Each section is flattened, multiplied by an embedding matrix, and added to a positional embedding before being fed to the transformer.

Pooling

After processing, the ViT produces embedding vectors that must be converted to a single class prediction. The original ViT and Masked Autoencoder used a dummy [CLS] token, following BERT. The output at [CLS] is processed by a LayerNorm-feedforward-softmax module.

Global average pooling (GAP) instead takes the average of all output tokens as the classification token, and was mentioned in the original paper as equally good. Multihead attention pooling (MAP) applies a multi-head attention block to pool the vectors, taking a list of vectors as input and producing a weighted combination.

Applications and Impact

ViTs have been applied to a wide range of computer vision tasks beyond classification, including object detection, segmentation, and video understanding. They are also used in medical imaging and remote sensing. The architecture has influenced the development of large-scale vision models and has been integrated into multimodal systems that combine vision and language.

Despite their advantages, ViTs require more data to train effectively than CNNs, which led to the development of hybrid models and techniques like knowledge distillation and self-supervised pretraining. The masked autoencoder approach enabled unsupervised learning, reducing the need for labeled data.

ViTs have also spurred innovation in hardware and software optimization, as their attention mechanisms are computationally intensive. Research continues on improving efficiency, interpretability, and robustness.