DeepViT: Towards Deeper Vision Transformer

Introduction

Recent studies have demonstrated that transformers can be successfully applied to vision tasks with competitive performance compared with convolutional neural networks (CNNs) . Different from CNNs that aggregate global information by stacking multiple convolutions (e.g., $3\times 3$) , vision transformers (ViTs) take advantages of the self-attention (SA) mechanism to capture spatial patterns and non-local dependencies. This allows ViTs to aggregate rich global information without handcrafting layer-wise local feature extractions as CNNs and thus achieves better performance. For example, as shown in , a 12-block ViT model with 22M learnable parameters achieves better results than the ResNet-101 model which has more than 30 bottleneck convolutional blocks in ImageNet classification.

The recent progress of deep CNN models is largely driven by training very deep models with a large number of layers which is enabled by novel model architecture designs . This is because a deeper CNN can learn richer and more complex representations for the input images and provide better performance on vision tasks . Thus, how to effectively scale CNNs to be deeper is an important theme in recent deep learning fields, which stimulates the techniques like residual learning .

Considering the remarkable performance of shallow ViTs, a natural question arises: can we further improve performance of ViTs by making it deeper, just like CNNs? Though it seems to be straightforward at the first glance, the answer may not be trivial since ViT is essentially different from CNNs in its heavy reliance on the self-attention mechanism. To settle the question, we investigate in detail the scalability of ViTs along depth in this work.

We start with a pilot study on ImageNet to investigate how the performance of ViT changes with increased model depth. In Fig. 1, we show the performance of ViTs with different block numbers (green line), ranging from 12 to 32. As shown, as the number of transformer blocks increases, the model performance does not improve accordingly. To our surprise, the ViT model with 32 transformer blocks performs even worse than the one with 24 blocks. This means that directly stacking more transformer blocks as performed in CNNs is inefficient at enhancing ViT models. We then dig into the cause of this phenomenon. We empirically observed that as the depth of ViTs increases, the attention maps, used for aggregating the features for each transformer block, tend to be overly similar after certain layers, which makes the representations stop evolving after certain layers. We name this specific issue as attention collapse. This indicates that as the ViT goes deeper, the self-attention mechanism becomes less effective for generating diverse attentions to capture rich representations.

To resolve the attention collapse issue and effectively scale the vision transformer to be deeper, We present a simple yet effective self-attention mechanism, named as Re-attention. Our Re-attention takes advantage of the multi-head self-attention(MHSA) structure and regenerates attention maps by exchanging the information from different attention heads in a learnable manner. Experiments show that, Without any extra augmentation and regularization policies, simply replacing the MHSA module in ViTs with Re-attention allows us to train very deep vision transformers with even 32 transformer blocks with consistent improvements as shown in Fig. 1. In addition, we also provide ablation analysis to help better understand of the role of Re-attention in scaling vision transformers.

To sum up, our contributions are as follows:

We deeply study the behaviour of vision transformers and observe that they cannot continuously benefit from stacking more layers as CNNs. We further identify the underlying reasons behind such a counter-intuitive phenomenon and conclude it as attention collapse for the first time.

We present Re-attention, a simple yet effective attention mechanism that considers information exchange among different attention heads.

To the best of our knowledge, we are the first to successfully train a 32-block ViT on ImageNet-1k from scratch with consistent performance improvement. We show that by replacing the self-attention module with our Re-attention, new state-of-the-art results can be achieved on the ImageNet-1k dataset without any pre-training on larger datasets.

Related Work

