Adder Attention for Vision Transformer

1 Introduction . Transformers [ 27 ] have prevailed in natural language processing ( NLP ) due to its superior capability in capturing long-distance dependencies based on self-attention mechanism . Some homologous large-scale models , e.g. , BERT [ 10 ] and GPT-3 [ 2 ] provide a significant performance improvement on learning powerful language representations from unlabeled text . The great breakthrough of transformers in NLP has sparked particular interest from the vision community in applying transformer to computer vision ( CV ) tasks such as image recognition [ 11 , 35 , 25 , 26 ] , object detection [ 3 , 38 , 33 ] , and image generation [ 4 , 12 , 17 , 15 ] . The high-power consumption of transformer-based models has blocked them from being deployed on mobile devices , e.g. , smart phone , camera , and micro-robots . Therefore , it is necessary to study efficient transformers which can be embedded on mobile devices with affordable computation resources . Wu et.al . [ 31 ] present an efficient Long-Short Range Attention with two groups of heads , respectively for local context and long-distance relationship modeling . Wang et.al . [ 28 ] decompose the original dot-product attention into smaller attentions through linear projections . Unfortunately , previous arts mainly focus on compressing and accelerating transformers for language processing tasks , which motivates us for building efficient transformers for computer vision applications . Various attempts have been made on compressing CNN models , including quantization [ 13 , 16 ] , pruning [ 19 , 30 , 23 ] and distillation [ 5 , 25 ] . Beyond these methods , Chen et.al . [ 6 ] propose Adder Neural Network ( AdderNet ) for building efficient deep learning models , which avoids massive computational cost by replacing the convolutional operation by using ` 1-distance between input * Equal contribution . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . signals and weights instead of correlation . Xu et.al . [ 32 ] further utilize knowledge distillation technique to make AdderNet achieve better performance than CNN . Song et.al . [ 24 ] successfully apply AdderNet on single image super-resolution ( SISR ) task by learning identity mapping with self-shortcuts and implementing high-pass filter with a learnable power activation . Wang et.al . [ 29 ] conduct practical implementations on hardware platform ( FPGA ) , and verify its superior performance in suppressing chip area and power consumption . Thus , we are motivated to investigate the feasibility of replacing multiplications by additions in transformer architectures . However , it is very hard to directly transport the existing success of AdderNet on CNNs to transformers , since their basic calculation paradigm is completely different . In particular , the existing works mainly focus on the additiveization of convolutional kernel filters which do not involve the unique linear transformations and self-attention mechanisms in transformers . To this end , we propose a general adder linear transformation operation and adder self-attention mechanism to replace the traditional multiplication version in both feed-forward modules and self-attention layers in transformer model . Specifically , we utilize ` 1-distance between query and key instead of scaled dot-product to measure the distance between them . For the attention module , we find that the directly using additions could cause more information concentrating on a few largest singular values , leading to a decay of the attention map rank , i.e. , the functionality of the attention mechanism can not be well preserved . Thus , we propose to insert an extra identity mapping in the adder attention module to solve the decay . This operation can also homogenize the information distribution of the adder attention map . Meanwhile , we provide the detailed feed-forward and back-propagation for optimizing the adder vision transformer with stable training process . As a result , we can obtain comparable performance using the proposed adder transformer models to that of original baselines while reducing about 3× of the overall energy consumption on various benchmark models and datasets . 2 Preliminaries and Motivation . In this section , we briefly revisit the basic related components , including AdderNet and transformer . Adder Neural Networks ( AdderNet ) . Denote filters in a convolutional layer of neural networks as W ∈ Rk×k×cin×cout , the input feature map as X ∈ RH×W×cin , and the output feature map as Y ∈ RH×W×cout , where k is the kernel size , cin and cout are the number of input channels and output channels , respectively . The traditional convolutional operation is defined as : Ym , n , q , X ∗ F = D∑ i=1 D∑ j=1 cin∑ k=1 S ( Xm+i , n+j , k , Fi , j , k , q ) . ( 1 ) where S ( · , · ) is a pre-defined similarity measure ( S ( x , y ) = x× y for the convolution multiplications operation ) . AdderNet [ 6 ] maximizes the use of additions by treating the ` 1-distance between the input feature map and filter as the similarity measure ( i.e. , S ( x , y ) = |x− y| in Eq . 1 ) : Ỹm , n , q , X ⊕ F = − D∑ i=1 D∑ j=1 cin∑ k=1 |Xm+i , n+j , k − Fi , j , k , q| . ( 2 ) However the pioneering work only focused on additiveization of CNNs , while ignoring transformers . Vision Transformers . Vision transformer [ 11 ] mainly consists of MHSA ( Multi-head Self-Attention ) , FFN ( Feed-Forward Network ) and LN ( Layer Normalization ) . For multi-head self-attention , the inputs X ∈ RN×D are applied with three different linear transformations and output queries Q ∈ RM×N×dq , keys K ∈ RM×N×dk and values V ∈ RM×N×dv , where M is total number of heads , dq , dq , dv is the dimension of queries , keys and values for each head , respectively . Generally , D =Mdq =Mdk =Mdv =Mdh and the multi-head self-attention is performed as follows : Am , j , i = exp { Qm , j , :K T m , i , :√ dt } Zj , where Zj = Nkv∑ i=1 exp { Qm , j , :K T m , i , :√ dt } , Om , j , v = Nkv∑ i=1 Am , j , i · Vm , i , v , O ′ j , d = Concat ( O1 , : , : , O2 , : , : , · · · , OM , : , : ) WO . ( 3 ) where dt = dh , WO ∈ RD×D denotes the weight matrix of a linear projection layer used to produce the output values O ′ ∈ RN×D . From the above statements , transformer has a different structure and computation paradigm compared to CNNs . Thus , it is necessary to study how to use additions to fulfill the modules . 3 The Proposed Model : Adder Transformer . In this section , we present the Adder Transformer— which implement the multi-head attention module and FFN module using adder operations . We will present the details in the following subsections and an illustration of the proposed method can be found in Section A.1 of the supplementary material . 3.1 Adder Linear Transformation . We first consider the linear transformation and FFN layers in transformers . The calculation paradigm of these two modules is based on fully connected layer , which is equivalent to a special 1x1 convolution . Therefore , we can use the original computation mechanism of adder convolutional kernel in this module . Considering weight matrices WQ ∈ RM×dq×D , WK ∈ RM×dk×D , WV ∈ RM×dv×D in a projection layer of the multi-head self-attention , we maximize the use of additions by taking ` 1 distance to measure the distance between the weight matrices and input embedding as : Qm , j , q = D∑ d=1 −|Xj , d −WQm , q , d | , Km , i , k = D∑ d=1 −|Xi , d −WKm , k , d | , Vm , i , v = D∑ d=1 −|Xi , d −WVm , v , d | , ( 4 ) wherein , m ∈ [ 1 , M ] , j ∈ [ 1 , Nq ] , i ∈ [ 1 , Nkv ] , q ∈ [ 1 , dq ] , k ∈ [ 1 , dk ] , v ∈ [ 1 , dv ] and generally we set N = Nq = Nkv as the resulting number of patches and D =Mdq =Mdk =Mdv =Mdh as the model embedding dimension . The feed-forward network can also be modified directly : FFNa ( X ) =W2 ⊕ ( W1 ⊕X + b1 ) + b2 , ( 5 ) As mentioned in AdderNet [ 6 ] , the adder layers output values with large magnitude and should be followed by a batch normalization layer . Similarly , we apply layer normalization ( LN ) to scale the output values of each layer , since LN is widely used in the transformer architecture . In fact , although LN introduces multiplications , the magnitude of the multiplication can be omitted compared to the total computation of a layer . The FLOPs of adder MHSA is DM ( Nq + Nkv ) ( dk + dv ) + NqNkvM ( dk + dv ) , and the FLOPs of adder FFN is 2NqDrD , where r is the dimension expansion ratio . The overall FLOPs of a standard adder transformer block is O ( 2ND ( 6D +N ) ) , for regular setting r = 4 , N = Nq = Nkv . The overall FLOPs of multiplication in Layer Normalization are O ( 12ND ) . In practice , given D = 768 and N = 197 , the ratio of a standard adder transformer block with the LN could be estimated as 2ND ( 6D+N ) 12ND ≈ 800 , which shows that LN can be easily applied in adder transformer with tiny extra energy cost . For back-propagation , directly calculating the partial derivative of output w.r.t input and weight will lead to a sign update of { −1 , 0 , +1 } value , respectively . Back-propagation using the sign gradient is detrimental to the efficient updating of the parameters because the direction of the gradient can never reach the steepest . Thus , we follow AdderNet [ 6 ] to compute the derivative of the ` 2-norm with a HardTanh function to prevent gradients from exploding . Back-propagation process is formulated as : ∂Qm , j , q ∂Xj , d = HT ( WQm , q , d −Xj , d ) , ∂Km , i , k ∂Xi , d = HT ( WKm , k , d −Xi , d ) , ∂Vm , i , v ∂Xi , d = HT ( WVm , v , d −Xi , d ) ( 6 ) where HT ( · ) denotes the HardTanh function , and adder linear transformation in FFN modules follows the same inference and optimization process . 3.2 Adder Multi-Head Self-attention . 3.2.1 Inference process in adder multi-head self-attention . Self-attention mechanism can be viewd as the measument of the similarity of input query and key matrices . The output is usually positively correlated with the value and the correlation is determined by the magnitude of the attention score after normalization . We design the adder multi-head selfattention in strict compliance with the above principles which measure the similarity between vectors with the help of the ` 1-norm , an effective measure to avoid multiplication operations . Hence , by calculating ` 1-distance between each query and key vector , adder multi-head self-attention can be formulated as : Am , j , i = exp { −‖Qm , j , :−Km , i , :‖1√ da } Zj where Zj = Nkv∑ i=1 exp { −‖Qm , j , : −Km , i , :‖1√ da } , da = 2dt ( 1− 2 π ) , ( 7 ) wherein , Q ∈ RM×Nq×dq , K ∈ RM×Nkv×dk , A ∈ RM×Nq×Nkv , m ∈ [ 1 , M ] indicates head index , j ∈ [ 1 , Nq ] , i ∈ [ 1 , Nkv ] and generally we set N = Nq = Nkv as number of patches , dt = dh indicates embedding dimension for each head . 1√ dt and 1√ da represent the scaling factor of the dot-product attention function and adder attention function , respectively . Theorem 1 . Assuming that the components of Qm , j , : and Km , i , : are independent random variables following normal distribution , the variance of dot-product and ` 1-distance between Qm , j , : and Km , i , : can be formulated respectively as follows : V ar ( QTm , j , :Km , i , : ) = dt , V ar ( −‖Qm , j , : −Km , i , :‖1 ) = 2dt ( 1− 2 π ) , ( 8 ) where the independent assumption is followed by [ 27 ] and the proof is provided in Section A.2 of the supplementary material . According to Theorem . 1 , scaling factor 1√ dt in dot-product attention is to counteract the variance explosion , and we adjust the scaling factor to match the adder circumstances . For the process of yielding output values , the model jointly attend to information at different positions according to the attention map . Specifically , transformer model multiply the attention map with the value matrix directly to allocate each value with corresponding attention weights as in Eq . 9 . For the projection to the output , we still use additions descibed in Sec . 3.1 , where WO ∈ RD×D is the weight matrix of the projection adder linear layer . Om , j , v = Nkv∑ i=1 Am , j , i · Vm , i , v , Ôm , j , v = LayerNorm ( Om , j , v ) , O ′ j , d = − D∑ s=1 ∣∣∣∣Concat ( Ô1 , : , : , Ô2 , : , : , · · · , ÔM , : , : ) j , s −WOs , d ∣∣∣∣ . ( 9 )

Summary:  This paper designs the novel Adder paradigm for vision transformer. Specifically, the attention module is conducted with additions which saves energy consumption. This paper demonstrate the expressive power of addition computation paradigm both theoretically and experimentally on vision transformers. Experiments results show comparable accuracy on classification tasks on cifar10，cifar100 and ImageNet while saving the energy cost by 2x-3x .