8. Transformers › 8.5 The Visual Transformer and Applications

Section 8.5
The Visual Transformer and Applications

Unlike language, images are already in a numeric format, but they do not present as a 1D sequence. One option is to serialize the image as a sequence of pixels (i.e., flatten the image in some raster scan order) and treating each pixel as a token. However, pixels by themselves are very uninformative. Compare the information in a single pixel, for example, to that of an English or Chinese word. Not to mention the fact that there are way too many pixels in an image to treat each one as a token. A better approach is to divide the image into patches and represent each patch as a visual word, the combination of which describe the image.

This is image tokenization. We divide a \((3 \times H \times W)\)-image into a grid of \(P \times P\) patches. Each patch can be represented by a vector of length \(3P^2\) by flattening the pixels in the patch. We can then arrange the patches in some pre-defined raster scan order (e.g., top-left to bottom-right) as shown in Figure 116. The \(3P^2\)-dimensional vectors representing each patch are passed through a learned linear layer to obtain the image tokens.

Tokenizing images
Figure 116: Tokenizing images.

The visual transformer (ViT) model proposed by Dosovitskiy et al. [23] uses this approach to tokenize images. It then uses a transformer encoder model (with special class token <cls>) as an image classifier as shown in Figure 117. Just like a language model transformer, the visual transformer includes position encoding added to the patch embeddings, and the output token associated with the <cls> token is passed through a multi-layer perceptron with softmax activation.

The visual transformer (ViT) [23]
Figure 117: The visual transformer (ViT) [23].

8.5.1 Self-supervised Image Feature Learning

We can train the ViT model using self-supervision by a process known as model (self-)distillation unlike language models, which train for next-token prediction (also in a self-supervised manner). Once trained the <cls> token gives a good image representation that can be used for downstream tasks such as image retrieval, video segmentation, and zero-shot prediction. Remarkably the attention maps in DINO have also been found to be useful for (unsupervised) image segmentation and salience detection.

The DINO algorithm [13] (later improved by DINOv2 [78]) uses a pair of models with the same architecture, known as the teacher and student models. The idea is for the two models to learn to produce the same representation of slightly different random crops and transformations of an image. Figure 118 shows the DINO training setup and pseudo-code. Mathematically, we can write

\begin{align} x_1, x_2 &= T_1(x), T_2(x) \tag{268}\\ \theta_{t+1}^{(1)} &= \theta_{t}^{(1)} - \eta \fracdd{\theta^{(1)}} L^{\text{ce}}\left(p(x_1; \theta^{(1)}_t), p(x_2; \theta^{(2)}_t)\right) \tag{269}\\ \theta^{(2)}_{t+1} &= \alpha \theta^{(2)}_{t} + (1 - \alpha) \theta^{(1)}_{t+1} \tag{270}\end{align}

where superscript 1 indicates the student network and 2 indicates the teacher network consistent with Figure 118. The first line extracts two different random crops from the same image; the second line updates the student network parameters via gradient descent on a cross-entropy loss between student and teacher outputs; the last line updates the paramaters of the teacher network as an exponential moving average (EMA) of the student network paremeters instead of gradient descent. Here \(\alpha\) is a momentum factor that controls how slowly the teacher parameters track the student parameters, and typically set close to one. This helps prevent trivial solutions such as the model collapsing to produce a constant feature output (i.e., ignoring the input). To further mitigate against model collapse the teacher output is sharpened using softmax with a low temperature parameter. The logits from the teacher network are also centered by subtracting the average of past logits. DINOv2 [78] does a comprehensive analysis of different model components and adds various other tricks to further improve learning and avoid collapse.

Self-supervised image feature learning using the DINO algorithm [13]
Figure 118: Self-supervised image feature learning using the DINO algorithm [13].

There is nothing in the DINO algorithm that is specific to transformers, and the method has been demonstrated on architectures other than ViT too.

8.5.2 Object Detection with Transformers

One really novel application of visual transformers is object detectors. Carion et al. [12] proposed the detection transformer, or DETR, that uses a transformer decoder with learnable queries (i.e., target tokens) whose outputs correspond to latent object detections. These outputs are passed through a prediction head to regress a bounding box and object category (similar to R-CNN).

The full DETR architecture is depicted in Figure 119. Briefly, the input tokens to the encoder are embedded image patches with positional encoding, much like the ViT model discussed above. The output tokens are passed to the decoder stack of the transformer where the tokens are learnable parameters and part of the model rather than data. DETR calls these object queries. You can think of them as latent anchor proposals that condition the model to make bounding box predictions. The output tokens are decoded into bounding boxes and object scores using a multi-layer perceptron.

The model parameters (including the decoder object queries) are trained end-to-end using a bipartite matching loss to avoid duplicate detections. Here a bipartite graph is constructed between predicted and ground-truth detections. The best matching predictions are considered positives and the remaining detections treated as negatives.

The detection transformer (DETR) [12]
Figure 119: The detection transformer (DETR) [12].

8.5.3 Other Applications

The use of transformers in computer vision is a very active research area. Other applications include video understanding, image and video generation, and 3D modelling. Novel applications and release of new pre-trained models are being reported almost every day.

Beyond language (including coding) and vision, transformers have begun to see applications in robotics and planning where foundation models combine interpreting natural language prompts and visual input, and then produce output actions such as motor control. These tend to go under the umbrella name of vision-language-action (VLA) models.