Transformers are initially used for machine translation which replace the recurrence and convolutions entirely with self-attention mechanisms and achieve outstanding performance. Later, transformers become the dominant models for various natural language processing (NLP) tasks . Motivated by their success on the NLP tasks, recent researchers attempted to combine the self-attention mechanism into CNNs for computer vision tasks .. Those achievements also stimulate interests of the community in building purely transformer-based models (without convolutions and inductive bias) for vision tasks. The vision transformer (ViT) is among the first attempt that uses the pure transformer architecture to achieve competitive performance with CNNs on the image classification task. However, due to the large model complexity, ViT needs to be pre-trained on larger-scale datasets (e.g., JFT300M) for performing well on the ImageNet-1k dataset. To solve the data efficiency issue, DeiT deploys knowledge distillation to train the model with a larger pre-trained teacher model. In this manner, vision transformer can perform well on ImageNet-1k without the need of pre-training on larger dataset. Differently, in this work, we target at a different problem with ViT, i.e., how to effectively scale ViT to be deeper. We propose a new design for the self-attention mechanism so that it can perform well on vision tasks without the need of extra data, teacher networks, and the domain specific inductive bias.

2 Depth Scaling of CNNs

Increasing the network depth of a CNN model is deemed to be an effective way to improve the model performance . However, very deep CNNs are generally harder to train to perform significantly better than the shallow ones in the past . How to effectively scale up the CNNs in depth was a long-standing and challenging problem . The recent progress of CNNs largely benefits from novel architecture design strategies that make training deep CNNs more effective . Transformer-alike models have modularized architectures and thus can be easily made deeper by repeating the basic transformer blocks or using larger embedding dimensions . However, those straightforward scaling strategies only work well with larger datasets and stronger augmentation policies to alleviate the brought training difficulties. In this paper, we observed that with the same dataset, the performance of vision transformers do saturate as the network depth rises. We rethink the self-attention mechanism and present a simple but effective approach to address the difficulties in scaling vision transformers.

Revisiting Vision Transformer

A vision transformer (ViT) model , as depicted in Fig. 2(a), is composed of three main components: a linear layer for patch embedding (i.e., mapping the high-resolution input image to a low-resolution feature map), a stack of transformer blocks with multi-head self-attention and feed-forward layers for feature encoding, and a linear layer for classification score prediction. In this section, we first review its unique transformer blocks, in particular the self-attention mechanism, and then we provide studies on the collapse problem of self-attention.

Transformers were extensively used in natural language for encoding a sequence of input word tokens into a sequence of embeddings. To comply with such sequence-to-sequence learning structure when processing images, ViTs first divide an input image into multiple patches uniformly and encode each patch into a token embedding. Then, all these tokens, together with a class token, are fed into a stack of transformer blocks.

Each transformer block consists of a multi-head self-attention (MHSA) layer and a feed-forward multi-layer perceptron (MLP). The MHSA generates a trainable associate memory with a query ($Q$) and a pair of key ($K$)-value ($V$) pairs to an output via linearly transforming the input. Mathematically, the output of a MHSA is calculated by:

2 Attention Collapse

Motivated by the success of deep CNNs , we conduct systematic study in the changes of the performance of ViTs as depth increases. Without loss of generality, we first fix the hidden dimension and the number of heads to 384 and 12 respectivelySimilar phenomenon can also be found when we vary the hidden dimension size according to our experiments., following the common practice in . Then we stack different number of transformer blocks (varying from 12 to 32) to build multiple ViT models corresponding to different depths. The overall performances for image classification are evaluated on ImageNet and summarized in Fig. 1. As evidenced by the performance curve, we surprisingly find that the classification accuracy improves slowly and saturates fast as the model goes deeper. More specifically, we can observe that the improvement stops after employing 24 transformer blocks. This phenomenon demonstrates that existing ViTs have difficulty in gaining benefits from deeper architectures.

Such a problem is quite counter-intuitive and worth exploration, as similar issues (i.e., how to effectively train a deeper model) have also been observed for CNNs at its early development stage , but properly solved later . By taking a deeper look into the transfromer architecture, we would like to highlight that the self-attention mechanism plays a key role in ViTs, which makes it significantly different from CNNs. Therefore, we start with investigating how the self-attention, or more concretely, the generated attention map $\mathbf{A}$ varies as the model goes deeper.

To measure the evolution of the attention maps over layers, we compute the following cross-layer similarity between the attention maps from different layers:

