Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention β€” part 1/3

Cleaned-text code points [0, 15260). Bibliography omitted here; source.md includes it.

# Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

###### Abstract

Transformers achieve remarkable performance in several tasks but due to their quadratic complexity, with respect to the input’s length, they are prohibitively slow for very long sequences. To address this limitation, we express the self-attention as a linear dot-product of kernel feature maps and make use of the associativity property of matrix products to reduce the complexity from π’ͺ ⁑ ( N 2 ) \mathcal{O}\left(N^{2}\right) to π’ͺ ⁑ ( N ) \mathcal{O}\left(N\right) , where N N is the sequence length. We show that this formulation permits an iterative implementation that dramatically accelerates autoregressive transformers and reveals their relationship to recurrent neural networks. Our linear transformers achieve similar performance to vanilla transformers and they are up to 4000x faster on autoregressive prediction of very long sequences.

## 1 Introduction

Transformer models were originally introduced by Vaswani et al. (2017) in the context of neural machine translation ( Sutskever et al., 2014 ; Bahdanau et al., 2015 ) and have demonstrated impressive results on a variety of tasks dealing with natural language ( Devlin et al., 2019 ) , audio ( Sperber et al., 2018 ) , and images ( Parmar et al., 2019 ) . Apart from tasks with ample supervision, transformers are also effective in transferring knowledge to tasks with limited or no supervision when they are pretrained with autoregressive ( Radford et al., 2018 ; Radford et al., 2019 ) or masked language modeling objectives ( Devlin et al., 2019 ; Yang et al., 2019 ; Song et al., 2019 ; Liu et al., 2020 ) .

However, these benefits often come with a very high computational and memory cost. The bottleneck is mainly caused by the global receptive field of self-attention, which processes contexts of N N inputs with a quadratic memory and time complexity π’ͺ ⁑ ( N 2 ) \mathcal{O}\left(N^{2}\right) . As a result, in practice transformers are slow to train and their context is limited . This disrupts temporal coherence and hinders the capturing of long-term dependencies. Dai et al. (2019) addressed the latter by attending to memories from previous contexts albeit at the expense of computational efficiency.

Lately, researchers shifted their attention to approaches that increase the context length without sacrificing efficiency. Towards this end, Child et al. (2019) introduced sparse factorizations of the attention matrix to reduce the self-attention complexity to π’ͺ ⁑ ( N ​ N ) \mathcal{O}\left(N\sqrt{N}\right) . Kitaev et al. (2020) further reduced the complexity to π’ͺ ⁑ ( N ​ log ⁑ N ) \mathcal{O}\left(N\log N\right) using locality-sensitive hashing. This made scaling to long sequences possible. Even though the aforementioned models can be efficiently trained on large sequences, they do not speed-up autoregressive inference.

In this paper, we introduce the linear transformer model that significantly reduces the memory footprint and scales linearly with respect to the context length. We achieve this by using a kernel-based formulation of self-attention and the associative property of matrix products to calculate the self-attention weights (Β§ 3.2 ). Using our linear formulation, we also express causal masking with linear complexity and constant memory (Β§ 3.3 ). This reveals the relation between transformers and RNNs, which enables us to perform autoregressive inference orders of magnitude faster (Β§ 3.4 ).

Our evaluation on image generation and automatic speech recognition demonstrates that linear transformer can reach the performance levels of transformer, while being up to three orders of magnitude faster during inference.

## 2 Related Work

In this section, we provide an overview of the most relevant works that seek to address the large memory and computational requirements of transformers. Furthermore, we discuss methods that theoretically analyze the core component of the transformer model, namely self-attention. Finally, we present another line of work that seeks to alleviate the softmax bottleneck in the attention computation.

### 2.1 Efficient Transformers

