Augmented Shortcuts for Vision Transformers
Year: 2021

Abstract:
Transformer models have achieved great progress on computer vision tasks recently. The rapid development of vision transformers is mainly contributed by their high representation ability for extracting informative features from input images. However, the mainstream transformer models are designed with deep architectures, and the feature diversity will be continuously reduced as the depth increases, i.e., feature collapse. In this paper, we theoretically analyze the feature collapse phenomenon and study the relationship between shortcuts and feature diversity in these transformer models. Then, we present an augmented shortcut scheme, which inserts additional paths with learnable parameters in parallel on the original shortcuts. To save the computational costs, we further explore an efficient approach that uses the block-circulant projection to implement augmented shortcuts. Extensive experiments conducted on benchmark datasets demonstrate the effectiveness of the proposed method, which brings about 1% accuracy increase of the state-of-the-art visual transformers without obviously increasing their parameters and FLOPs.

heading: 1 Introduction
text: Originating from the natural language processing field, transformer models [35, 9] have recently made great progress in various computer vision tasks such as image classification [9, 34, 23], object detection [3] and image processing [4]. Wherein, the ViT model [9] divides the input images as visual sequences and obtains an 88.36% top-1 accuracy which is competitive to the SOTA convolutional neural network (CNN) models (e.g., EfficientNet [31]). Compared to CNNs which are usually customized for vision tasks with prior knowledge (e.g., translation equivalence and locality), vision transformers introduce less inductive bias and have a larger potential to achieve better performance on different visual tasks. Besides, considering the high performance of transformers in various fields (e.g., natural language processing and computer vision), we may only need to support transformer for processing different tasks which can significantly simplify the software/hardware design.
Besides the self-attention layers in vision transformers, a shortcut is often included to directly connect multiple layers with identity projection [15, 35, 9]. The introduce of shortcut is motivated by the architecture designs in CNNs, and has been demonstrated to be beneficial for a stable convergence and better performance [16, 33, 22, 39, 38]. There is a series of works on interpreting and understanding the role of shortcut. For example, Balduzzi et al. [1] analyze the gradients of deep networks and show that shortcut connection can effectively alleviate the problem of gradient vanishing and exploding. Veit et al. [36] reckon ResNet [15] as the ensemble of a collection of paths with different lengths, and the gradient vanishing problem is addressed by short paths. From a theoretical perspective, Liu et al. [22] prove that shortcut connection can avoid network parameters trapped by spurious local optimum and help them converge to a global optimum. Apart from CNNs, shortcut connection is
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
consequently widely-used in other deep neural networks such as transformer [35, 9], LSTM [11] and AlphaGo Zero [29].
In vision transformers, the shortcut connection bypasses the multihead self-attention (MSA) and multilayer perceptron (MLP) modules, which plays a critical role towards the success of vision transformers. A transformer without shortcut suffer extremely low performance (Table 1). Empirically, removing the shortcut results in features from different patches becoming indistinguishable as the network going deeper (shown in Figure 3(a)), and such features have limited representation capacity for the downstream prediction. We name this phenomenon as feature collapse. Fortunately, adding shortcut in transformers can alleviate the phenomenon (Figure 3(b)) and make the generated features diverse. However, the conventional shortcut simply copies the input feature to the output, limiting its ability for enhancing the feature’s diversity.
In this paper, we introduce a novel augmented shortcut scheme for improving feature diversity in vision transformers (Figure 1). Besides the conventional identity shortcut, we propose to parallel the MSA module with multiple parameterized shortcuts, which provide more alternative paths to bypass the attention mechanism. In particular, an augmented shortcut connection is constructed as the sequence of a linear projection with learnable parameters and a nonlinear activation function. We theoretically demonstrate that the transformer model equipped with augmented shortcuts can avoid feature collapse and produce more diverse features. For the efficiency reason, we further replace the original dense matrices with block-circulant matrices, which have lower computational complexity in the Fourier frequency domain while high representation ability in the spatial domain. Owing to the compactness of circulant projection, the parameters and computational costs introduced by the augmented shortcuts are negligible compared to those of the MSA and MLP modules. We empirically evaluate the effectiveness of augmented shortcut with the ViT model [9, 34] and its SOTA variants (PVT [37], T2T [40]) on ImageNet dataset. Equipped with our augmented shortcut, the performances (top-1 accuracies) of these models can be enhanced by about 1% with comparable computational costs (e.g., Figure 2).

