Augmented Shortcuts for Vision Transformers

1 Introduction . 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 ) . 2 Preliminaries and Motivation . 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 .

Summary: 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.