Existing works seek to improve memory efficiency in transformers through weight pruning ( Michel et al., 2019 ) , weight factorization ( Lan et al., 2020 ) , weight quantization ( Zafrir et al., 2019 ) or knowledge distillation. Clark et al. (2020) proposed a new pretraining objective called replaced token detection that is more sample efficient and reduces the overall computation. Lample et al. (2019) used product-key attention to increase the capacity of any layer with negligible computational overhead.

Reducing the memory or computational requirements with these methods leads to training or inference time speedups, but, fundamentally, the time complexity is still quadratic with respect to the sequence length which hinders scaling to long sequences. In contrast, we show that our method reduces both memory and time complexity of transformers both theoretically (Β§ 3.2 ) and empirically (Β§ 4.1 ).

Another line of research aims at increasing the β€œcontext” of self-attention in transformers. Context refers to the maximum part of the sequence that is used for computing self-attention. Dai et al. (2019) introduced Transformer-XL which achieves state-of-the-art in language modeling by learning dependencies beyond a fixed length context without disrupting the temporal coherence. However, maintaining previous contexts in memory introduces significant additional computational cost. In contrast, Sukhbaatar et al. (2019) extended the context length significantly by learning the optimal attention span per attention head, while maintaining control over the memory footprint and computation time. Note that both approaches have the same asymptotic complexity as the vanilla model. In contrast, we improve the asymptotic complexity of the self-attention, which allows us to use significantly larger context.

More related to our model are the works of Child et al. (2019) and Kitaev et al. (2020) . The former ( Child et al., 2019 ) introduced sparse factorizations of the attention matrix reducing the overall complexity from quadratic to π’ͺ ⁑ ( N ​ N ) \mathcal{O}\left(N\sqrt{N}\right) for generative modeling of long sequences. More recently, Kitaev et al. (2020) proposed Reformer. This method further reduces complexity to π’ͺ ⁑ ( N ​ log ⁑ N ) \mathcal{O}\left(N\log{N}\right) by using locality-sensitive hashing (LSH) to perform fewer dot products. Note that in order to be able to use LSH, Reformer constrains the keys, for the attention, to be identical to the queries. As a result this method cannot be used for decoding tasks where the keys need to be different from the queries. In comparison, linear transformers impose no constraints on the queries and keys and scale linearly with respect to the sequence length. Furthermore, they can be used to perform inference in autoregressive tasks three orders of magnitude faster, achieving comparable performance in terms of validation perplexity.

### 2.2 Understanding Self-Attention

There have been few efforts to better understand self-attention from a theoretical perspective. Tsai et al. (2019) proposed a kernel-based formulation of attention in transformers which considers attention as applying a kernel smoother over the inputs with the kernel scores being the similarity between inputs. This formulation provides a better way to understand attention components and integrate the positional embedding. In contrast, we use the kernel formulation to speed up the calculation of self-attention and lower its computational complexity. Also, we observe that if a kernel with positive similarity scores is applied on the queries and keys, linear attention converges normally.

More recently, Cordonnier et al. (2020) provided theoretical proofs and empirical evidence that a multi-head self-attention with sufficient number of heads can express any convolutional layer. Here, we instead show that a self-attention layer trained with an autoregressive objective can be seen as a recurrent neural network and this observation can be used to significantly speed up inference time of autoregressive transformer models.

### 2.3 Linearized softmax

For many years, softmax has been the bottleneck for training classification models with a large number of categories ( Goodman, 2001 ; Morin & Bengio, 2005 ; Mnih & Hinton, 2009 ) . Recent works ( Blanc & Rendle, 2017 ; Rawat et al., 2019 ) , have approximated softmax with a linear dot product of feature maps to speed up the training through sampling. Inspired from these works, we linearize the softmax attention in transformers. Concurrently with this work, Shen et al. (2020) explored the use of linearized attention for the task of object detection in images. In comparison, we do not only linearize the attention computation, but also develop an autoregressive transformer model with linear complexity and constant memory for both inference and training. Moreover, we show that through the lens of kernels, every transformer can be seen as a recurrent neural network.