heading: 2 Preliminaries and Motivation
text: The vision transformer (ViT) [9] adopts the typical architecture proposed in [35], with MSA and MLP modules alternatively stacked. For an image, ViT splits it into N patches and projects each patch into a d-dimension vector. Thus the input of a vision transformer can be seen as a matrix Z0 ∈ RN×d. An MSA module with H heads is defined as
MSA(Zl) = Concat([AlhZlW v lh] H h=1)W o l , (1)
where Zl ∈ RN×d is the features of the l-th layer, Alh ∈ RN×N and W vlh ∈ Rd×(d/H) are the corresponding attention map and value projection matrix in the h-th head, respectively. Concat(·)
denotes the concatenating for features of theH heads and W ol ∈ Rd×d is the output projection matrix. The attention matrix Alh is calculated by the self-attention mechanism, i.e.,
Alh = softmax
( (ZlW q lh)(ZlW k lh) >
√ d
) , (2)
Alh = softmax
( (ZlW q lh)(ZlW k lh) >
√ d
) , (3)
where W qlh ∈ Rd×(d/H) and W klh ∈ Rd×(d/H) are the query and value projection matrices, respectively. Attention Alh reflects the relation between different patches, and a larger value A ij lh indicate that patch i and patch j have a stronger relationship.
The MLP module extracts features from each patch independently, which is usually constructed by stacking two linear projection layers with weights W a ∈ Rd×dhidden , W b ∈ Rdhidden×d and non-linear activation function σ, i.e., MLP(Zl) = σ(ZlW a)W b. A vision transformer model is constructed by stacking the MLP and MSA modules alternatively.
Layer 12
(c) Aug-ViT
Layer 6
Layer 1
(a)ViT w/o Shortcut (b) ViT with Shortcut
Figure 3: The similarity matrices over patches in the ViTBase model. 2 4 6 8 10 12 Layer
0.0
0.2
0.4
0.6
0.8
r( Z)
Aug-ViT ViT w shortcut ViT w/o shortcut
Figure 4: Diversity metric r(Zl) varies w.r.t. layers.
Besides MSA and MLP, the shortcut connection also plays a vital role to achieve high performance. Empirically, removing shortcut severely harms the accuracies of ViT models (as shown in Table 1). To explore the reason behind, we propose to analyze the intermediate features in the ViT-like models. Specially, we calculate the cosine similarity between different patch features and show the similarity matrices of low (Layer 1), middle (Layer 6) and deep layers (Layer 12) in Figure 3. Features from different patches in a layer quickly become indistinguishable as the network depth increasing. We call this phenomenon feature collapse, which greatly restrict the representation capacity and then prevent high performance.
The feature collapse mainly results from the attention mechanism, which aggregates features from different patches layer-by-layer. Denoting Z̃lh as the output feature after the attention map, the attention mechanism can be formulated as:
Z̃ilh = N∑ j=1 AijlhZ j l , s.t. N∑ j=1 Aijlh = 1, i = [1, 2, · · · , N ], (4)
where the feature Z̃ilh from the i-th patch is the weighted average of features from all other patches, and the weights are values in attention map Alh. Though attention mechanisms capture global relationship, the diversity of patch features are reduced as well. Theoretically, the diversity of a feature Zl can be measured by the difference between the feature and a rank-1 matrix, i.e.,
r(Zl) = ‖Zl − 1z>l ‖, where zl = argminz′l‖Zl − 1z ′ l >‖. (5)
where ‖ · ‖ denotes the matrix norm1. zl, z′l ∈ Rd are vectors and 1 is an all-ones vector. The rank of matrix 1z>l is 1. Intuitively, if zl ∈ Rd can represent Zl ∈ RN×d of N patches, the generated
1Here ‖ · ‖ is defined as the `1, `∞-composite norm for the convenience of theoretical derivation.
feature will be redundant. A larger r(Zl) implies stronger diversity of the given feature Zl. We use r(Zl) as the diversity metric in the following. Because MSA module incurs the feature collapse, we focus on the a model stacked by attention modules and have the following theorem [8], i.e.,
Theorem 1 Given a model stacked by the MSA modules, the diversity r(Xl) of feature in the l-th layer can be bounded by that of input data Z0, i.e.,
r(Zl) ≤ ( Hγ√ d ) 3l−1 2 r(Z0) 3l , (6)
where H is number of heads, d is feature dimension and γ is a constant related to the norms of weight matrices in the MSA module. Hγ/ √ d and r(Z0) are usually smaller than 1, so the feature diversity r(Zl) will decrease rapidly as the network depth increases [8]. We also empirically show how r(Zl) varies w.r.t. the network depth in the ViT models (Figure 4), and the empirical results are accordant to the Theorem 1.
Fortunately, adding a shortcut connection parallel to the MSA module can empirically preserve the feature diversity especially in deep layers (see Figure 3(b) and Figure 4). The MSA module with shortcut can be formulated as:
ShortcutMSA(Zl) = MSA(Zl) + Zl, (7)
where the identity projection (i.e., Zl) is parallel to the MSA module. Intuitively, the shortcut connection bypasses the MSA module and provides another alternative path, where features can be directly delivered to the next layer without interference of other patches. Adding the shortcut connections can also theoretically improve the bound of feature diversity r(Zl) (as discussed in Section 3.1). The success of shortcut shows that bypassing the attention layers with extra paths is an effective way to enhance feature diversity and improve the performance of transformer-like models.
However, in general vision transformer models [9, 13, 37, 14, 12, 40, 23], there is only a single shortcut connection with identity projection for each MSA module, which only copies the input features to the outputs. This simple formulation may not have enough representation capacity to improve the feature diversity maximally. In the following chapters, we aim to refine the existing shortcut connections in vision transformers and explore efficient but powerful augmented shortcuts to produce visual features with higher diversity.

