Vision Transformers: State of the Art and Research Challenges

In order to learn variant representations at different positions, the input $X$ is transformed into $n_{h}$ different representations (heads), denoted by $h_{1},\cdots,h_{i},\cdots,h_{n_{h}}$. The attention is first computed for each head $h_{i}$ with $Q^{i}$, $K^{i}$, and $V^{i}$ with projection matrices $W_{q}^{i}$, $W_{k}^{i}$, and $W_{v}^{i}$, respectively.

$$h_{i}=\textit{Attention}(Q^{i},K^{i},V^{i}).$$

(3)

Afterward, all the heads are concatenated to form the multi-head representations $H$ as follows:

$$H=\textit{Concat}(h_{1},h_{2},...,h_{n_{h}}).$$

(4)

Finally, $H$ is transformed back to dimension $d$ with projection matrix $W^{O}$, i.e., $O=HW^{O}$. As the number of relations between inputs is usually unknown, multi-head attention is commonly-used to capture different relations between input elements in a data-driven manner.

Feed-Forward Networks (FFN) are mainly composed of Fully-Connected layers and can be written in:

$$\textit{FFN}(X)=\delta(XW_{1}+b_{1})W_{2}+b_{2},$$

(5)

where $W_{1}\in\mathbb{R}^{d\times h}$ and $W_{2}\in\mathbb{R}^{h\times d}$ are the weight and $b_{1}$ and $b_{2}$ are the bias terms. $\delta$ is ReLU activation function. $h$ is the hidden dimension and is usually set to $4d$.

Since the self-attention relaxes the sequential order, transformers use positional encoding to maintain the ordinal information. The author proposes to use sine and cosine functions with different frequencies to derive the positional encoding, which allows the model to learn the relative positions since any positional encoding can be linearly combined by other positional encodings. Specifically, the $i$-th dimension of the positional encoding at position $pos$, denoted by $\textit{PE}_{(pos,i)}$ can be calculated as follows:

$$\textit{PE}_{(pos,i)}=\begin{cases}\sin(pos/10000^{i/d})&\text{if }i\text{ is even,}\\ \cos(pos/10000^{(i-1)/d})&\text{otherwise,}\end{cases}$$

(6)

where $d$ is the dimension of the input features.

A single encoder block has two sub-layers, i.e., multi-head attention and position-wise feed-forward networks. Each sub-layer is followed by residual connection He et al. (2016) and layer normalization Ba et al. (2016). The outputs of the encoder will be sent into the decoder as $K,V$ in the second multi-head attention layer.

Each decoder block is made by three sub-layers. The first two layers are multi-head attention blocks, and the final one is position-wise feed-forward networks. Similar to encoder block, all three layers are followed by residual connection and layer normalization as well. Since the decoder is auto-regressive, i.e., sequentially predicting a new result only based on previous predictions, the multi-head attention layers in the decoder utilize the mask operation to only attend the predicted results for preventing the violation of causality.

In order to transform an image $X\in\mathbb{R}^{C\times H\times W}$ into a 1D sequence of vectors, an image will be transformed into $X_{p}\in\mathbb{R}^{N\times(P^{2}C)}$ as illustrated in Fig. 1(b), where $P\times P$ is the resolution of each patch, and the number of patches is $N=HW/P^{2}$. The patches are then projected into patch embeddings by a linear layer. Akin to BERT Devlin et al. (2018), a learnable class token $x_{cls}\in\mathbb{R}^{1\times(P^{2}C)}$ is added to $X_{p}$. Afterward, the positional embeddings are added to the output $[x_{cls};X_{p}]$ to obtain the positional information.

Instead of using both encoder and decoder, ViT only utilizes the Transformer encoder since the goal is to find a better representation rather than autoregressive prediction. The image patches are transformed into sequences and then sent into the encoder. The only difference is that the Layer Normalization is added before the sub-layers (pre-norm) Xiong et al. (2020).

The class token is trained to include the class information and can be used to represent the entire features. Therefore, it can be used to classify the data into different categories.

MLP head is commonly-used for the downstream tasks. Specifically, let $Z=[Z_{0},Z_{1},\cdots,Z_{N}]$ denotes the output of the Transformer encoder. As the class token prepended to the input sequence can be regarded as the image representation, the final result $y$ is determined based on the first token of the output $Z_{0}$ from encoders through a single MLP layer, i.e.,

	$\displaystyle Z$	$\displaystyle=\text{TransEncoder}([x_{cls};X_{p}]),$		(7)
	$\displaystyle y$	$\displaystyle=\text{MLP}(Z_{0}).$		(8)

Touvron et al. (2021a) uses a CNN as a teacher model to train a vision transformer, which utilizes knowledge distillation Hinton et al. (2015) to transfer the inductive bias to a vision transformer and applies stronger data augmentation for input data. This approach allows others to train a vision transformer from scratch without the requirement of pre-training on a large dataset.

d’Ascoli et al. (2021) is similar to DeiT as ConViT also combines the inductive bias of CNN into models. Instead of using knowledge distillation, ConViT includes Gated Positional Self-attention (GPSA), which can be initialized as a convolutional layer Cordonnier et al. (2020) for capturing the local information at the beginning of the training stage. As such, ConViT can utilize the advantages of soft inductive bias of CNN without being limited to CNN. In other words, GPSA allows vision transformers to be the same as CNN to improve the data efficiency on small datasets but can be better than CNN when there are infinite data used in training.

