LMSA: Low-relation Mutil-head Self-Attention Mechanism in Visual Transformer

1 INTRODUCTION . The self-attention mechanism proposed by Vaswani et al . ( 2017 ) is originated from the natural language processing ( NLP ) Transformer network . Due to its superior performance , the self-attention mechanism has been widely used in the NLP field . The standard self-attention formula is defined as : Attention ( Q , K , V ) = softmax ( QKT√ dim ) V ( 1 ) The standard self-attention mechanism converts the input features linearly to obtain the variables Query ( Q ) , Key ( K ) , and Value ( V ) with the same dimension ( dim ) , then normalizes the result from the vector dot product of the Q and K , and then obtains the potential relationship between every two features , finally uses this potential relationship to re-combine the features to get the output features . Compared to recurrent neural network ( RNN ) ( Lipton et al. , 2015 ) series , self-attention is characterized by directly calculating dependencies regardless of the distance between features and can be calculated in parallel . ( Tang et al. , 2018 ) demonstrate that the self-Attention mechanism is better at Word Sense Disambiguation ( WSD ) tasks ( Raganato et al. , 2017 ) . The authors believe that the powerful semantic feature extraction ability is the key reason why self-Attention has better performance than RNN . Most recent studies have focused on improving the efficiency of the self-attention mechanism as it brings the computational complexity of O ( N2 ) to reach a considerable performance gain . Therefore , a batch of variants that pay more attention to computational efficiency are proposed on the basis of original transformer backbone , such as LongFormer ( Beltagy et al. , 2020 ) , Reformer ( Kitaev et al. , 2020 ) , Faster-Transformer , Turbo-Transformers ( Fang et al. , 2020 ) etc .. Through changing the feature sequence , partial self-attention , hardware acceleration , and other methods , the Transformer model can be better used in the NLP field . In the meanwhile , the excellent performance of the self-attention mechanism in the field of NLP has attracted great attention in computer vision . The Vision Transformer ( ViT ) backbone model proposed by ( Dosovitskiy et al. , 2020 ) divides the image information into patches for parallel input into the model , discards the decoder in the original transformer model , and retains its encoder as the feature extractor connected to the MLP layer to get output features . The results show that ViT is potentially competitive to the CNN network series ( He et al. , 2016 ) , which increasingly makes the Self-Attention mechanism popular for visual tasks . In addition to processing input and output features , ViT obeys the conventional Transformer Block , which makes ViT have high computational complexity , and coupled with the diversity of image feature information , it makes ViT training difficult . In essence , Transformer is data-hungry ( Khan et al. , 2021 ) , which makes it a consensus to establish more efficient and extensive attention fusion image features . ( Liu et al. , 2021 ) proposes Swin Transformer , which further divides the features into several windows based on patch division . Swin introduces Non-Local neural network ( Wang et al. , 2018 ) to model the relationship of each patch in windows and realizes the cross-window relationship modeling through the shift window operation . Recently , Cswin ( Dong et al. , 2021 ) abandons the shift windows method , set windows as stripe , and models the horizontal and vertical relationships of features at the same time . Even there are many similar applications of Transformer to the field of vision , but the works aim at optimizing the visual self-attention mechanism is relatively rare . This paper proposes a novel self-attention mechanism Low-relation Mutil-head Self-Attention ( LMSA ) to reduce the number of feature mapping Query and Key dimensions while maintaining the V alue dimension . This exploration breaks the information consistency of the traditional self-attention mechanism and greatly saves the computational complexity of a single self-Attention block . The diagram of LMSA is shown in Figure 1 , which also illustrates the difference between the proposed LMSA and the newly proposed multi-head self-attention mechanism . Specifically , we modified the self-attention mechanism based on recently proposed models such as Swin and Cswin , and we conducted a controlled experiment on the two attention mechanisms . In the Swin model , by appropriately compressing the number of dimensions , the image classification performance of the model has been further improved . In the Cswin model , we maximize the dimensionality compression of Query and Key , while its performance on different datasets remains competitive to the traditional self-attention mechanism . More importantly , through experiments , we demonstrate that the self-attention mechanism doesn ’ t have to align the dimensions of Query , Key , and V alue to be same . The excessively high dimensions of Query and Key not only cause model data redundancy but even have a negative impact on the final model performance . 2 RELATED WORK . 2.1 MUTIL-HEAD SELF-ATTENTION MACHINISM . For a long time , CNN and RNN as the first choice for sequence task models have the disadvantage that they can not be calculated in parallel . The input of each subsequence depends on the output of the previous subsequence , and as the sequence length increases , it will be accompanied by long-term memory failure . The Transformer was proposed by ( Vaswani et al. , 2017 ) . Transformer abandons the recurrent network structure and uses scaled dot-Product attention ( see Equation 1 ) as its computing core to obtain the correlation of features at any two positions , and perform feature fusion based on this . In order to enable the model to pay more attention to the information of different representation subspaces at different positions , the author embeds the multi-head mechanism into self-attention . The formula of the multi-head self-attention mechanism is as follows : MutilHead ( Q , K , V ) = Concat ( head1 , ... , headh ) W 0 Where headi = Attention ( QW Q i , KW K i , V W V i ) ( 2 ) Where WQi ∈ Rdmodel×dk , WKi ∈ Rdmodel×dk , WVi ∈ Rdmodel×dv , WO ∈ Rhdv×dmodel . Each head uses a different weight transformation , all heads are calculated in parallel , and the calculation results are concatenated to obtain the final output . Adjust the output scale by controlling the dimension of V alue . Generally , the dimension of V alue is set to be the same as the input feature . After a complete multi-head attention calculation , the input and output sizes can be guaranteed to be consistent . It is convenient to embed various networks for multi-head attention . 2.2 EFFICIENT ATTENTION . The multi-head self-attention mechanism has achieved great success in the field of NLP , but it has also brought excessive complexity , especially in long text sequence tasks . Memory Compressed Transformer ( Liu et al. , 2018 ) tries to divide a long text sequence into several modules of similar size , and perform self-attention within the modules , and reduce the size of the self-attention matrix through stride convolution , so as to achieve the goal of improving efficiency . ( Child et al. , 2019 ) proposed the Sparse Transformer , which simplifies the dense attention to sparse attention through explicit selection , and improves the concentration of the model . Longformer ( Beltagy et al. , 2020 ) which based on Sparse Transformer combines sliding window and dilated sliding window to achieve local and global Balance of attention . Axial Transformer ( Ho et al. , 2019 ) applies multiple attention along a single axis of the input tensor . Since the length of any single axis is usually much smaller than the total number of elements , this model can significantly save computation and memory . Unlike the previously mentioned Efficient Transformer , Refomer ( Kitaev et al. , 2020 ) has a more in-depth discussion of the self-attention mechanism itself . By designing a reversible Transformer Block , Reformer greatly saves the space used to store intermediate results and proposed that locality sensitive hashing only calculates similar feature relationships , which effectively reduces the computational complexity . Synthesizer ( Tay et al . ) changes the self-attention mechanism more boldly . It directly abandoned the generation and calculation of Query and Key , and instead used random initialization or linear projection of input features to directly replace the calculation results of QKT . The author concludes through experiments that this method has advantages and disadvantages with the traditional self-attention mechanism , and has no affiliation with it , which makes people interested in the effective principle of the self-attention mechanism . 2.3 VISUAL TRANSFORMER . With the continuous development of Transformer in the NLP field , people have also begun to try to introduce the self-attention mechanism into visual tasks . After ViT ( Dosovitskiy et al. , 2020 ) was proposed , more efficient visual Transformers have also been proposed . Among them , the most concerned are Swin Tranformer ( Liu et al. , 2021 ) and Cswin Transformer ( Dong et al. , 2021 ) . Swin Transformer realizes a more flexible local fusion of information through multi-level division of image features and then realizes information interaction between windows through shift operations . Cswin sets the windows in Swin as a bidirectional stripe with cross-feature capability , which is a successful application of Axial Attention in the computer vision field . Both Swin and CSwin adopted the strategy proposed by PiT ( Heo et al. , 2021 ) , dividing the entire model into several Layers , and down-sampling between the two layers . As the features move forward in the model , their scale decreases while their dimensions increase until the final High-Level feature is obtained to continue downstream tasks . The network structure of the Swin and Cswin model is shown in Figure 2 . 2.4 CONVOLUTIONAL VISION TRANSFORMER . Visual Tranformer ’ s extensive research makes the Self-Attention mechanism often used to compare with convolution , Convolutional Vision Transformer ( Wu et al. , 2021 ) groundbreakingly combines convolution and self-attention mechanism . CvT discards position embeding and uses convolutional projection instead of linear projection to generate the Q , K , and V required for self-attention calculations , as shown in Figure 3 . These changes enable CvT to well inherit the advantages of CNN while maintaining the dynamic self-encoding characteristics of Transformer : translation , scaling , and distortion invariance . 3 METHODOLOGY . Unlike text sequence features , visual features often contain more nonobjective information . The development of Transfomer in the computer vision field shows that the self-attention mechanism can effectively extract image features , and there is still much room for improvement . However , the standard self-attention mechanism has a computational complexity of O ( N2 ) . As the image size increases , the self-attention mechanism consumes a lot more computing time and storage space , which limits the efficiency of the self-attention mechanism in visual tasks . To reduce the complexity of the self-attention mechanism and improve the computational efficiency , in this paper , we propose Low-Relation Mutil-Head Self-Attention , which aims to reduce the low-quality data that has a small effect on the calculation result in the self-attention . 3.1 LOW-RELATION MUTIL-HEAD SELF-ATTENTION . In order to solve the problem of the high complexity of the self-attention mechanism , a large number of Efficient Attentions have been proposed ( Tay et al. , 2020 ) . The most common way to increase calculation speed is to control the number of Keys and V alues . Associate the Query generated by the feature map with the Key whose number is much smaller than Query , and act on the scaled V alue to achieve the purpose of improving efficiency by losing a certain accuracy . In the visual task , the most concerned optimization method is to adjust the receptive field of the self-attention mechanism to achieve a balance of accuracy and efficiency . The above self-attention optimization methods mostly focus on adjusting the number of features used for calculation . From the perspective of feature dimensions , we try to compress the number of Query and Key dimensions as much as possible while keeping the V alue dimension unchanged and explore the impact of this compression on the accuracy of self-attention . In the proposed LMSA , the consistency of Query , Key , and V alue dimensions is destroyed . We control the overall complexity by adjusting the dimensions of Query and Key . The core formula of LMSA is expressed as follows : LMSA ( Q , K , V ) = MutilHead ( CP ( Q ) , CP ( K ) , V ) ( 3 ) where CP ( X ) indicates the compress operation of X . In LMSA , we still use the multi-head mechanism to divide the feature into several heads , and each head is calculated in parallel . In the case of extreme compression , after Query and Key are divided into several heads , each head can only retain one dimension , that is , dimQi =1 . In specific tasks , we need to adjust the dimensions of Query and Key to meet the actual needs .

Summary: This paper presents a mechanism to reduce the computation costs of a standard self-attention module, named LMSA.  The basic idea of LMSA is to reduce the dimension of key&query of self-attention(SA) while keeping the dimension of value unchanged. Therefore, the computational complexity of SA will be reduced from O(N^2D) to O(N^2D_q,k).  They study the proposed mechanism on image classification and semantic segmentation.