where $M^{p,q}$ is the cosine similarity matrix between the attention map of layers $p$ and $q$. Each element $M^{p,q}_{h,t}$ measures the similarity of attention for head $h$ and token $t$. Consider one specific self-attention layer and its $h$-th head, $\mathbf{A}^{*}_{h,:,t}$ is a $T$-dimensional vector representing how much the input token $t$ contributes to each of the $T$ output tokens. $M^{p,q}_{h,t}$, thus, provides an appropriate measurement on how the contribution of one token varies from layer $p$ to $q$. When $M^{p,q}_{h,t}$ equals one, it means that token $t$ plays exactly the same role for self-attention in layers $p$ and $q$.

Given Eqn. (2), we then train a ViT model with 32 transformer blocks on ImageNet-1k and investigate the above similarity among all the attention maps. As shown in Fig. 3(a), the ratio of similar attention maps in $M$ after the 17th block is larger than 90% . This indicates that the learned attention maps afterwards are similar and the transformer block may degenerate to an MLP. As a result, further stacking such degenerated MHSA may introduce the model rank degeneration issue (i.e., the rank of the model parameter tensor from multiplying the layer-wise parameters together will decrease) and limits the model learning capacity. This is also validated by our analysis on the degeneration of learned features as shown below. Such observed attention collapse could be one of the reasons for the observed performance saturation of ViTs. To further validate the existence of this phenomenon for ViTs with different depths, we conduct the same experiments on ViTs with 12, 16, 24 and 32 transformer blocks respectively and calculate the number of blocks with similar attention maps. The results shown in Fig. 3(b) clearly demonstrate the ratio of the number of similar attention map blocks to the total number of blocks increases when adding more transformer blocks.

To understand how the attention collapse may hurt the ViT model performance, we further study how it affects feature learning of the deeper layers. For a specific 32-block ViT model, we compare the final output features with the outputs of each intermediate transformer block by investigating their cosine similarity. The results in Fig. 4 demonstrate that the similarity is quite high and the learned features stop evolving after the 20th block. There is a close correlation between the increase of attention similarity and feature similarity. This observation indicates that attention collapse is responsible for the non-scalable issue of ViTs.

Re-attention for Deep ViT

As revealed above, one major obstacle in scaling up ViT to a deeper one is the attention collapse problem. In this section, we present two solution approaches, one is to increase the hidden dimension for computing self-attention and the other one is a novel re-attention mechanism.

One intuitive solution to conquer attention collapse is to increase the embedding dimension of each token. This will augment the representation capability of each token embedding to encode more information. As such, the resultant attention maps can be more diverse and the similarity between each block’s attention map could be reduced. Without loss of generality, we verify this approach empirically by conducting a set of experiments based on ViT models with 12 blocks for quick experiments. Following previous transformer based works , four embedding dimensions are selected, ranging from 256 to 768. The detailed configurations and the results are shown in Tab. 1.

From Fig. 5 and Tab. 1, one can see that the number of blocks with similar attention maps is reduced and the attention collapse is alleviated by increasing the embedding dimension. Consequently, the model performance is also increased accordingly. This validates our core hypothesis—the attention collapse is the main bottleneck for scaling ViT. Despite its effectiveness, increasing the embedding dimension also increases the computation cost significantly and the brought performance improvement tends to diminish. Besides, a larger model (with higher embedding dimension) typically needs more data for training, suffering the over-fitting risk and decreased efficiency.

2 Re-attention

It has been demonstrated in Sec. 3 that the similarity between attention maps across different transformer blocks is high, especially for deep layers. However, we find the similarity of attention maps from different heads of the same transformer block is quite small, as shown in Fig. 3(c). Clearly, different heads from the same self-attention layer focus on different aspects of the input tokens. Based on this observation, we propose to establish cross-head communication to re-generate the attention maps and train deep ViTs to perform better.

