Understanding the Vision Transformer and Counting Its Parameters
By: Andrei-Cristian Rad
Tags: Computer Vision, Deep Neural Networks, Transformers, Deep Learning

Photo by Simon Migaj on Unsplash

In this post, I will share my understanding of the Vision Transformer architecture, which I have researched for a report on Computer Vision for my Master’s degree. All the drawings in this post are original content, based on the knowledge from the paper and other tutorials which will be referred to where appropriate.

Transformer architectures have been a major breakthrough in Natural Language Processing (NLP) tasks since being proposed in 2016. Google’s BERT and Open AI’s GPT-2 / GPT-3 architectures have been the state-of-the-art solutions for various tasks, including language modeling, text summarization, and question answering.

The aim was to prove that the recurrent neural networks can be completely replaced, and solutions can be developed using only attention mechanisms — hence the pun in the Transformer paper title. Detailing the way the original Transformer architecture operates is beyond the scope of this article, but there are plenty of solid tutorials and videos that explain the way it works.

The focus of this article is presenting an overview of the Vision Transformer (ViT) architecture, proposed in a new paper for Google, submitted for review for ICLR 2021. Although there were previous attempts at using attention mechanisms for computer vision tasks (this, this, or this), ViT is the most promising architecture, both from the scalability and the efficiency point of view.

The Vision Transformer architecture, from the original paper

The architecture contains 3 main components.

Patch embedding. Feature extraction via stacked transformer encoders. The classification head.

Each component will be detailed in the next paragraphs, with the focus on the transformations and the trainable parameters. The first dimension of some tensors is b, which represents the batch size, but this will be ignored in the comments.

Patch Embedding

Patch Embedding

In the first step, an input image of shape (height, width, channels) is embedded into a feature vector of shape (n+1, d), following a sequence of transformations. This corresponds to equation (1) from the paper:

The image is split into n square patches of shape (p, p, c), where p is a pre-defined parameter, in raster order (left to right, up to down). The patches are flattened, resulting in n line vectors of shape (1, p²*c). The flattened patches are multiplied with a trainable embedding tensor of shape (p²*c, d), which learns to linearly project each flat patch to dimension d. This dimension d is constant in the architecture and is used in most of the components. The result is n embedded patches of shape (1, d). A learnable [cls] token of shape (1, d) is prepended to the sequence of patch embeddings. The idea of this token is from the BERT paper, where only the last feature representation corresponding to this token is fed through the classification layers. Intuitively, this represents an aggregate of the representations of the patches. A trainable positional embedding tensor, Eₚₒₛ, with the same shape, (n+1, d), is added to the concatenated sequence of projections. This tensor learns 1D positional information for each of the patches, in order to add a spatial representation of each patch within the sequence.

The result, z₀, is the first input to the stacked transformer encoders. The L stacked encoders represent the second component of the architecture. Each transformer takes as input features represented as an (n+1, d) tensor, and produces an output of the same dimension.

Transformer Encoder

Transformer Encoder

In the second step, the network learns more abstract features from the embedded patches, using a stack of L transformer encoders. This corresponds to equations (2) and (3) from the paper.

The encoder component contains a multi-headed attention (MHA) mechanism and a 2-layer MLP, with layer normalization and residual connections in between. The trainable weights in this component lie inside the MHA mechanism and the MLP weights. Since the MLP has 2 layers (hidden and output), there will be two weight matrices:

Wₕ — The weight tensor mapping the input layer to the hidden layer, of shape (d, dₘₗₚ)

Wₒ — The weight tensor mapping the hidden layer to the output layer, of shape (dₘₗₚ, d)

Multi-Head Attention Mechanism

The multi-head attention (MHA) step, included in each of the L stacked transformers, corresponds to equations (5), (6), (7), and (8) from the paper’s appendix.

The hidden state from the previous encoder is multiplied with the K trainable matrices of shape (d, 3dₕ), each corresponding to one of the K heads.

On each head, the result of this multiplication is an (n+1, 3dₕ) tensor. This one contains a concatenation of the keys (k), queries (q), and values (v), each of the dimensions (n+1, d). From the multi-head attention intuition, multiple heads allow the mechanism to learn from different aspects of the abstract representation.

Following, on each head, the scaled dot-product attention tensor (A) is computed as a softmax of the multiplication between the k and q matrices, normalized with the square root of the dimension of the head. The self-attention is the product between A and v, which has the shape (n+1, dₕ). The self-attention matrices are concatenated on the second dimension, resulting in an (n+1, d) tensor, which is then run through a single linear layer, by effectively multiplying it with a (d, d) trainable tensor.

Classification Head

Classification Head

As stated previously, only the last representation of the [cls] token is used in the classification head. For pre-training, a 2-layer MLP is used, therefore there are two weight matrices — Wₕ of shape (d, dₘₗₚ) and Wₒ of shape (dₘₗₚ, d). For fine-tuning, a single linear layer is used, therefore there is only a single tensor, of shape (d, n_cls). In each case, the final output of the network is a vector of shape (1, n_cls), containing the probabilities associated with each of the n_cls classes.

Trainable Parameters

Let us take the ViT-Base architecture and calculate the number of parameters. This number is equal to the sum of the number of parameters in each trainable tensor in the architecture.

In the patch embedding step, the two embedding matrices account for 786.432 parameters.

p²*c*d + (n+1)*d = 256*3*768 + 256*768 = 786.432

In the encoder stack, there are a total of 84.943.656 parameters.

L*(k*d*3*dₕ + d*d + d*dₘₗₚ + dₘₗₚ*d)

= 12*(12*768*3*64 + 768*768 + 2*768*3072)

= 12*(1.769.472 + 589.824 + 4.718.592)

= 12*589.824(3+1+8)

= 12*589.824*12

= 84.934.656

For fine-tuning (on ImageNet, for instance), there are other 768.000 parameters.

d*ncls = 768 * 1000 = 768.000

Adding those, we obtain around 86M, as listed in the paper.

Conclusions

For me, being able to visualize the transformations and the flow in this architecture has been of great help. Hopefully, it is the same in your case. Cheers!