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