heading: 3 Approach

heading: 3.1 Augmented Shortcuts
text: We propose augmented shortcuts to alleviate the feature collapse by paralleling the original identity shortcut with more parameterized projections. The MSA module equipped with T augmented shortcuts can be formulated as:
AugMSA(Zl) = MSA(Zl) + Zl + T∑ i=1 Tli(Zl; Θli), l ∈ [1, 2, · · · , L], (8)
where Tli(·) is the i-th augmented shortcut connection of the l-th layer and Θli denotes its parameters. Besides the original shortcut, the augmented shortcuts provide more alternative paths to bypass the attention mechanism. Different from the identity projection directly copying the input patches to the corresponding outputs, the parameterized projection Tli(·) can transform input features into another feature space. Actually, projections Tli(·) will make different transformations on the input feature as long as their weight matrices Θli are different, and thus paralleling more augmented shortcuts has potential to enrich the feature space.
A simple formulation for Tli(·) is the sequence of a linear projection and an activation function i.e.,
Tli(Zl; Θli) = σ(ZlΘli), l ∈ [1, · · · , L], i ∈ [1, 2, · · · , T ], (9)
where Θli ∈ Rd×d is the weight matrix and σ is the non-linear activation function (e.g., GELU). In Eq. 9, Tli(·) tackles each patch independently and preserves their specificity, which is complement to
the MSA modules aggregating different patches. Note that the identity mapping is a special case of Eq. 9, i.e., σ(x) = x and Θli is the identity matrix.
Recall that in a transformer-like model without shortcut, the upper bound of feature diversity r(Zl) decreases dramatically as the increase of network depth(Theorem 1). In the following, we analyze how the diversity r(Zl) changes w.r.t. the layer l in the model stacked by the AugMSA modules, and we has the follow theorem.
Theorem 2 Given a model stacked by the AugMSA modules, the diversity r(Xl) of feature in the l-th layer can be bounded by that of input data Z0, i.e.,
r(Zl) ≤ max 0≤m≤l ( Hγ√ d ) 3m−1 2 (2Hαm) 3m(l−m)r(Z0) 3m , (10)
where αm = 1 + ∑T
i=1 λ‖Θmi‖. Θmi is the weight matrix in the i-th augmented shortcut of the m-th layer, and λ is the Lipschitz constant of activation function σ(·).
Compared with Theorem 1, the augmented shortcuts introduce an extra term (2Hαm)3 m(l−m), which will increase doubly exponentially as 2Hαm is usually larger than 1. This tends to suppress the diversity decay incurred by attention mechanism. The term αm (0 ≤ m ≤ l) is determined by the norms of weight matrices Θmi of the augmented shortcuts in the m-th layer, and then bound of diversity r(Zl) in the l-th layer can be affected by all the augmented shortcuts in the previous layers. For the ShortcutMSA module (Eq. 7) with only a identity shortcut, we have αm = 1. Adding more augmented shortcuts can increase the magnitude of αm, which further improves the bound. Detailed proof for Theorem 2 is represented in the supplemental material.
Considering that shortcut connections exist in both MSA and MLP modules, the proposed augmented shortcuts can also be embedded into MLP similarly, i.e.,
AugMLP(Z ′l) = MLP(Z ′ l) + Z ′ l + T∑ i=1 Tli(Z ′l ; Θ′li), l ∈ [1, 2, · · · , L], (11)
where Z ′l is the input feature of the MLP module in the l-th layer and Θ ′ li denotes the parameters in augmented shortcuts. Paralleling the MLP module with the augmented shortcuts can further improve the diversity, which is analyzed detailedly in the supplemental material. Stacking the AugMSA and AugMLP modules constructs the Aug-ViT model, whose feature has stronger diversity as shown in Figure 3(c) and Figure 4.