Advantages: The advantages of the proposed Re-attention are two-fold. First of all, compared with other possible attention augmentation methods, such as randomly dropping some elements of the attention map or tuning SoftMax temperature, our Re-attention exploits the interactions among different attention heads to collect their complementary information and better improves the attention map diversity. This is also verified by our following experiments. Furthermore, our Re-attention is effective and easy to implement. It needs only a few lines of code and negligible computational overhead compared to the original self-attention. Thus it is much more efficient than the approach of increasing embedding dimension.

Experiments

In this section, we first conduct experiments to further demonstrate the attention collapse problem. Then, we give extensive ablation analysis to show the advantages of the proposed Re-attention. By incorporating Re-attention into the transformers, we design two modified version of vision transformers and name them as deep vision transformers (DeepViT). Finally, we compare the proposed DeepViT models against the latest state-of-the-arts (SOTA).

To make a fair comparison, we first tuned a set of parameters for training the ViT base model and then use the same set of hyper-parameters for all the ablation experiments. Specifically, we use AdamW optimizer and cosine learning rate decay policy with an initial learning rate of 0.0005. We use 8 Telsa-V100 GPUs and train the model for 300 epochs using Pytorch library. The batch size is set to 256. We use 3 epochs for learning rate warm-up . We also use some augmentation techniques such as mixup and random augmentation to boost the performance of baseline models following . When comparing with other methods, we adopt the same set of hyper-parameters as used by the target models. We report results on the ImageNet dataset . For all experiments, the image size is set to be 224$\times$224. To study the scaling capability of current transformer blocks, we set the embedding dimension to 384 and the expansion ratio 3 for the MLP layers. We use 12 heads for all the models. More detailed configurations are shown in Tab. 2.

2 More Analysis on Attention Collapse

In this section, we show more analysis on the attention map similarity and study how the collapsed attention maps affect the model performance.

Attention reuse: As discussed above, when the model goes deeper, the attention maps of the deeper blocks become highly similar. This implies that adding more blocks on a deep ViT model may not improve the model performance. To further verify this claim, we design an experiment to reuse the attention maps computed at an early block of ViT to replace the ones after it. Specifically, we run experiments on the ViT models with 24 blocks and 32 blocks but share the $Q$ and $K$ values (and the resulted attention maps) of the last “unique” block to all the blocks afterwards. The “unique” block is defined as the block whose attention map’s similarity ratio with adjacent layers is smaller than 90%. More implementation details can be found in the supplementary material. The results are shown in Tab. 3. Surprisingly, for a ViT model with 32 transformer blocks, when we use the same $Q$ and $K$ values for the last 15 blocks, the performance degradation is negligible. This implies the attention collapse problem indeed exists and reveals the inefficacy in adding more blocks when the model is deep.

Visualization: To more intuitively understand the attention map collapse across layers, we visualize the learned attention maps from different blocks of the original ViT . We take a 32-block ViT model as an example and pre-train it on ImageNet. The visualization of the attention maps with original MHSA and Re-attention are shown in Fig. 6. It can be observed that the original MHSA learns the local relationship among the adjacent patches in the shallow blocks and the attention maps tend to expand to cover more patches gradually. In the deep blocks, the MHSA learns nearly uniform global attention maps with high similarity. Differently, after implementing Re-attention, the attention maps at deep blocks keep the diversity and have small similarities from adjacent blocks.

3 Analysis on Re-attention

In this subsection, we present two straightforward modifications to the current self-attention mechanism as baselines. We then conduct a series of comparison experiments to show the advantages of our proposed Re-attention.

Re-attention v.s. Self-attention: We first evaluate the effectiveness of Re-attention by comparing to the pure ViT models using the same set of training hyper-parameters. We directly replace the self-attention module in ViT with Re-attention and show the results in Tab. 4 with different number of transformer blocks. As can be seen, the vanilla ViT architecture suffers performance saturation when adding more transformer blocks. This phenomenon coincides with our observations that the number of blocks with similar attention maps increases with the depth as shown in Fig. 3(b). Interestingly, when replacing the self-attention with our proposed Re-attention, the number of similar blocks are all reduced to be zero and the performance rises consistently as the model depth increases. The performance gain is especially significant for deep ViT with 32 blocks. This might be explained by the fact that the 32 block ViT model has the largest number of blocks with similar attention maps and the improvements should be proportional to the number similar blocks in the model. These experiments demonstrate that the proposed Re-attention can indeed solve the attention collapse problem and thus enables training a very deep vision transformer without extra datasets or augmentation policies.