Graham et al. (2021) gets the embeddings from four convolutional layers, which the model can extract the local features at the beginning and also reduce the input size. On top of that, LeViT further lowers the input size in some attention blocks, which speeds up the inference time. These advantages help strike the balance between accuracy, data efficiency and the training speed.

Yuan et al. (2021a) obtains the embedding features directly from convolutional blocks to add the locality in the beginning of the model, similar to LeViT. Moreover, the authors include a depth-wise convolutional layer in FFN to encourage the model to extract the local features. In order to exchange the class information in different layers, CeiT utilizes Layer-wise Class-token Attention to collect different class representations by computing the self-attention on class tokens.

Yuan et al. (2021a) also obtains the embeddings with convolutional layers. Besides, CvT uses convolutional layers to create query, key, and value before self-attention operation, which provides the local spatial information. Furthermore, CvT also shrinks the size of the key and the value by using $stride=2$ to accelerate the computation.

Li et al. (2021c) is designed to use convolutional layers in FFN to extract the local features in every transformer blocks. The authors also try to apply different activation functions (ReLU6, h-swish) and architectures (SE-Block, ECA module) in FFN layers to improve the performance.

Hassani et al. (2021) is proposed to solve the data-hungry problem to make vision transformers perform well on small datasets. Specifically, CCT gets the embeddings with convolutional layers and makes the input shape flexible by removing the positional embeddings. On top of that, CCT incorporates the Sequence Pooling at the end of the transformer layers to compute the weights of the output sequence.

Zhou et al. (2021) is created after the authors examine the attention maps of ViT and find out that the attention collapsing takes place in the deeper layer. This problem hinders the model to be representative and would lower the performance. The authors alter the self-attention layers and provide a learnable transformation matrix after the attention layer to address the issue by stimulating the model to generate a new set of attention maps.

Yuan et al. (2021b) is designed after the authors observe the feature maps reshaped from the tokens and point out that most of the token maps are meaningless in ViT. To improve the diversity of the token maps, the authors inserts a T2T-module in the beginning of the model. T2T-module is made up of few T2T-Transformers. These T2T-Transformers can be viewed as the original transformer layers or can be replaced by the Performer used in Choromanski et al. (2021). Between the T2T-Transformers is the T2T-process, which crops the images into several overlapping patches. The overlapping strategy enables the model to share the information between the neighbors to improve the features diversity.

Wang et al. (2021) is mainly designed for solving dense prediction problems (e.g., object detection and semantic segmentation). It uses Spatial Reduction Layer to reduce the computation by decreasing the dimension of $K$ and $V$. The feature size is reduced by a patch embedding layer at the beginning of each stage. Due to its pyramid structure, the model can generate multi-scale feature maps and can be trained faster with its smaller feature size.

Heo et al. (2021) includes Pooling Layer, which uses a depth-wise convolutional layer to achieve the dimension reduction. The authors also test PiT on different robustness benchmarks and get outstanding performance.

Liu et al. (2021) is proposed to derive a general-purpose backbone based on ViT, which can be used for different applications, e.g., image classification and semantic segmentation. Since applying self-attention pixel-by-pixel results in a tremendous computation complexity, Swin-Transformer forms a hierarchical structure, which merges patches after each Swin-Transformer block and therefore has approximate linear computation time complexity to input image size due to the computation of self-attention only within each local shifted window.

Chu et al. (2021) also computes the attention within a shifted window, while the global attention is evaluated after the local window attention. This helps each window retains the outside-window information, similar to the overlapping strategy. Likewise, the patches are merged at the beginning of the stages to form the hierarchical shape.

Zhang et al. (2021) is created on a different strategy to perform the hierarchical computation. Specifically, an image is first partitioned into $n$ blocks and every four blocks are merged after a transformer layer. Moreover, Gradient-based Class-aware Tree-traversal is proposed to visualize the most critical path from child to root, which reveals how the model makes the decision given an input image. It is worth noting that NexT can be used to generate images by turning the tree upside down due to its tree structure.

Chen et al. (2021a) can extract multi-scale features from two branches. These branches are L-Branch and S-Branch, where the former uses a larger patch size and the latter uses a smaller patch size. By using the dual-branch structure, the model is able to obtain different scales of spatial information. To fuse the spatial information between the branches, Cross Attention is applied by inserting the class token from one to the other.

Mao et al. (2021) is designed after studying different components in ViT with regards to the robustness, e.g., accuracy under adversarial attacks. Based on the results, RVT decides to 1) remove the class token, which is not important to a vision transformer, by averaging the features and 2) add CNN to the embedding and the FFN layers for increasing locality and 3) use more attention heads for obtaining different features. Moreover, the original self-attention is replaced by Position-Aware Attention Scaling (PAAS), which adds a learnable matrix to the self-attention for showing the importance of each $Q$-$K$ pair. PAAS is proved to suppress the unrelated signal in the noise input. Moreover, Patch-wise Augmentation applies different augmentations on different patches to diversify the training data.

Touvron et al. (2021b) employs the normalization with LayerScale, which uses learnable factors to adaptively normalize the features. The normalization speeds up the converge rate and allows the deep model to be trained well. CaiT also includes Class Attention Layer for computing the attention between the class embedding and the overall features to have better knowledge on the inputs.

Ali et al. (2021) is proposed to reduce the time complexity of the self-attention. The reduction is done by Cross-Covariance Attention, which uses the transposed version of self-attention, i.e., the self-attention is computed on the feature channels but not on the tokens. As such, the time complexity is reduced from $O(N^{2}d)$ to $O(N^{2}d/h)$. Moreover, Local Patch Interaction Block is employed to use depth-wise convolutional networks to further extract the information between different patches.