heading: 3.2 Efficient Implementation via Circulant Projection
text: As discussed above, Paralleling multiple augmented shortcuts with the MSA and MLP modules in a vision transformer can improve the feature diversity for higher performance. However, directly implementing Tli(·) (Eq. 9) involves a lot of matrix multiplications, which are computationally expensive [19, 43, 32]. For example, given feature Zl ∈ Rn×d and weight matrix Θli ∈ Rd×d, the matrix multiplication ZlΘli consumes nd2 FLOPs, where d is usually large in vision transformers (e.g., 798 in ViT-B). Therefore, we propose to implement the augmented shortcuts with blockcirculant matrices, whose computational costs are negligible compared to other modules in a vision transformer. In the following, we omit the subscripts l and i indicating layers and paths for brevity.
Circulant matrix [19, 7] is a typical structured matrix with extremely few parameters and low computational complexity on the Fourier domain. A circulant matrix C ∈ Rd′×d′ only has d′ parameters and the product between C and a vector only has O(d′ log d′) computational complexity via the fast Fourier transformation (FFT). The projection with circulant matrix also has theoretical guarantee to be a good approximation of unstructured projection and preserve the critical projection properties such as `2 distance and angles [17, 10]. Consequently, we take advantage of circulant matrices to implement augmented shortcuts. Specially, the original matrix Θ is split into b2 submatrices Cij ∈ Rd′×d′ , i.e.,
Θ =  C11 C12 · · · C1b C21 C22 · · · C2b ... ... ... ...
Cb1 Cb2 · · · Cbb
, (12)
where d′ is the size of sub-matrices and bd′ = d. Each sub-matrix Cij is a circulant matrix generated by circulating the elements in a d′-dimension vector cij = [cij1 , c ij 2 , · · · , c ij d′ ], i.e.,
Cij = circ(cij) =  cij1 c ij d′ · · · c ij 3 c ij 2 cij2 c ij 1 c ij d′ c ij 3 ... cij2 c ij 1 . . . ...
cijd′−1 . . . . . . cijd
cijd′ c ij d′−1 · · · c ij 2 c ij 1
. (13)
For efficient implementation of the projection T (Z) = σ(ZΘ), we first split the input Z into b slices Zj ∈ RN×d′ , i.e., Z = [Z1;Z2; · · · ;Zb], and then multiply each slice Zj by a circulant matrix Cij . The product between circulant matrix and vector in the original domain is equivalent to element-wise multiplication in the Fourier domain, i.e., F(ZjCij) = F(Zj) F(cij), where F(·) is the discrete Fourier transformation and denotes that each row in Zj is multiplied by the according element in cij . The discrete Fourier transform and its inverse transformation both can be efficiently calculated with fast Fourier transform FFT and its inverse IFFT with only O(d′ log d′) computational complexity. The output is calculated as:
T (Z)i = σ  b∑ j=1 ZjCij  = σ  b∑ j=1 IFFT(FFT(Zj) ◦ FFT(cij))  , (14) where T (Z)i ∈ RN×d′ is a slice of T (Z). Finally, T (Z) is obtained by concatenating slices T (Z)i, i.e., T (Z) = [T (Z)1; T (Z)2; · · · ; T (Z)b]. Complexity Analysis. For a d × d matrix Θ split into multiple b2 sub-matrices, the numbers of parameters is only bd, which has linear complexity with matrix size d. Note that the learnable parameters cij can be stored as its Fourier form, and both the FFT and IFFT operations in Eq. 14 are conducted one time. Thus the computational cost (i.e., FLOPs) of Eq. 14 is about (d log(d/b) + 2bd), where d log(d/b) comes from the FFT and IFFT transformation and 2bd is consumed by the elementwisely operation in the Fourier domain [26, 2]2. The size of matrix d in a transformer is usually large while the number of sub-matrices b is small, and thus the parameters and computation costs incurred by the augmented shortcuts can be negligible. For example, in the ViT-S model with d = 384, b is set to 4 in our implementation. The augmented shortcuts only adds 0.07 M parameters and 0.08 G FLOPs, which is negligible considering the ViT-S model has 22.1 M parameters and 4.6 G FLOPs.

heading: 4 Experiments
text: In this section, we conduct extensive experiments to demonstrate the effectiveness of the proposed augmented shortcuts. The vision transformer [9] and its SOTA variants are equipped with the augmented shortcuts for performance improving. We first compare the performances of different models on the ImageNet dataset for the image classification task. Then ablation studies are conducted to analyze the algorithm. To validate its generalization ability, we further test the proposed Aug-ViT model on the object detection and transfer learning tasks.

heading: 4.1 Experiments on ImageNet
text: Dataset. ImageNet (ILSVRC-2012) dataset [6] contains 1.3 M training images and 50k validation images from 1000 classes, which is a widely used image classification benchmark.
Implementation details. We use the same training strategy with DeiT [34] for a fair comparison. Specifically, the model is trained with AdamW [24] optimizer for 300 epochs with batchsize 1024. The learning rate is initialized to 10−3 and then decayed with the cosine schedule. Label smoothing [30], DropPath [21] and repeated augmentation [18] are also implemented following DeiT [34]. The data augmentation strategy contains Rand-Augment [5], Mixup [42] and CutMix [41]. Besides the original
2The constants are approximated by considering the relation between complex operation and real operation, as well as the symmetry property [2]
shortcuts, one or two augmented shortcuts are added, where the hyper-parameter b for partitioning matrices is set to 4 empirically. The models are trained from scratch on ImageNet and no extra data are used. All experiments are conducted with PyTorch [28] and MindSpore 3 on NVIDIA V100 GPUs. The model architecture can be found in the MindSpore model zoo 4.
Backbones. We apply the augmented shortcuts on multiple vision transformer models. ViT [9] is the typical transformer model for the vision tasks, which splits an image to multiple 16 × 16 patches. Deit [34] adopts the same architecture with ViT but improves the training strategy for better performance. T2T [40] and PVT [37] are two recently proposed SOTA variants of ViT. T2T [40] improves the process of producing patches by considering the structured information in images. PVT [37] designs a pyramid-like structure by partitioning the model into multiple stages. In Table 2, ‘ViT(DeiT)-S’ ,‘ViT(DeiT)-B’ and ‘PVT-S’, ‘PVT-B’ denote the DeiT and PVT models with different size. ‘T2T-14’, ‘T2T-19’ and ‘T2T-24’ are the T2T models with different depths.
Experimental Results. The validation accuracies of different models on ImageNet are shown in Table 2. 5 We firstly equip the ViT models with the augmented shortcuts and get Aug-ViT, which show large superiority to the plain counterparts, i.e., more than 1% accuracy improvement without noticeable parameter and computational complexity increasing. For example, the proposed ’Aug-ViTS’ achieves 80.9% top-1 accuracy with 4.6G FLOPs, which suppresses the baseline (’ViT(DeiT)-S’ with 4.6G FLOPs) by 1.2% top-1 accuracy while the computational cost is barely changed. For a large model with higher input image resolution (e.g., ViT-B with input resolution 384 × 384) and high performance, equipping it with the augmented shortcuts can still further improve the performance (e.g., 83.1%→ 84.2%). Besides the typical ViT model, the augmented shortcuts can be embedded into multiple vision transformers flexibly and improve their performance as well. For example, equipping the PVT-M with augmented shortcut can improve its accuracy form 81.2% to 82.3%. For the T2T-14 model, the performance improvement is even more obvious, i.e., 1.3% accuracy improvement from 82.3% to 83.6%.
Performance Improvement w.r.t. Network Depth. It is interesting to see that the augmented shortcuts improve the performance of deeper models more obviously. ‘T2T-14’, ‘T2T-19’ and ‘T2T24’ compose of the same blocks but have different depths. For the baseline, increasing the depth of T2T model from 14 to 24, the accuracy is only improved by 0.7% (from 81.5% to 82.3%). While with the augmented shortcuts, the Aug-T2T model with 24 layers can achieve 83.6%, which achieves more obvious performance improvement than models with 14 layers. Deeper models tend to suffer more serve feature collapse suffer more serve feature collapse as features from different patches are aggregated with more attention layers. The proposed augmented shortcuts alleviate the feature collapse and release the potential of deep models, which results in higher performance.

heading: 4.2 Ablation Studies
text: To better understand the proposed augmented shortcuts for vision transformers, we conduct extensive experiments to investigate the impact of each component. All the ablation experiments are conducted based on ViT(DeiT)-S model on the ImageNet dataset.
The number of augmented shortcuts. The performance varies w.r.t. the number of augmented shortcuts as shown in Table 3. Besides the original identity shortcut, adding only one augmented shortcut can significantly improves the performance of the ViT model (e.g., 0.8% top-1 accuracy improvement compared to the baseline). Further increasing the number of augmented shortcuts will further improve the performance, and the improvement margin will be saturated gradually. We empirically find that two augmented shortcuts are enough to achieve obvious performance improvement.
Location for implementing the augmented shortcuts. As discussed before, the augmented shortcuts can be paralleled with both MSA and MLP modules to increase the feature diversity. Table 4 shows how the implementation location affects the final performance. Paralleling MSA with the
3https://www.mindspore.cn 4https://gitee.com/mindspore/models/tree/master/research/cv/augvit 5Note that DeiT [34] adopts the same model architecture with ViT [9] but achieves higher performance by
adjusting the training strategy, which is used as the baseline model.
augmented shortcuts significantly improves the performance (e.g., 0.8% top-1 accuracy), which we attribute it to the increasing of feature diversity. Enhancing MLP module also bring the performance improvement, which is accordant to our analysis in supplementary materials. Combining them together can achieve the highest performance (1.1% accuracy improvement), which we adopt in our implementation.
Efficiency of the block-circulant projection. In the block-wisely circulant projection, the hyperparameter b controls the number of sub-matrices Cij partitioned by original matrix Θ. Table 5 shows how the performance varies w.r.t. the parameter b. A larger b implies the matrix Θ will be partitioned into more circulant matrix with smaller sizes, which brings more parameters and higher computational cost, as well as the performance improvement. The unstructured projection can also be used as augmented shortcuts and brings performance improvement, but it incurs obvious increasing of parameters and computational cost. Using the block-circulant projection with b = 4 can achieve very similar performance with the unstructured projection but has much few parameters (e.g., 0.07M vs.7.1M), implying that the block-circulant projection is an efficient and effective formulation to transform features in the augmented shortcuts.
Formulation of the augmented shortcuts. The augmented shortcuts are implemented sequentially with the block-circulant projection and activation function (e.g., GeLU). Table 6 shows the impact of each component. Only using the activation function introduces no learnable parameters, which only achieves similar accuracy with the baseline model, implying that the learnable parameters are vital to produce diverse features. If only the linear circulant projection is kept, multiple augmented shortcut can be merged to a single one. Swapping the sequential order of the circulant projection and activation function has negligible influence on the final performance.
Drop rates of stochastic depth and dropout. We change the parameters of stochastic depth/dropout for both ViT(DeiT)-S model and Aug-ViT-S, and the results are shown in Table 7 and Table 8. We empirically find the default parameters for DeiT (drop rates of stochastic depth/Dropout are set as 0.1/0.0) also work well for the Aug-VIT models. In addition, the models equipped with augmented shortcuts tend to have higher tolerance to over-large drop rates. For example, with drop rate=0.3 for
Table 5: Efficiency of the block cifculant projection.
Type b Top1-Acc. Params FLOPs(%) (M) (G) Baseline - 79.8 22.1 4.6
Block-Circulant 1 80.5 (+0.7) +0.02 +0.07 Block-Circulant 2 80.7 (+0.9) +0.03 +0.07 Block-Circulant 4 80.9 (+1.1) +0.07 +0.08 Block-Circulant 8 80.9 (+1.1) +0.15 +0.10
Unstructured - 81.0 (+1.2) +7.1 +1.4
Table 6: Formulation of the augmented shortcuts. ‘Act’ denotes the activation function.
Form Top1-Acc. (%) Baseline 79.8
Act. 79.9 Circulant 80.6 Act. + Circulant 80.8 Circulant + Act. 80.9
dropout, the accuracy of Aug-ViT-S can still achieve 75.1% accuracy, while the accuracy of DeiT-S decreases to 69.8%.
Feature Visualization. We intuitively show the features of different models in Figure 5, From top to bottom are features in low, middle and deep layers of the ViT-Small model. The input image is scaled to 1024× 1024 for better visualization and the patch embeddings are reshaped to their spatial positions to construct the features maps. Without shortcut, the feature maps in deep layers conveys no effective information (Figure 5 (a)), and adding a shortcut connection make the feature maps informative ((b)). Compared with them, the features in the Aug-ViT model are further enriched, especially for the deep layers.

heading: 4.3 Object Detection with Pure Transformer
text: We also validate the effectiveness of the augmented shortcuts on the object detection task. A pure transformer detector can be constructed by combining the vision transformer backbone and the DETR [3] head, and we equip the backbones with the augmented shortcuts. For a fair comparison, we follow the training strategy in PVT [37] and fine-tune the models for 50 epochs on the COCO train2017 dataset. Random flip and random scale are used as the data augmentation strategy. The results on COCO val2017 are shown in Table 9. The detector equipped with augmented shortcut achieve better performance than the base model. For the DeiT-S model with 33.9 AP, the augmented shortcut can improve 1.8% AP and achieve 35.7%.

heading: 4.4 Transfer Learning
text: To validate the generalization ability of vision transformers equipped with the projection, we conduct experiments on the transfer learning tasks. Specifically, the models trained on ImageNet is further
fine-tuned on the downstream tasks, containing superordinate-level image recognition dataset (CIFAR10 [20], CIFAR-100 [20]) and fine-grained image recognition dataset( Oxford 102 Flowers [25] and Oxford-IIIT Pets [27]). Following [34], we fine-tune the models on images with resolution 384×384, and use the same fine-tuning strategy as [34]. Table 10 shows the performances of different models on the downstream tasks, and the model equipped with the augmented shortcut always achieves higher accuracies than the baseline on different tasks.

heading: 5 Conclusion
text: We presented augmented shortcuts for resolving the feature collapse issue in vision transformers. The augmented shortcuts are parallel with the original identity shortcuts, and each connection has its own learnable parameters to make diverse transformations on the input features. Efficient circulant projections are used to implement the augmented shortcuts, whose memory and computational cost are negligible compared with other components in the vision transformer. Similar to the widely used identity shortcuts, the augmented shortcuts do not depend on the specific architecture design either, which can be flexibly embedded into various variants of vision transformers (e.g., ViT [9], T2T [40], PVT [37]) for enhancing their performance on different tasks such as image classification, object detection and transfer learning. In the future, we plan to research designing deeper vision transformers with the help of augmented shortcuts.

heading: Acknowledgment
text: This work is supported by National Natural Science Foundation of China under Grant No. 61876007, and Australian Research Council under Project DE180101438 and DP210101859.

Review summaries:
This paper introduce a universal module to improve vision transformers. They show that vision transformers suffer feature collapse, which mean that features in deep layers will become similar. To address, this paper presents an augmented shortcut module to  be flexibly embedded into various transformer models. Extensive experiments are conducted on multiple tasks (image classification, object detection, and transfer learning), which show that augmented shortcuts can bring obvious performance improvements. That is a good paper to me because of the new and effective module which is a new investigation to the field.

This work proposes to boost the performance of vision transformers from shortcuts. As analyzed, the feature diversity of tokens declines as the layers go deeper while the shortcut connection provide a fix towards to potential ``feature collapse'' issue. To this end, additional branches of augmented shortcut is proposed to further improve the feature diversity as well as the final recognition power.

This paper observes the problem of feature collapse in transformer, where the tokens in deep layers collapse to the same representation. It then proposes to add additional augmented shortcuts in parallel to the original shortcut connection in the multihead self-attention modules and MLP modules, based on the observation that simple identity shortcut alleviates the problem of feature collapse. In the experiments, the proposed method brings 1.1% performance gain in the ImageNet classification task compared to the baseline model. Performance gains can also be seen from COCO detection task and other downstream finegrained classification tasks, as compared to the baseline models.

This paper alleviates feature collapse in vision transformers by introducing an efficient module, named augmented shortcut. As shown in the theoretical and empirical results, the feature diversity is enhanced and obvious performance improvement is achieved. The authors conduct experiments with various SOTA vision transformer models and validate the effectiveness of the proposed augmented shortcuts.