Comparison to adding temperature in self-attention: The most intuitive way to mitigate the over-smoothing phenomenon is to sharpen the distribution of the elements in the attention map of MHSA. We could achieve this by assigning a temperature $\tau$ to the Softmax layer of MHSA:

As the attention collapse is observed to be severe on deep layers (as shown in Fig. 3), we design two sets of experiments on a ViT model with 32 transformer blocks: (a) linearly decaying the temperature $\tau$ in each block such that the attention map distribution is sharpened and (b) making the temperature learnable and optimized together with the model training. We first check the impact of the SoftMax temperature on reducing the attention map similarity. As shown in Fig. 8(a), the number of similar blocks are still large. Correspondingly, the feature similarity among blocks are also large as shown in Fig. 8(b). Thus, adding a temperature to the SoftMax only reduces the attention map similarity by a small margin. The classification results on ImageNet are shown in Tab. 5. As shown, using a learnable temperature could improve the performance but the improvement is marginal.

Another baseline we have attempted to differentiate the self-attention maps across layers is to use random dropout on the attention maps $\mathbf{A}$. As the dropout will mask out different positions on the attention maps for different blocks, the similarity between attention maps could be reduced. The impacts on the attention maps and the output features of each block are shown in Fig. 8(a-b). It is observed that dropping attention does reduce the cross layer similarity of the attention maps. However, the similarity among features are not reduced by much. This is because the difference between attention maps comes from the zero positions in the generated mask. Those zero values do reduce the similarity between attention maps but not contribute to the features. Thus, the improvement is still not significant as shown in Tab. 5.

Advantages of Re-attention: Our proposed Re-attention brings more significant improvements over the temperature-tuning and the attention-dropping methods. This is because both adding temperature and dropping attention are regularizing the distribution of the originally over-smoothed self-attention maps, without explicitly encouraging them to be diverse. However, our proposed Re-attention mechanism uses different heads (whose attention maps are dissimilar) as basis and re-generate the attention maps via the transformation matrix $\Theta$. This process incorporates the inter-head information communication and the generated attention maps can encode richer information. It is worth noting that the original MHSA design can be thought as a special case of Re-attention with an identity transformation matrix. By making $\Theta$ learnable for each block, an optimized pattern could be learned end to end. As shown in Fig. 8(c), the learned transformation matrix assigns a diverse set of weights for each newly generated head. It clearly shows that the combination for each new heads takes different weights from the original heads in the re-attention process and thus reduces the similarity between their attention maps. As shown in Fig. 8(a), our proposed Re-attention achieves the lowest cross layer attention map similarity. Consequently, it also reduces the feature map similarity across layers as shown in Fig. 8(b).

4 Comparison with other SOTA models

With Re-attention, we design two ViT variants, i.e., DeepViT-S and DeepViT-L, based on the ViT with 16 and 32 transformer blocks respectively. For both models, we use Re-attention to replace the self-attention. To have a similar number of parameters with other ViT models, we adjust the embedding dimension accordingly. The hidden dimensions of DeepViT-S and DeepViT-L models are set as 396 and 408 respectively. More details on the model configuration are given in the supplementary material. Besides, motivated by , we add three CNN layers for extracting the token embeddings, using the same configurations as . It is worth noting that we do not use the optimized training recipes and the repeated augmentation as for training our models. The results are shown in Tab. 6. Clearly, our DeepViT model achieves higher accuracy with less parameters than the recent CNN and ViT based models. Notably, without any complicated architecture change as made by T2T-ViT (adopting a deep-narrow architecture) or DeiT (introducing token distillation), simply using the Re-attention makes our DeepViT-L outperforms them by 0.4 points with even smaller model size (55M vs. 64M & 86 M).

