cosFormer: Rethinking Softmax In Attention

1 INTRODUCTION . With years of development , the transformer model ( Vaswani et al. , 2017 ) and its variants ( Zaheer et al. , 2020 ; Wang et al. , 2020 ; Tay et al. , 2020a ) have been successfully adapted to three most popular artificial intelligence ( AI ) fields : i.e. , natural language processing ( Devlin et al. , 2019 ; Liu et al. , 2019 ) , computer vision ( Dosovitskiy et al. , 2020 ; Carion et al. , 2020 ; Liu et al. , 2021 ) and audio processing ( Schneider et al. , 2019 ; Baevski et al. , 2020 ) . Compared with conventional ∗Indicates the corresponding author . † Indicates equal contribution . recurrent ( Hochreiter & Schmidhuber , 1997 ) and convolutional architectures ( He et al. , 2016 ) , transformer-based architectures are generally more scalable to data volumes ( Brown et al. , 2020 ) and stronger in capturing global information with less inductive bias , thus excelling on many tasks . Dot-product attention with softmax normalization is the cornerstone of the transformer to capture long-range dependencies . However , its quadratic space and time complexity with regard to the length of the sequence make its computational overhead prohibitive , especially for long inputs . To address this issue , numerous methods are proposed recently , such as the sparse attention matrix ( Zaheer et al. , 2020 ; Beltagy et al. , 2020 ; Tay et al. , 2020a ; Kitaev et al. , 2019 ; Child et al. , 2019 ) , lowrank representations ( Wang et al. , 2020 ) or kernel-based methods ( Peng et al. , 2020 ; Choromanski et al. , 2020 ; Katharopoulos et al. , 2020 ) , among many others . These methods achieve reduced computational complexity with comparable performances when compared with the vanilla attention architecture on several selected tasks or corpus . However , the improved efficiency is usually achieved via introducing additional yet often impractical assumptions on the attention matrix ( Wang et al. , 2020 ) or with valid approximation of softmax operation only within constrained theoretical bounds ( Choromanski et al. , 2020 ; Peng et al. , 2020 ) Therefore , when their assumptions are unsatisfied or when approximation errors get accumulated , these methods may not always be advantageous over the vanilla architecture ( Narang et al. , 2021 ) . Consequently , performance deficiencies in a broad application spectrum are often observed in these transformer variants , especially those with linear complexity . For example , the Performer ( Choromanski et al. , 2020 ) , RFA ( Peng et al. , 2020 ) and Reformer ( Kitaev et al. , 2019 ) show less satisfactory performance on the GLUE benchmark ( Wang et al. , 2018 ) when compared with the vanilla architecture as suggested in our preliminary experiments ( Tab . 2 ) . Furthermore , many of these aforementioned methods are not applicable to casual attentions , which are critical for auto-regressive training . For example , techniques proposed in Linformer ( Wang et al. , 2020 ) and BigBird ( Zaheer et al. , 2020 ) are specific to cross attentions . Since the softmax operator appears to be the main hurdle while efficient yet accurate approximation to softmax is difficult to achieve , one question naturally arises : “ Can we replace the softmax operator with a linear function instead , while maintaining its key properties ? ” . By digging into the softmax attention , we find two key properties that affect its empirical performance : ( i ) elements in the attention matrix are non-negative ( Tsai et al. , 2019 ; Katharopoulos et al. , 2020 ) ; ( ii ) the non-linear re-weighting scheme acts as a stabilizer for the attention weights ( Titsias , 2016 ; Gao & Pavel , 2017 ; Jang et al. , 2016 ) . These findings reveal some new insights of the current approaches . For example , the linear transformer ( Katharopoulos et al. , 2020 ) achieves property ( i ) using an exponential linear unit ( Clevert et al. , 2016 ) activation function . However , due to lack of the re-weighting scheme , it underperforms other efficient transformer variants on the Long-Range Arena benchmark as shown in Figure 1 as well as the language modeling task ( Table 2 ) based on our controlled experiments . In this paper , we propose a new variant of linear transformer called COSFORMER that satisfies both of the above properties . Specifically , we enforce the non-negative property by passing the features to a ReLU ( Agarap , 2018 ) activation function before computing the similarity scores . In this way , we encourage the model to avoid aggregating negatively-correlated contextual information . Further , we adopt a cos re-weighting scheme to stabilize the attention weights . This helps the model to amplify local correlations , which usually contain more relevant information for natural language tasks . Thanks to the Ptolemy ’ s theorem , our attention can be exactly decomposed into a linear form . We perform extensive experiments on both autoregressive language models and bidirectional models on five public benchmarks , including WikiText-103 ( Merity et al. , 2017 ) , GLUE ( Wang et al. , 2018 ) , IMDB ( Maas et al. , 2011 ) , AMAZON ( Ni et al. , 2019 ) and Long-Range Arena benchmark ( Tay et al. , 2020b ) . Our model shows much better inference speed and smaller memory footprint , while achieving on par performance with the vanilla transformer . It is noteworthy that our method ranks 1st on the Long-Range Arena benchmark , showing favorable performance than other competitors , which well demonstrates its strong capacity in modeling long sequence inputs . 2 OUR METHOD . In this section , we provide technique details of our linear transformer called COSFORMER . The key insight of the COSFORMER is to replace the non-decomposable non-linear softmax operation by a linear operation with decomposable non-linear re-weighting mechanism . Our model is applicable to both casual and cross attentions with a linear time and space complexity with regard to the input sequence length , thus exhibiting strong capacity in modeling long-range dependencies . 2.1 THE GENERAL FORM OF TRANSFORMER . Given an input sequence x with length of N , we first represent it in the embedding space x ∈ RN×d with feature dimension of d. A transformer block T : RN×d → RN×d with input x is defined as : T ( x ) = F ( A ( x ) + x ) , ( 1 ) where F is a feedforward network that contains a residual connection ; A is the self-attention function that computes the attention matrix A , which has quadratic space and time complexity with respect to N , thus becoming the computation bottleneck of T on long inputs . There are three key components in A , namely , query ( Q ) , key ( K ) , value ( V ) computed through three learnable linear matrices WQ , WK , WV : Q = xWQ , K = xWK , V = xWV . We use Mi to represent the i-th row of a matrix M , then the output O ∈ RN×d of A ( x ) can be computed as : O = A ( x ) = [ O1 , . . . , ON ] T , Oi = ∑ j S ( Qi , Kj ) ∑ j S ( Qi , Kj ) Vj , ( 2 ) where S ( · ) measures the similarity between queries . If S ( Q , K ) = exp ( QKT ) , the Eq . 2 becomes the dot-product attention with softmax normalization . In this case , the space and time complexity to compute one row of the output Oi is O ( N ) . Therefore , the total space and time complexity for computing O grows quadratically with respect to the input length . 2.2 LINEARIZATION OF SELF-ATTENTION . According to Eq . 2 , we can select any similarity functions to compute the attention matrix . In order to maintain a linear computation budget , one solution is to adopt a decomposable similarity function such that : S ( Qi , Kj ) = φ ( Qi ) φ ( Kj ) T , ( 3 ) where φ is a kernel function that maps the queries and keys to their hidden representations . Then one can rewrite Eq . 2 in the form of kernel functions as : Oi = ∑N j=1 ( φ ( Qi ) φ ( Kj ) T ) Vj∑N j=1 ( φ ( Qi ) φ ( Kj ) T ) . , ( 4 ) After that , attention operation in linear complexity is achieved via the matrix product property : ( φ ( Q ) φ ( K ) T ) V = φ ( Q ) ( φ ( K ) TV ) . ( 5 ) In this form ( Eq . 5 ) , instead of explicitly computing the attention matrix A = QKT ∈ RN×N , we calculate the φ ( K ) TV ∈ Rd×d first , and then multiplying φ ( Q ) ∈ RN×d . By using this trick , we only incurs a computation complexity of O ( Nd2 ) . Note that in typical natural language tasks , the feature dimension of one head d is always much smaller than the input sequence lengthN ( d N ) , so we can safely omit d and achieve computation complexity of O ( N ) , as illustrated in Figure 2 . Previous Solutions As aforementioned , the key to the linear attentions is to find a decomposable similarity function S ( · ) that generalizes well to different tasks . Most existing linear transformers are trying to find an unbiased estimation of the softmax attention . For example , RFA ( Peng et al. , 2020 ) approximates the softmax operation with random feature maps using theorem of random fourier features ( Rahimi & Recht , 2008 ) and the Performer ( Choromanski et al. , 2020 ) utilizes positive random features to approximate it . However , we empirically find that these methods are sensitive to the selection of sampling rate and becomes unstable if the sampling rate gets too high . Also , to accommodate recency bias , gating mechanisms are employed to better exploit more recent context . Another group of works attempt to directly replace the softmax with a linear operation . For example , the linear transformer ( Katharopoulos et al. , 2020 ) model replaces the softmax similarity function with a pure dot product S = QKT , and use a non-linear activation function φ ( · ) = elu ( · ) + 1 to model the pairwise relation between features . However , our controlled experiments show that their solution does not necessarily generalize well on many downstream tasks ( Tab . 2 ) or the Long-Range Arena benchmark ( Tab . 4 ) . In this paper , we propose a new replacement of softmax that not only achieves comparable or better performance than the softmax attention in a wide range of tasks , but also enjoys linear space and time complexity . d x d Q ’ K ’ T V O ( Nd2 ) O ( Nd2 + Nd2 ) ≈ O ( N ) O ( Nd2 ) N x N Q KT V d x NN x d O ( N2d ) O ( N2d + N2d ) ≈O ( N2 ) Softmax O ( N2d ) Vanilla self attention Linearized self attention N x d N x d N x d d x N N x d N x d Figure 2 : Illustration of the computations for vanilla self attention ( left ) and linearized attention ( right ) . The input length is N and feature dimension is d , with d N . Tensors in the same box are associated for computation . The linearized formulation allows O ( N ) time and space complexity .

Summary: This work introduces a simple and elegant way to incorporate recency bias for kernelized linear attention mechanism using cos-based re-weighting mechanism.  This cosine re-weighting can be effectively implemented via query and key matrix transformations for linear attention mechanisms. The authors experiment with LRA, and both language modeling with autoregressive and bidirectional setups to show the efficacy of proposed method.