## 3 Linear Transformers

In this section, we formalize our proposed linear transformer . We present that changing the attention from the traditional softmax attention to a feature map based dot product attention results in better time and memory complexity as well as a causal model that can perform sequence generation in linear time, similar to a recurrent neural network.

Initially, in Β§ 3.1 , we introduce a formulation for the transformer architecture introduced in ( Vaswani et al., 2017 ) . Subsequently, in Β§ 3.2 and Β§ 3.3 we present our proposed linear transformer and finally, in Β§ 3.4 we rewrite the transformer as a recurrent neural network.

### 3.1 Transformers

Let x ∈ ℝ N Γ— F x\in\mathbb{R}^{N\times F} denote a sequence of N N feature vectors of dimensions F F . A transformer is a function T : ℝ N Γ— F β†’ ℝ N Γ— F T:\mathbb{R}^{N\times F}\to\mathbb{R}^{N\times F} defined by the composition of L L transformer layers T 1 ​ ( β‹… ) , … , T L ​ ( β‹… ) T_{1}(\cdot),\dots,T_{L}(\cdot) as follows, T l ​ ( x ) = f l ​ ( A l ​ ( x ) + x ) . T_{l}(x)=f_{l}(A_{l}(x)+x). (1) The function f l ​ ( β‹… ) f_{l}(\cdot) transforms each feature independently of the others and is usually implemented with a small two-layer feedforward network. A l ​ ( β‹… ) A_{l}(\cdot) is the self attention function and is the only part of the transformer that acts across sequences.

The self attention function A l ​ ( β‹… ) A_{l}(\cdot) computes, for every position, a weighted average of the feature representations of all other positions with a weight proportional to a similarity score between the representations. Formally, the input sequence x x is projected by three matrices W Q ∈ ℝ F Γ— D W_{Q}\in\mathbb{R}^{F\times D} , W K ∈ ℝ F Γ— D W_{K}\in\mathbb{R}^{F\times D} and W V ∈ ℝ F Γ— M W_{V}\in\mathbb{R}^{F\times M} to corresponding representations Q Q , K K and V V . The output for all positions, A l ​ ( x ) = V β€² A_{l}(x)=V^{\prime} , is computed as follows, Q \displaystyle Q = x ​ W Q , \displaystyle=xW_{Q}, (2) K \displaystyle K = x ​ W K , \displaystyle=xW_{K}, V \displaystyle V = x ​ W V , \displaystyle=xW_{V}, A l ​ ( x ) = V β€² \displaystyle A_{l}(x)=V^{\prime} = softmax ​ ( Q ​ K T D ) ​ V . \displaystyle=\text{softmax}\left(\frac{QK^{T}}{\sqrt{D}}\right)V. Note that in the previous equation, the softmax function is applied rowwise to Q ​ K T QK^{T} . Following common terminology, the Q Q , K K and V V are referred to as the β€œqueries”, β€œkeys” and β€œvalues” respectively.

Equation 2 implements a specific form of self-attention called softmax attention where the similarity score is the exponential of the dot product between a query and a key. Given that subscripting a matrix with i i returns the i i -th row as a vector, we can write a generalized attention equation for any similarity function as follows, V i β€² = βˆ‘ j = 1 N sim ​ ( Q i , K j ) ​ V j βˆ‘ j = 1 N sim ​ ( Q i , K j ) . V^{\prime}_{i}=\frac{\sum_{j=1}^{N}\text{sim}\left(Q_{i},K_{j}\right)V_{j}}{\sum_{j=1}^{N}\text{sim}\left(Q_{i},K_{j}\right)}. (3) Equation 3 is equivalent to equation 2 if we substitute the similarity function with sim ​ ( q , k ) = exp ⁑ ( q T ​ k D ) \text{sim}\left(q,k\right)=\exp\left(\frac{q^{T}k}{\sqrt{D}}\right) .

