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 substitutes softmax self-attention with a Gaussian kernel function. This allows to approximate the self-attention matrix with a low-rank approximation. The robustness of the approximation is supported by theoretical guarantees, while previous works that don't provide these guarantees. The proposed method is evaluated on ImageNet, arguably the most popular benchmark for evaluating computer vision models, and shows impressive results when compared against other Linear Transformer alternatives and even when compared against the standard (full) self-attention approach.