Random Feature Attention

1 INTRODUCTION . Transformer architectures ( Vaswani et al. , 2017 ) have achieved tremendous success on a variety of sequence modeling tasks ( Ott et al. , 2018 ; Radford et al. , 2018 ; Parmar et al. , 2018 ; Devlin et al. , 2019 ; Parisotto et al. , 2020 , inter alia ) . Under the hood , the key component is attention ( Bahdanau et al. , 2015 ) , which models pairwise interactions of the inputs , regardless of their distances from each other . This comes with quadratic time and memory costs , making the transformers computationally expensive , especially for long sequences . A large body of research has been devoted to improving their time and memory efficiency ( Tay et al. , 2020c ) . Although better asymptotic complexity and prominent gains for long sequences have been achieved ( Lee et al. , 2019 ; Child et al. , 2019 ; Beltagy et al. , 2020 , inter alia ) , in practice , many existing approaches are less well-suited for moderatelength ones : the additional computation steps required by some approaches can overshadow the time and memory they save ( Kitaev et al. , 2020 ; Wang et al. , 2020 ; Roy et al. , 2020 , inter alia ) . This work proposes random feature attention ( RFA ) , an efficient attention variant that scales linearly in sequence length in terms of time and space , and achieves practical gains for both long and moderate length sequences . RFA builds on a kernel perspective of softmax ( Rawat et al. , 2019 ) . Using the well-established random feature maps ( Rahimi & Recht , 2007 ; Avron et al. , 2016 ; §2 ) , RFA approximates the dot-then-exponentiate function with a kernel trick ( Hofmann et al. , 2008 ) : exp ( x · y ) ≈ φ ( x ) · φ ( y ) . Inspired by its connections to gated recurrent neural networks ( Hochreiter & Schmidhuber , 1997 ; Cho et al. , 2014 ) and fast weights ( Schmidhuber , 1992 ) , we further augment RFA with an optional gating mechanism , offering a straightforward way of learning with recency bias when locality is desired . ∗The majority of this work was done while these authors were at DeepMind . RFA and its gated variant ( §3 ) can be used as a drop-in substitute for the canonical softmax attention , and increase the number of parameters by less than 0.1 % . We explore its applications in transformers on language modeling , machine translation , and long text classification ( §4 ) . Our experiments show that RFA achieves comparable performance to vanilla transformer baselines in all tasks , while outperforming a recent related approach ( Katharopoulos et al. , 2020 ) . The gating mechanism proves particularly useful in language modeling : the gated variant of RFA outperforms the transformer baseline on WikiText-103 . RFA shines in decoding , even for shorter sequences . In our head-to-head comparison on machine translation benchmarks , RFA decodes around 2× faster than a transformer baseline , without accuracy loss . Comparisons to several recent efficient transformer variants on three long text classification datasets show that RFA is competitive in terms of both accuracy and efficiency . Our analysis ( §5 ) shows that more significant time and memory efficiency improvements can be achieved for longer sequences : 12× decoding speedup with less than 10 % of the memory for 2,048-length outputs . 2 BACKGROUND . 2.1 ATTENTION IN SEQUENCE MODELING . The attention mechanism ( Bahdanau et al. , 2015 ) has been widely used in many sequence modeling tasks . Its dot-product variant is the key building block for the state-of-the-art transformer architectures ( Vaswani et al. , 2017 ) . Let { qt } Nt=1 denote a sequence of N query vectors , that attend to sequences of M key and value vectors . At each timestep , the attention linearly combines the values weighted by the outputs of a softmax : attn ( qt , { ki } , { vi } ) = ∑ i exp ( qt · ki/τ ) ∑ j exp ( qt · kj/τ ) v > i . ( 1 ) τ is the temperature hyperparameter determining how “ flat ” the softmax is ( Hinton et al. , 2015 ) .1 Calculating attention for a single query takes O ( M ) time and space . For the full sequence of N queries the space amounts to O ( MN ) . When the computation can not be parallelized across the queries , e.g. , in autoregressive decoding , the time complexity is quadratic in the sequence length . 2.2 RANDOM FEATURE METHODS . The theoretical backbone of this work is the unbiased estimation of the Gaussian kernel by Rahimi & Recht ( 2007 ) . Based on Bochner ’ s theorem ( Bochner , 1955 ) , Rahimi & Recht ( 2007 ) proposed random Fourier features to approximate a desired shift-invariant kernel . The method nonlinearly transforms a pair of vectors x and y using a random feature map φ ; the inner product between φ ( x ) and φ ( y ) approximates the kernel evaluation on x and y . More precisely : Theorem 1 ( Rahimi & Recht , 2007 ) . Let φ : Rd → R2D be a nonlinear transformation : φ ( x ) = √ 1/D [ sin ( w1 · x ) , . . . , sin ( wD · x ) , cos ( w1 · x ) , . . . , cos ( wD · x ) ] > . ( 2 ) When d-dimensional random vectors wi are independently sampled from N ( 0 , σ2Id ) , Ewi [ φ ( x ) · φ ( y ) ] = exp ( −‖x− y‖2 /2σ2 ) . ( 3 ) Variance of the estimation is inversely proportional to D ( Appendix A.2 ; Yu et al. , 2016 ) . Random feature methods proved successful in speeding up kernel methods ( Oliva et al. , 2015 ; Avron et al. , 2017 ; Sun , 2019 , inter alia ) , and more recently are used to efficiently approximate softmax ( Rawat et al. , 2019 ) . In §3.1 , we use it to derive an unbiased estimate to exp ( 〈· , ·〉 ) and then an efficient approximation to softmax attention . 3 MODEL . This section presents RFA ( §3.1 ) and its gated variant ( §3.2 ) . In §3.3 we lay out several design choices and relate RFA to prior works . We close by practically analyzing RFA ’ s complexity ( §3.4 ) . 1M = N in self-attention ; they may differ , e.g. , in the cross attention of a sequence-to-sequence model . 3.1 RANDOM FEATURE ATTENTION . RFA builds on an unbiased estimate to exp ( 〈· , ·〉 ) from Theorem 1 , which we begin with : exp ( x · y/σ2 ) = exp ( ‖x‖2 /2σ2 + ‖y‖2 /2σ2 ) exp ( −‖x− y‖2 /2σ2 ) ≈ exp ( ‖x‖2 /2σ2 + ‖y‖2 /2σ2 ) φ ( x ) · φ ( y ) . ( 4 ) The last line does not have any nonlinear interaction between φ ( x ) and φ ( y ) , allowing for a linear time/space approximation to attention . For clarity we assume the query and keys are unit vectors.2 attn ( qt , { ki } , { vi } ) = ∑ i exp ( qt · ki/σ2 ) ∑ j exp ( qt · kj/σ2 ) v > i ≈ ∑ i φ ( qt ) > φ ( ki ) v > i∑ j φ ( qt ) · φ ( kj ) = φ ( qt ) > ∑ i φ ( ki ) ⊗ vi φ ( qt ) · ∑ j φ ( kj ) = RFA ( qt , { ki } , { vi } ) . ( 5 ) ⊗ denotes the outer product between vectors , and σ2 corresponds to the temperature term τ in Eq . 1 . RFA can be used as a drop-in-replacement for softmax-attention . ( a ) The input is revealed in full to cross attention and encoder self-attention . Here RFA calculates attention using Eq . 5 . ( b ) In causal attention RFA attends only to the prefix.3 This allows for a recurrent computation . Tuple ( St ∈ R2D×d , zt ∈ R2D ) is used as the “ hidden state ” at time step t to keep track of the history , similar to those in RNNs . Then RFA ( qt , { ki } i≤t , { vi } i≤t ) = φ ( qt ) > St/ ( φ ( qt ) · zt ) , where St = St−1 + φ ( kt ) ⊗ vt , zt = zt−1 + φ ( kt ) . ( 6 ) 2D denotes the size of φ ( · ) . Appendix A.1 summarizes the computation procedure of RFA , and Figure 1 compares it against the softmax attention . Appendix A.3 derives causal RFA in detail . Analogously to the softmax attention , RFA has its multiheaded variant ( Vaswani et al. , 2017 ) . In our experiments we use causal RFA in a transformer language model ( §4.1 ) , and both cross and causal RFA in the decoder of a sequence-to-sequence machine translation model . 3.2 RFA-GATE : LEARNING WITH RECENCY BIAS . The canonical softmax attention does not have any explicit modeling of distance or locality . In learning problems where such inductive bias is crucial ( Ba et al. , 2016 ; Parmar et al. , 2018 ; Miconi et al. , 2018 ; Li et al. , 2019 , inter alia ) , transformers heavily rely on positional encodings . Answering to this , many approaches have been proposed , e.g. , learning the attention spans ( Sukhbaatar et al. , 2This can be achieved by ` 2-normalizing the query and keys . See §3.3 for a related discussion . 3It is also sometimes called “ decoder self-attention ” or “ autoregressive attention. ” 2019 ; Wu et al. , 2020 ) , and enhancing the attention computation with recurrent ( Hao et al. , 2019 ; Chen et al. , 2019 ) or convolutional ( Wu et al. , 2019 ; Mohamed et al. , 2019 ) components . RFA faces the same issue , but its causal attention variant ( Eq . 6 ) offers a straightforward way of learning with recency bias . We draw inspiration from its connections to RNNs , and augment RFA with a learned gating mechanism ( Hochreiter & Schmidhuber , 1997 ; Cho et al. , 2014 ; Peng et al. , 2018 , inter alia ) : gt = sigmoid ( wg · xt + bg ) , St = gt St−1 + ( 1− gt ) φ ( kt ) ⊗ vt , zt = gt zt−1 + ( 1− gt ) φ ( kt ) . ( 7 ) wg and bg are learned parameters , and xt is the input representation at timestep t.4 By multiplying the learned scalar gates 0 < gt < 1 against the hidden state ( St , zt ) , history is exponentially decayed , favoring more recent context . The gating mechanism shows another benefit of RFA : it would be otherwise more difficult to build similar techniques into the softmax attention , where there is no clear sense of “ recurrence ” ( Appendix A.5 ) . It proves useful in our language modeling experiments ( §4.1 ) .

Summary: The paper presents a linear time and space attention mechanism based on random features to approximate the softmax. The paper is clearly written and easy to follow. The results are convincing: not chasing SOTA, but comparing to sensible baselines, namely [Baevski & Auli 2019] for language modeling on Wikitext-103, and [Vaswani et al. 2017] for machine translation on WMT14 EN-DE/EN-FR and IWSLT14 DE-EN.