Vision Transformers for Image Recognition at Scale
By: Abhilash Pattnaik
Tags: Image Recognition, Transformation, NLP, Artificial Intelligence

While convolutional neural networks have been used in computer vision since the 1980s, they were not at the forefront until 2012 when AlexNet surpassed the performance of contemporary state-of-the-art image recognition methods by a large margin.

Two factors helped enable this breakthrough: (i) The availability of training sets like ImageNet, and (ii) The use of commoditized GPU hardware, which provides significantly more compute for training.

As such, since 2012, convolutional neural networks have become the go-to model for vision tasks.

The benefit of using a convolutional neural network is to avoid the need for hand-designed visual features, instead of learning to perform tasks directly from data end-to-end. However, while the neural network avoids hand-crafted feature-extraction, the architecture itself is designed specifically for images and can be computationally demanding. Looking forward to the next generation of scalable vision models, one might ask whether this domain-specific design is necessary, or if one could successfully leverage more domain agnostic and computationally efficient architectures to achieve state-of-the-art results.

As the first step in this direction is the Vision Transformer, a vision model-based as closely as possible on the Transformer architecture which was originally designed for text-based tasks. Vision Transformer represents an input image as a sequence of image patches, similar to the sequence of word embeddings used when applying Transformers to text, and directly predicts class labels for the image. The transformer will demonstrate excellent performance when trained on sufficient data, outperforming a comparable state-of-the-art convolutional neural network with four times fewer computational resources.

The Vision Transformer treats an input image as a sequence of patches, akin to a series of word embeddings generated by a Natural Language Processing Transformer.

Vision Transformer

The original text Transformer takes as input a sequence of words, which it then uses for classification, translation, or other Natural Language Processing (NLP) tasks. The transformer can be designed to make it operate directly on images instead of words, and observe how much about image structure the model can learn on its own.

Vision Transformer divides an image into a grid of square patches. Each patch is flattened into a single vector by concatenating the channels of all pixels in a patch and then linearly projecting it to the desired input dimension. Because Transformers are agnostic to the structure of the input elements, there’s an added learnable position embedding to each patch, which allows the model to learn about the structure of the images. A prior model would not comprehend the relative location of patches in the image, or even that the image has a 2D structure, it must learn such relevant information from the training data and encode structural information in the position embeddings.

Left: Performance of Vision Transformer when pre-trained on different datasets. Right: Vision Tranformeryields a good performance/compute trade-off.

High-Performing Large-Scale Image Recognition

The data suggest that (1) with sufficient training the Vision Transformer can perform very well, and (2) It yields an excellent performance/compute trade-off at both smaller and larger compute scales. Therefore, to see if performance improvements carried over to even larger scales, it is trained a 600M-parameter Vision Transformer model.

This large model attains state-of-the-art performance on multiple popular benchmarks, including 88.55% top-1 accuracy on ImageNet and 99.50% on CIFAR-10. It also performs well on the cleaned-up version of the ImageNet evaluations set “ImageNet-Real”, attaining 90.72% top-1 accuracy. Finally, the Transformer works well on diverse tasks, even with few training data points. For example, on the VTAB-1k suite (19 tasks with 1,000 data points each), ViT attains 77.63%, significantly ahead of the single-model state of the art (SOTA) (76.3%), and even matching SOTA attained by an ensemble of multiple models (77.6%).