An Attention Free Transformer

We introduce Attention Free Transformer ( AFT ) , an efficient variant of Transformers ( Vaswani et al. , 2017 ) that eliminates the need for dot product attention . AFT offers great simplicity and efficiency compared with standard Transformers , where the multi-head attention operation is replaced with the composition of element-wise multiplications/divisions and global/local pooling . During training time , AFT has linear time and space complexity w.r.t . both the sequence length and feature dimension ; in the autoregressive decoding mode , AFT has constant memory and time complexity per step . We show that , surprisingly , we are able to train AFT effectively on challenging benchmarks , and also to match or surpass the standard Transformer counterparts and other efficient variants . In particular , AFT achieves the state-of-the-art result on CIFAR10 autoregressive modeling with much reduced complexity , and also outperforms several efficient Transformer variants on Enwik8 . 1 INTRODUCTION . Attention mechanisms , represented by Transformers ( Vaswani et al. , 2017 ) , have driven the advancement of various machine learning problems , including language modeling ( Devlin et al. , 2018 ; Radford et al . ) , image modeling ( Chen et al . ) , and set modeling ( Lee et al. , 2019 ) . Different from other well known model architectures such as Convolutional Neural Nets ( CNNs ) or Recurrent Neural Nets ( RNNs ) , Transformers enable direct interaction between every pair of elements within a sequence , which makes them especially powerful at capturing long term dependencies . However , Transformers require high computational costs . The root cause of this challenge is the need to perform attention operations that have quadratic time and space complexity w.r.t the context size . This makes it especially difficult for Transformers to scale to inputs with large context sizes . A number of recent works have been dedicated to addressing the scalability issue of Transformers ( Child et al. , 2019 ; Kitaev et al. , 2020 ; Rae et al. , 2020 ; Wang et al. , 2020b ; Katharopoulos et al. , 2020 ; Tay et al. , 2020a ; Choromanski et al. , 2020 ) . While the techniques adopted in the literature range from sparsity , locality sensitive hashing , low rank decomposition , kernel approximation and etc. , most of them are trying to approximate the full attention operation . In this paper , we take a bolder step towards the same goal , by proposing a computational module that does not use or approximate the standard dot product attention . We hence name our model the attention free transformer ( AFT ) . Similar to dot product attention , AFT is composed of the interaction of three quantities , namely the query , key and value . What ’ s different , however , is that AFT operates solely based on element-wise operations . To be more concrete , they key and value are first multiplied element-wise , the result of which is then pooled over the context dimension ( in the causal model , this corresponds to a cumulative sum ) . The query is then multiplied with the reduced key-value representation element-wise to produce the final output . See Figure 1a for an illustration . AFT maintains the full advantage of dot product attention , namely direct interaction between any two elements in a sequence ( up to proper masking ) . However , the computational cost is drastically reduced to a O ( Td ) complexity for time and space , where T , d are the context length and feature dimension , respectively . In the autoregressive decoding mode , AFT also provides constant decoding time and space complexity per step , compared to O ( T ) for standard transformers . To the best of our knowledge , AFT is the first model that achieves such efficiency in the context of Transformers . See Table 1 for the complexity analysis of AFT in comparison to other variants . We show that we can interpret AFT as an extreme case of multi head dot product attention ( MHA ) . In particular , we show that by 1 ) setting the number of heads equal to the feature dimension in MHA and 2 ) using relu in place of softmax as the non-linearity , MHA can be decomposed into the summation of two AFT modules ( see Equation 6 ) . However , this relationship is not true in a general sense , i.e. , by varying the non-linearity injected after the query and key in AFT , we can obtain models that do not have a MHA counterpart . This realization allows us to freely explore the design choices ( e.g. , nonlinearity ) of AFT to achieve the best performance . This philosophy is in direct contrast with previous and concurrent “ linearized attention ” works ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2020 ) , which are constrained by the design space of MHA . We perform experiments with AFT on several benchmarks , including unconditional image modeling , image super-resolution , language modeling , machine translation and point cloud generation . We show that AFT works very well as an alternative to the standard Transformer , providing competitive results as well as excellent efficiency . To summarize , our contributions are as follows : • We propose AFT , a new family of Transformer models that achievesO ( Td ) time and space complexity in training , as well asO ( d ) time and space complexity in autoregressive decoding . • We show strong performance of AFT as a drop in replacement of MHA on various benchmarks , including setting the state-of-the-art result on CIFAR10 in the standard setting and outperforming other efficient Transformer variants . 2 MULTI-HEAD ATTENTION . At the core of Transformers is the Multi-Head Attention ( MHA ) operation . Given three sequences , namely the query Q ∈ RT×d , key K ∈ RT×d and value V ∈ RT×d , and the number of heads h , MHA performs a scaled dot product attention for each head i , defined as : fi ( Q , K , V ) = σ ( Q′i ( K ′ i ) T √ dk ) V ′i , s.t . Q ′ i = QW Q i , K ′ i = KW K i , V ′ i = VW V i , ( 1 ) where WQi ∈ Rd×dk , WKi ∈ Rd×dk , WVi ∈ Rd×dvare linear transformations for head i , and σ is the non-linearity by default set as the softmaxr function ( subscript r indicates softmax is applied to each row of a matrix ) . dk , dv are dimensions for key and value , respectively . MHA concatenates the output of h attention heads along the channel dimension , resulting in feature dimension hdv . Unless otherwise mentioned , we assume dk = dv and h = ddk . This means the query , key and value are the same dimension within each head , and output dimension matches that of the input . 3 METHODOLOGY . 3.1 ATTENTION FREE TRANSFORMER . We now define Attention free transformer ( AFT ) , which provides an alternative to MHA . Given Q , K , V , AFT first linearly transforms them into Q′ = QWQ , K ′ = KWK , V ′ = VWV , then performs following operation : f ( Q , K , V ) = σq ( Q ′ ) T∑ t=1 ( σk ( K ′ ) V ′ ) t , ( 2 ) where is the element-wise product , with support for broadcasting when the operands ’ dimensions don ’ t exactly match 1 ; σq , σk are nonlinearities applied to the query and key , respectively . Explained in words , the key and value are first combined with an element-wise multiplication , the result of which is then pooled over the context dimension , yielding a fixed length context vector ∈ Rd . This context vector is then multiplied with each row of the query , which forms the final output of an AFT layer . One particularly useful variant of MHA is masked attention , oftentimes presented in the form of causal attention . Specifically , in auto-regressive models , queries are constrained to not be able to interact with keys and values beyond the curret position . In standard attention , this is usually implemented with an explicit binary masking matrix of shape T×T , with non-causal entries masked as 0 . We show that it is also straightforward to extend AFT to the causal mode while maintaining its efficiency . We denote an AFT layer ’ s output as Yt = f ( Q≤t , K≤t , V≤t ) , t = 1 , ... , T 2 . We formulate the casual AFT as : Yt = σq ( Q ′ t ) t∑ t′=1 ( σk ( K ′ ≤t ) V ′≤t ) t′ , t = 1 , ... , T , ( 3 ) where the subscript Xt indexes the tth row of matrix X . Discussions : The design philosophy of AFT is to promote extreme efficiency , while keeping the benefit of standard Transformers . Concretely , AFT enables direct interaction of any two elements within the sequence , which is arguably the biggest advantage of Transformers over other types of models such as RNNs and ConvNets . However , AFT gets rid of the need of performing the costly spatial dot product attention , by computing a reduced value representation with the weights only depending on the keys . The resulting operation has an extremely efficiency of O ( Td ) w.r.t . both time and space , which is the first model that achieves linear complexity along both context and feature dimensions . Moreover , the causal mode of AFT has an additional advantage of a constant decoding cost per step , similar to ( Katharopoulos et al. , 2020 ) . To see this , from Equation 3 , we have a simple recursion of Yt = σq ( Q′t ) ( σk ( K ′t ) V ′t +KVt−1 ) with KVt = ∑t t′=1 ( σk ( K ′ t′ ) V ′t′ ) , assuming σq , σk are both element-wise functions . One thus only need to keep KVt in memory , and update it with constant cost per step . Selecting nonlinearies : σq , σk provide additional nonlinearity which helps to increase model ’ s capacity . Empirically , we have found that one particularly strong setting is to let σk = softmax which is normalized along the context dimension . This choice brings an interesting benefit especially in the causal mode , which we can explicitly write as : Yt = σq ( Q ′ t ) ( gt ( t ) V ′t + t−1∑ t′=1 gt ( t ′ ) V ′t′ ) , gt ( t ′ ) = exp ( K ′t ) ∑t t′=1 exp ( K ′ t′ ) . ( 4 ) Here gt ( t ) acts as a role similar to that of an input gate in an LSTM , and gt ( t′ ) is operating like the forget gate , which depends on the input of time t , dynamically downweights the contribution of past time steps . When augmented with standard position embeddings as commonly used in Transformers , this allows the model to be able to learn the notion of recency at the same time of having access to the full context in the history . From this view , σq can also be interpreted as the output gate , for which we found that both sigmoid and relu work well , with the former being slightly better . Also note that the same space and time complexity still holds for σk = softmax , both in training and decoding . In our experiments , unless otherwise mentioned , we use the sigmoid+ softmax setting for σq and σk by default . Relation to MHA : Although AFT performs fundamentally different operations than standard attention , we show that the two family of models overlap in the extreme case . To see this , we explore 1We adopt Numpy styled broadcasting convention : https : //numpy.org/doc/stable/user/theory.broadcasting.html 2We assume here that Yt includes input information at the current position t , the version where the current position is excluded can be obtained by shifting the outputs to the right . the limit of number of heads in MHA , which amounts to letting dk = 1 for each head . In this case , the dot product operation within each head reduces to a scalar product . Next , we set σ to be relu instead of softmax in Equation 1 . In this case , we have : fi ( Q , K , V ) = [ Q ′ i ( K ′ i ) T ] +V ′ i = ( [ Q′i ] + [ ( K ′ i ) T ] + + [ −Q′i ] + [ − ( K ′i ) T ] + ) V ′i = [ Q′i ] + ( [ K ′i ] T +V ′ i ) + [ −Q′i ] + ( [ −KTi ] +V ′i ) , ( 5 ) where [ · ] + denotes the relu operator , and Q′i , K ′i , V ′i ∈ RT×1 by definition . The concatenated output of the attention heads can then be concisely written as : f ( Q , K , V ) = [ Q′ ] + T∑ t=1 ( [ K ′ ] + V ′ ) t + [ −Q′ ] + T∑ t=1 ( [ −K ′ ] + V ′ ) t , ( 6 ) which consists of two terms , each of which is an AFT operation , with σq = σk = [ · ] + and σq = σk = [ −· ] + , respectively . However , note that this correspondence is not general , i.e. , AFT does not need to approximate any MHA counterpart and can indeed have very different inductive biases than that of a standard Transformer . Relation to Linearized Attention : There are a few recent works proposing to linearize the dot product attention ( Linear Attention ) from the view of kernel approximation , first proposed in Katharopoulos et al . ( 2020 ) and also in concurrent work ( Choromanski et al. , 2020 ) . ( Katharopoulos et al. , 2020 ) proposes the linear attention operation in the form : Yt = φ ( Q′t ) ∑T t′=1 ( φ ( K ′t′ ) TV ′t′ ) φ ( Q ) ′t ∑t t′−1 φ ( K ′ t ) T , ( 7 ) where Q′t , K ′ t , V ′ t are all row vectors of R d. Equation 7 is similar to AFT , in the sense that the key and value are first combined and reduced in both cases . However AFT differs in two aspects : 1 ) the time complexity of Linear Attention is O ( Td2 ) , which is linear in the sequence length but has difficulty scaling to wide networks 2 ) Linear Attention is designed to approximate MHA , where the nonlinearity on query and key are shared . In AFT however , we show that it is beneficial to search for different nonlinearities for both the query and key .

Summary: The paper introduces a method to replace qkv-attention by a simpler, efficient building block. This is done by element-wise multiplication of a query representation with a compressed kv-memory. Per-channel attention pooling is used to compress the kv-memory. The model is derived from a softmax-free version of self-attention. The results show good performance on a couple of standard image and language modeling tasks while occasionally exhibiting favorable training speed. The results are largely on par with transformer baselines.