Conclusion

In this work, we found the attention collapse problem of vision transformers as they go deeper and propose a novel Re-attention mechanism to solve it with minimum amount of computation and memory overhead. With our proposed Re-attention, we are able to maintain an increasing performance when increasing the depth of ViTs. We hope our observations and methods could facilitate the development of vision transformers in future.

References

Appendix A Experiment Implementation Details

As shown in Fig. 3(b) and Tab. 3 in the main paper, the vision transformers with 24 blocks and 32 blocks have 11 and 15 blocks with similar attention maps, respectively. To verify the effectiveness of the attention maps from those blocks, we directly force those blocks to share the same attention map as the last ‘unique’ block as defined in Sec. 5.2. Specifically, we take the attention map of the last ‘unique’ block and denote it as $\mathbf{A}_{unique}$. For all the following blocks, the attention output is calculated by:

where $\Theta$ is used to simulate the small variance between attention maps across layers since they are not identical. Norm is batch normalization used to adjust the variance across layers. As shown in Tab. 3, for a ViT with 32 blocks, forcing the top 15 blocks to share the same attention map causes negligible degradation on the classification accuracy on ImageNet. This proves that adding those blocks do not contribute to the accuracy improvement.

Training loss:

We use the cross-entropy (CE) loss as the training loss for all experiments. To minimize the similarity of the attention maps across layers, we add the cosine similarity between layers into the loss function when training the model.

where $\text{Similarity}(\mathbf{A}^{l},\mathbf{A}^{l+1})$ denotes the cosine similarity between layer $l$ and $l+1$ and $\mathbf{A}^{l}$ denotes the attention map of layer $l$. B denotes the number of bottom blocks used for regularization and is a hyper-parameter. We set B to 4, 8 and 12 for training ViT models with 16, 24 and 32 blocks respectively.

Appendix B DeepViT architecture design

As observed in Fig. 3(a), the attention maps of the transformer blocks become similar only at the top blocks. Thus, it is not necessary to apply re-attention to all blocks. To study the optimal number of blocks with re-attention, we conduct a set of experiments on a ViT model with 16 transformer blocks. For each experiment, we only apply re-attention on the top $K$ blocks where $K$ ranges from 5 to 15. The rest of the blocks are using the original transformer block structure. We train each model on ImageNet with the same set of training hyper-parameters as those for baseline models as detailed in Sec. 5 in the main paper. The results are shown in Fig. 9.

It is observed that, as the number of re-attention blocks varies, the top-1 classification accuracy changes correspondingly. The highest accuracy appears at the position where the number of re-attention blocks is the same as the number of similar attention map blocks. Based on this observation, we define the architecture of DeepViT-S and DeepViT-L with 5 and 12 re-attention blocks respectively. Detailed configurations are shown in Tab. 7. Note that we adjust the embedding dimension to have a comparable size with other methods.

Appendix C Impacts of hyper-parameters

In the main paper, all experiments are run with the same set of training hyper-parameters as the one used for reproducing ViT models. However, as shown in , an improved training recipe could improve the performance of ViT models significantly. In Tab. 8, we present the performance of DeepViT-S and DeepViT-L with the same set of training recipes as DeiT except that we do not use repeated augmentation. In Tab. 8, it is clearly shown that the performance of DeepViT could be further improved with optimized training hyper-parameters.

Appendix D Similarity calculation

Cosine similarity between layers To measure the similarity between the attention maps, we define the similarity $S^{p,q}$ between the attention maps of two layers, $p$ and $q$, as the ratio of the number of similar vector pairs to the total number of pairs between two attention maps:

where $\tau$ is a hyper-parameter and used as a threshold for deciding similar vectors0.5 is selected as a threshold for visualization purpose in this paper.

Definition of similar blocks A block is counted as a similar block if the similarity between it’s attention map and the adjacent block’s attention map is larger than 80%. To measure the block similarity for a ViT model with $B$ blocks, we take the ratio of the number of similar blocks to the total number of blocks as a measurement.