SOFT: Softmax-free Transformer with Linear Complexity

1 Introduction . Recently the step change brought by Transformers [ 33 ] in natural language processing ( NLP ) [ 10 , 4 ] seems to have arrived in vision [ 11 , 41 , 47 , 46 ] . Indeed , with less inductive bias in its architecture design than Convolution neural networks ( CNNs ) , pure Vision Transformer ( ViT ) [ 11 ] and its variants have shown to be able to outperform CNNs on various vision tasks [ 8 , 15 ] . However , there is a bottleneck in any Transformer based model , namely its quadratic complexity in both computation and memory usage . This is intrinsic to the self-attention mechanism : given a sequence of tokens ( e.g. , words or image patches ) as input , the self-attention module iteratively learns the feature representations by relating one token to all other tokens . This results in a quadratic complexity O ( n2 ) with the token sequence length n in both computation ( time ) and memory ( space ) since an n × n sized attention matrix needs to be computed and saved during inference . This problem is particularly acute in vision : a 2D image after tokenization will produce a far longer sequence than those in NLP even with a moderate spatial resolution . This quadratic complexity thus prevents a ViT model from modeling images at high spatial resolutions , which are often crucial for visual recognition tasks . ∗Li Zhang ( lizhangfd @ fudan.edu.cn ) is the corresponding author with School of Data Science , Fudan University . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A natural solution is to reduce the complexity of self-attention computation via approximation . Indeed , there have been a number of attempts in NLP [ 34 , 5 , 18 , 39 ] . For example , [ 34 ] takes a naive approach by shortening the length of Key and Value via learnable projections . Such a coarse approximation would inevitably cause performance degradation . In contrast , [ 5 , 17 ] both leverage the kernel mechanism to approximate softmax normalization to linearize the computation in self-attention . [ 18 ] instead adopts a hashing strategy to selectively compute the most similar pairs . Recently , [ 39 ] uses Nyström matrix decomposition to reconstruct the full attention matrix with polynomial iteration for approximating the pseudo-inverse of the landmark matrix . Nonetheless , softmax normalization is simply duplicated across the matrix decomposition process , which is theoretically unsound . We empirically found that none of these methods are effective when applied to vision ( see Sec . 4.2 ) . In this work , we identify that the limitations of existing efficient Transformers are caused by the use of softmax self-attention , and for the first time propose a softmax-free Transformer . More specifically , in all existing Transformers ( with or without linearization ) , a softmax normalization is needed on top of scaled dot-product between token feature vectors [ 33 ] . Keeping this softmax operation challenges any subsequent linearization efforts . To overcome this obstacle , we introduce a novel softmax-free self-attention mechanism , named as SOFT , with linear complexity O ( n ) in both space and time . Specifically , SOFT uses Gaussian kernel to define the similarity ( self-attention ) function without the need for subsequent softmax normalization . With this softmax-free attention matrix , we further introduce a novel low-rank matrix decomposition algorithm for approximation . The robustness of the approximation is theoretically guaranteed by employing a Newton-Raphson method for reliably computing the Moore-Penrose inverse of the matrix . We make the following contributions . ( I ) We introduce a novel softmax-free Transformer with linear space and time complexity . ( II ) Our attention matrix approximation is achieved through a novel matrix decomposition algorithm with theoretical guarantee . ( III ) To evaluate our method for visual recognition tasks , we design a family of generic backbone architectures with varying capacities using SOFT as the core self-attention component . Extensive experiments show that with a linear complexity ( Figure 1b ) , our SOFT models can take in as input much longer image token sequences . As a result , with the same model size , our SOFT outperforms the state-of-the-art CNNs and ViT variants on ImageNet [ 9 ] classification in the accuracy/complexity trade-off ( Figure 1a ) . 2 Related work . Vision Transformers There is a surge of research interests recently in exploiting Transformers for visual recognition tasks [ 36 , 35 , 41 , 31 , 44 ] , inspired by their remarkable success in NLP [ 33 , 10 , 4 ] . Core to these NLP and vision transformers is the same self-attention mechanism [ 33 ] that computes a self-attention matrix by exhaustively comparing token pairs . This means a quadratic complexity with the sequence length in both space and time , which thus limits the scalability of Transformers in dealing with long sequences . This limitation is more serious in vision than NLP : To process an image with at least thousands of pixels , patch-wise tokenization is a must for Transformers to control the computational cost . Given higher resolution images , the patch size also needs to be enlarged proportionally sacrificing the spatial resolution . This limits the capability of Transformers , e.g. , learning fine-grained feature representation as required in many visual recognition tasks . Linear Transformers Recently , there have been a number of linear/efficient variants [ 5 , 34 , 17 , 18 , 30 , 24 , 16 ] of Transformers in NLP . For example , [ 34 ] learns to shrink the length of Key and Value based on a low-rank assumption . [ 18 ] adopts a hashing strategy to selective the most similar pairs and only compute attention among them . [ 5 , 17 ] utilize different kernel functions for approximating softmax-based self-attention matrix . [ 24 ] applies random feature mapping on the sequences to approach the original softmax function . [ 16 ] decreases the time and memory consumption of the attention matrix by replacing the softmax function with its linear-complexity recurrent alternative . When applied to visual recognition tasks , however , we show that these models have considerable performance degradation compared to the standard Transformers [ 33 ] ( see Sec . 4.2 ) . The most related work to SOFT is [ 39 ] which uses the Nyström matrix decomposition to avoid computing the full attention matrix . However , this method suffers from several theoretical defects : ( 1 ) As the standard self-attention needs to apply row-wise softmax normalization on the full attention matrix , a direct application of matrix decomposition is infeasible . As a workaround , softmax is simply applied to all the ingredient matrices in [ 39 ] . Such an approximation is not guaranteed theoretically . ( 2 ) With a polynomial iteration method , it is not guaranteed that the generalized attention matrix inverse can be computed when the matrix is a nearly singular one in practice . In contrast to all the above methods , in this paper we propose a softmax-free self-attention mechanism that facilitates matrix decomposition for complexity minimization with theoretical guarantees . 3 Method . 3.1 Softmax-free self-attention formulation . A schematic illustration of our model is given in Figure 2 . Let ’ s first look at our attention module design . Given a sequence of n tokens X ∈ Rn×d with each token represented by a d-dimensional feature vector , self-attention [ 33 ] aims to discover the correlations of all token pairs exhaustively . Formally , X is first linearly projected into three de-dimensional spaces ( query , key , and values ) as : Q = XWq ∈ Rn×de , K = XWk ∈ Rn×de , V = XWv ∈ Rn×de , ( 1 ) where Wq , Wk , Wv ∈ Rd×de are learnable matrices . Self-attention can be expressed in a generic formulation as : yi , : = n∑ j=1 α ( Qi , : , Kj , : ) Vj , : , ( 2 ) where is the Hadamard product , and i , j ∈ { 1 , · · · , n } index the tokens . The key self-attention function α : Rde × Rde → R is composed of a nonlinear function β : R → R and a relation function γ : Rde × Rde → R. A dominant instantiation of α is the scaled dot-product based softmax self-attention [ 33 ] , defined as β ( · ) = softmax ( · ) , γ ( Qi , : , Kj , : ) = 1√ de ·Q > i , :Kj , : . ( 3 ) Whilst this softmax self-attention has been the de facto choice and seldomly questioned , as discussed earlier it is not necessarily suited for linearization . To facilitate the design of linear self-attention , we introduce a softmax-free self-attention function with the dot-product replaced by a Gaussian kernel as : β′ ( · ) = exp ( · ) , γ′ ( Qi , : , Kj , : ) = − 1 2 √ de · ‖Qi , : −Kj , :‖22 . ( 4 ) To preserve the symmetric property of attention matrix as in Eq ( 3 ) , we set the project matrices Wq and Wk in Eq ( 1 ) identical ( i.e. , Q = K ) . Our self-attention matrix is then written as : Si , j = exp ( − 1 2 √ de · ‖Qi , : −Kj , :‖22 ) . ( 5 ) For notation simplicity , we define the matrix formulation as : S = exp ( Q K ) . Remarks Our self-attention matrix S has three important properties : ( 1 ) It is symmetric ; ( 2 ) All the elements lie in a unit range of [ 0 , 1 ] ; ( 3 ) All diagonal elements hold the largest value 1 ( selfreinforced ) , with the bottom ones ( corresponding to most dissimilar token pairs ) being close to 0 . As Gaussian kernel is a positive definite kernel [ 12 ] , S is deemed a Gram matrix . However , we find that when using our kernel-based self-attention matrix S without linearization , the training of a transformer fails to converge . This might explain why softmax dot-product based self-attention [ 33 ] is so popular in vanilla transformers . 3.2 Low-rank regularization via matrix decomposition with linear complexity . To solve the convergence and quadratic complexity problems , we leverage matrix decomposition as a unified solution with low-rank regularization . In particular , we consider Nyström [ 38 ] , which is originally a low-rank matrix approximation algorithm . This enables our model ’ s complexity to be reduced significantly without computing the full self-attention matrix S. We make this choice because our S is positive semi-definite ( i.e. , a Gram matrix ) without follow-up normalization which are all necessary conditions for Nyström . In contrast , [ 39 ] totally ignores these requirements , leading to theoretical flaw in its approximation . To define the Nyström method formally , let us express S = exp ( Q K ) as a block matrix : S = [ A B B > C ] ∈ Rn×n , ( 6 ) where A ∈ Rm×m , B ∈ Rm× ( n−m ) , C ∈ R ( n−m ) × ( n−m ) with m n. Through Nyström decomposition ( see derivative details in Supplementary A.1 ) , an approximation can be represented as : Ŝ = [ A B > ] A† [ A B ] = P > A†P , where P = [ A B ] , ( 7 ) and A† is the Moore-Penrose ( a generalized ) inverse of A . Sampling In the standard Nyström formulation , A and B are sub-matrices of S obtained by randomly sampled m tokens , denoted as Q̃ . We call the sampled Q̃ as bottleneck tokens . However , Algorithm 1 : SOFT : Softmax-free attention Input : Q ∈ Rn×de , sampling function fs Sampling Q̃← fs ( Q ) ; A← exp ( Q̃ Q̃ ) , P ← exp ( Q̃ Q ) ; Ŝ ← P > NR ( A ) P ; Output : Ŝ Algorithm 2 : NR : Newton-Raphson iteration Input : A ∈ Rm×m , and T ∈ Z+ α = 2/‖A‖21.Initialize A0 ← αA ; for k from 1 to T do Ak ← 2Ak−1 −Ak−1AAk−1 end Output : AT we find empirically that random sampling is considerably sensitive to the choice of m. We hence explore two additional options by leveraging the structural prior of visual data : ( 1 ) Using one convolutional layer with kernel size k and stride k to learn Q̃ , and ( 2 ) Using average pooling with kernel size k and stride k to generate Q̃ . For both , we need to reshape Q to the form of RH×W×de . Each slide of convolution or pooling produces a token . We set k according to the length of Q such that m tokens can be obtained . Our experiments show that a convolution layer performs better in accuracy . We therefore use a convolution layer by default . As K is identical to Q , we have K̃ = Q̃ . Given these m tokens , we then compute A and P as : A = exp ( Q̃ K̃ ) , P = exp ( Q̃ K ) . ( 8 ) We finally obtain the regularized self-attention matrix Ŝ of SOFT as : Ŝ = exp ( Q K̃ ) ( exp ( Q̃ K̃ ) ) † exp ( Q̃ K ) , ( 9 ) leading to Algorithm 1 . The low-rank regularization is conducted as follows . For computing the attention score between any two tokens , we first correlate each of them with sampled tokens using our self-attention function ( Eq ( 5 ) ) ; With this correlation representation we then compute their similarity under the modulation of the generalized inverse of Q̃ ’ s correlation matrix . Similar as standard Nyström , our design associates the input tokens w.r.t . a small space spanned by sampled tokens , giving a proper estimation of the original attention relationships subject to a low-rank constraint . The correctness of this method is proved in Supplementary A.1 . Moore-Penrose inverse An accurate and commonly used way to calculate the Moore-Penrose inverse is to use Singular Value Decomposition ( SVD ) . Given A ∈ Rm×m and its SVD form A = UΣV > where U , V are m × m unitary matrices and Σ is a m × m diagonal matrix , the Moore-Penrose inverse of A is A† = V Σ†U > . Nevertheless , SVD is not friendly to the training process on GPU hence harming the model training efficiency . To solve this issue , we adopt the Newton–Raphson method . It is an iterative algorithm with the ( k + 1 ) -th iteration formulated given the previous iteration as : Ak+1 = 2Ak −AkAAk , and A0 = αA . ( 10 ) We now prove that Ak finally converges to Moore-Penrose inverse of Am×m , if α is sufficiently small [ 3 ] . Theorem 1 When α is sufficiently small , Ak+1 = 2Ak −AkAAk , Ak converges to A† . Though α = 2/‖A‖21 which ensures good convergence behavior in Algorithm 2 ( see more details in Supplementary A.2.1 ) , in practice , we find that using an alternative form gives more stable training and faster convergence . Specifically , in ‖I − A 2β n ‖A‖21 ‖1 ≤ 1 where β equals to 0.5 , we find the smallest ni that holds this inequality . Then , we initialize α as α = 2β ni ‖A‖21 . The following proposition comes with the proof of Theorem 1 : Proposition 1 ‖AAkA−A‖ and ‖Ak −A†‖ decreases to 0 monotonously , if α is sufficiently small . The detail of proposition 1 is shown in Supplementary A.2.2 . This ensures that our estimated inverse is sufficiently accurate for matrix decomposition , subject to that our SOFT attention is regularized . Complexity We summarize the complexity of SOFT in space and time . For time complexity , it involves : ( 1 ) Sampling : O ( nde ) . ( 2 ) Calculating three decomposed matrices : O ( nmde +mnde + m2de ) = O ( 2mnde + m2de ) ; ( 3 ) Moore-Penrose inverse : O ( T ×m3 ) = O ( Tm3 ) , where T is the iteration steps . ( 4 ) All matrix multiplication : O ( nm2 +mnde +mnde ) = O ( nm2 + 2mnde ) . The total time complexity is O ( ( de + 4mde + m2 ) n + Tm3 + dem2 ) . The space complexity is decided by four decomposed matrices with O ( n×m ) +O ( m×m ) +O ( m× n ) +O ( n× de ) = O ( ( 2m+ de ) n+m2 ) . As we keep m ( m n ) a fixed constant in our model , both time and space complexity are O ( n ) , making SOFT a linear self-attention .

Summary: The paper proposes a softmax-free self-attention that replaces the previous dot-product between the query matrix and the key matrix with a Gaussian kernel. To facilitate the proposed softmax-free self-attention, this paper further proposes to decompose the self-attention matrix with Singular Value Decomposition. The robustness of the proposed low-rank decomposition can be theoretically proved by computing the Moore-Penrose inverse with the Newton-Raphson method.