### 3.2 Linearized Attention

The definition of attention in equation 2 is generic and can be used to define several other attention implementations such as polynomial attention or RBF kernel attention ( Tsai et al., 2019 ) . Note that the only constraint we need to impose to sim ​ ( β‹… ) \text{sim}\left(\cdot\right) , in order for equation 3 to define an attention function, is to be non-negative. This includes all kernels k ⁑ ( x , y ) : ℝ 2 Γ— F β†’ ℝ + k(x,y):\mathbb{R}^{2\times F}\to\mathbb{R}_{+} .

Given such a kernel with a feature representation Ο• ⁑ ( x ) \phi\left(x\right) we can rewrite equation 2 as follows, V i β€² = βˆ‘ j = 1 N Ο• ​ ( Q i ) T ​ Ο• ​ ( K j ) ​ V j βˆ‘ j = 1 N Ο• ​ ( Q i ) T ​ Ο• ​ ( K j ) , \displaystyle V^{\prime}_{i}=\frac{\sum_{j=1}^{N}\phi\left(Q_{i}\right)^{T}\phi\left(K_{j}\right)V_{j}}{\sum_{j=1}^{N}\phi\left(Q_{i}\right)^{T}\phi\left(K_{j}\right)}, (4) and then further simplify it by making use of the associative property of matrix multiplication to V i β€² = Ο• ​ ( Q i ) T ​ βˆ‘ j = 1 N Ο• ⁑ ( K j ) ​ V j T Ο• ​ ( Q i ) T ​ βˆ‘ j = 1 N Ο• ⁑ ( K j ) . \displaystyle V^{\prime}_{i}=\frac{\phi\left(Q_{i}\right)^{T}\sum_{j=1}^{N}\phi\left(K_{j}\right)V_{j}^{T}}{\phi\left(Q_{i}\right)^{T}\sum_{j=1}^{N}\phi\left(K_{j}\right)}. (5) The above equation is simpler to follow when the numerator is written in vectorized form as follows, ( Ο• ⁑ ( Q ) ​ Ο• ​ ( K ) T ) ​ V = Ο• ⁑ ( Q ) ​ ( Ο• ​ ( K ) T ​ V ) . \left(\phi\left(Q\right)\phi\left(K\right)^{T}\right)V=\phi\left(Q\right)\left(\phi\left(K\right)^{T}V\right). (6) Note that the feature map Ο• ⁑ ( β‹… ) \phi\left(\cdot\right) is applied rowwise to the matrices Q Q and K K .

From equation 2 , it is evident that the computational cost of softmax attention scales with π’ͺ ⁑ ( N 2 ) \mathcal{O}\left(N^{2}\right) , where N N represents the sequence length. The same is true for the memory requirements because the full attention matrix must be stored to compute the gradients with respect to the queries, keys and values. In contrast, our proposed linear transformer from equation 5 has time and memory complexity π’ͺ ⁑ ( N ) \mathcal{O}\left(N\right) because we can compute βˆ‘ j = 1 N Ο• ⁑ ( K j ) ​ V j T \sum_{j=1}^{N}\phi\left(K_{j}\right)V_{j}^{T} and βˆ‘ j = 1 N Ο• ⁑ ( K j ) \sum_{j=1}^{N}\phi\left(K_{j}\right) once and reuse them for every query.

#### 3.2.1 Feature Maps and Computational Cost

For softmax attention, the total cost in terms of multiplications and additions scales as π’ͺ ⁑ ( N 2 ​ max ⁑ ( D , M ) ) \mathcal{O}\left(N^{2}\max\left(D,M\right)\right) , where D D is the dimensionality of the queries and keys and M M is the dimensionality of the values. On the contrary, for linear attention, we first compute the feature maps of dimensionality C C . Subsequently, computing the new values requires π’ͺ ⁑ ( N ​ C ​ M ) \mathcal{O}\left(NCM\right) additions and multiplications.