Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention — part 2/3

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

The previous analysis does not take into account the choice of kernel and feature function. Note that the feature function that corresponds to the exponential kernel is infinite dimensional, which makes the linearization of exact softmax attention infeasible. On the other hand, the polynomial kernel, for example, has an exact finite dimensional feature map and has been shown to work equally well with the exponential or RBF kernel ( Tsai et al., 2019 ) . The computational cost for a linearized polynomial transformer of degree 2 is 𝒪 ⁡ ( N ​ D 2 ​ M ) \mathcal{O}\left(ND^{2}M\right) . This makes the computational complexity favorable when N > D 2 N>D^{2} . Note that this is true in practice since we want to be able to process sequences with tens of thousands of elements.

For our experiments, that deal with smaller sequences, we employ a feature map that results in a positive similarity function as defined below, ϕ ​ ( x ) = elu ​ ( x ) + 1 , \phi\left(x\right)=\text{elu}(x)+1, (7) where elu ​ ( ⋅ ) \text{elu}(\cdot) denotes the exponential linear unit ( Clevert et al., 2015 ) activation function. We prefer elu ​ ( ⋅ ) \text{elu}(\cdot) over relu ​ ( ⋅ ) \text{relu}(\cdot) to avoid setting the gradients to 0 when x x is negative. This feature map results in an attention function that requires 𝒪 ⁡ ( N ​ D ​ M ) \mathcal{O}\left(NDM\right) multiplications and additions. In our experimental section, we show that the feature map of equation 7 performs on par to the full transformer, while significantly reducing the computational and memory requirements.

### 3.3 Causal Masking

The transformer architecture can be used to efficiently train autoregressive models by masking the attention computation such that the i i -th position can only be influenced by a position j j if and only if j ≤ i j\leq i , namely a position cannot be influenced by the subsequent positions. Formally, this causal masking changes equation 3 as follows, V i ′ = ∑ j = 1 i sim ​ ( Q i , K j ) ​ V j ∑ j = 1 i sim ​ ( Q i , K j ) . V^{\prime}_{i}=\frac{\sum_{j=1}^{i}\text{sim}\left(Q_{i},K_{j}\right)V_{j}}{\sum_{j=1}^{i}\text{sim}\left(Q_{i},K_{j}\right)}. (8)

Following the reasoning of § 3.2 , we linearize the masked attention as described below, V i ′ = ϕ ​ ( Q i ) T ​ ∑ j = 1 i ϕ ⁡ ( K j ) ​ V j T ϕ ​ ( Q i ) T ​ ∑ j = 1 i ϕ ⁡ ( K j ) . V^{\prime}_{i}=\frac{\phi\left(Q_{i}\right)^{T}\sum_{j=1}^{i}\phi\left(K_{j}\right)V_{j}^{T}}{\phi\left(Q_{i}\right)^{T}\sum_{j=1}^{i}\phi\left(K_{j}\right)}. (9) By introducing S i S_{i} and Z i Z_{i} as follows, S i \displaystyle S_{i} = ∑ j = 1 i ϕ ⁡ ( K j ) ​ V j T , \displaystyle=\sum_{j=1}^{i}\phi\left(K_{j}\right)V_{j}^{T}, (10) Z i \displaystyle Z_{i} = ∑ j = 1 i ϕ ⁡ ( K j ) , \displaystyle=\sum_{j=1}^{i}\phi\left(K_{j}\right), (11) we can simplify equation 9 to V i ′ = ϕ ​ ( Q i ) T ​ S i ϕ ​ ( Q i ) T ​ Z i . V^{\prime}_{i}=\frac{\phi\left(Q_{i}\right)^{T}S_{i}}{\phi\left(Q_{i}\right)^{T}Z_{i}}. (12) Note that, S i S_{i} and Z i Z_{i} can be computed from S i − 1 S_{i-1} and Z i − 1 Z_{i-1} in constant time hence making the computational complexity of linear transformers with causal masking linear with respect to the sequence length.

#### 3.3.1 Gradient Computation

A naive implementation of equation 12 , in any deep learning framework, requires storing all intermediate values S i S_{i} in order to compute the gradients. This increases the memory consumption by max ⁡ ( D , M ) \max\left(D,M\right) times; thus hindering the applicability of causal linear attention to longer sequences or deeper models. To address this, we derive the gradients of the numerator in equation 9 as cumulative sums. This allows us to compute both the forward and backward pass of causal linear attention in linear time and constant memory . A detailed derivation is provided in the supplementary material.

Given the numerator V ¯ i \bar{V}_{i} and the gradient of a scalar loss function with respect to the numerator ∇ V ¯ i ℒ \nabla_{\bar{V}_{i}}\mathcal{L} , we derive ∇ ϕ ⁡ ( Q i ) ℒ \nabla_{\phi\left(Q_{i}\right)}\mathcal{L} , ∇ ϕ ⁡ ( K i ) ℒ \nabla_{\phi\left(K_{i}\right)}\mathcal{L} and ∇ V i ℒ \nabla_{V_{i}}\mathcal{L} as follows, ∇ ϕ ⁡ ( Q i ) ℒ \displaystyle\nabla_{\phi\left(Q_{i}\right)}\mathcal{L} = ∇ V ¯ i ℒ ​ ( ∑ j = 1 i ϕ ⁡ ( K j ) ​ V j T ) T , \displaystyle=\nabla_{\bar{V}_{i}}\mathcal{L}\left(\sum_{j=1}^{i}\phi\left(K_{j}\right)V_{j}^{T}\right)^{T}, (13) ∇ ϕ ⁡ ( K i ) ℒ \displaystyle\nabla_{\phi\left(K_{i}\right)}\mathcal{L} = ( ∑ j = i N ϕ ⁡ ( Q j ) ​ ( ∇ V ¯ j ℒ ) T ) ​ V i , \displaystyle=\left(\sum_{j=i}^{N}\phi\left(Q_{j}\right)\left(\nabla_{\bar{V}_{j}}\mathcal{L}\right)^{T}\right)V_{i}\ , (14) ∇ V i ℒ \displaystyle\nabla_{V_{i}}\mathcal{L} = ( ∑ j = i N ϕ ⁡ ( Q j ) ​ ( ∇ V ¯ j ℒ ) T ) T ​ ϕ ​ ( K i ) . \displaystyle=\left(\sum_{j=i}^{N}\phi\left(Q_{j}\right)\left(\nabla_{\bar{V}_{j}}\mathcal{L}\right)^{T}\right)^{T}\phi\left(K_{i}\right). (15) The cumulative sum terms in equations 9 , 13 - 15 are computed in linear time and require constant memory with respect to the sequence length. This results in an algorithm with computational complexity 𝒪 ⁡ ( N ​ C ​ M ) \mathcal{O}\left(NCM\right) and memory 𝒪 ⁡ ( N ​ max ⁡ ( C , M ) ) \mathcal{O}\left(N\max\left(C,M\right)\right) for a given feature map of C C dimensions. A pseudocode implementation of the forward and backward pass of the numerator is given in algorithm 1 .

#### 3.3.2 Training and Inference

When training an autoregressive transformer model the full ground truth sequence is available. This makes layerwise parallelism possible both for f l ​ ( ⋅ ) f_{l}(\cdot) of equation 1 and the attention computation. As a result, transformers are more efficient to train than recurrent neural networks. On the other hand, during inference the output for timestep i i is the input for timestep i + 1 i+1 . This makes autoregressive models impossible to parallelize. Moreover, the cost per timestep for transformers is not constant; instead, it scales with the square of the current sequence length because attention must be computed for all previous timesteps.

Our proposed linear transformer model combines the best of both worlds . When it comes to training, the computations can be parallelized and take full advantage of GPUs or other accelerators. When it comes to inference, the cost per time and memory for one prediction is constant for our model. This means we can simply store the ϕ ⁡ ( K j ) ​ V j T \phi\left(K_{j}\right)V_{j}^{T} matrix as an internal state and update it at every time step like a recurrent neural network. This results in inference thousands of times faster than other transformer models.

### 3.4 Transformers are RNNs

In literature, transformer models are considered to be a fundamentally different approach to recurrent neural networks. However, from the causal masking formulation in § 3.3 and the discussion in the previous section, it becomes evident that any transformer layer with causal masking can be written as a model that, given an input, modifies an internal state and then predicts an output, namely a Recurrent Neural Network (RNN). Note that, in contrast to Universal Transformers ( Dehghani et al., 2018 ) , we consider the recurrence with respect to time and not depth.

In the following equations, we formalize the transformer layer of equation 1 as a recurrent neural network. The resulting RNN has two hidden states, namely the attention memory s s and the normalizer memory z z . We use subscripts to denote the timestep in the recurrence. s 0 \displaystyle s_{0} = 0 , \displaystyle=0, (16) z 0 \displaystyle z_{0} = 0 , \displaystyle=0, (17) s i \displaystyle s_{i} = s i − 1 + ϕ ⁡ ( x i ​ W K ) ​ ( x i ​ W V ) T , \displaystyle=s_{i-1}+\phi\left(x_{i}W_{K}\right)\left(x_{i}W_{V}\right)^{T}, (18) z i \displaystyle z_{i} = z i − 1 + ϕ ⁡ ( x i ​ W K ) , \displaystyle=z_{i-1}+\phi\left(x_{i}W_{K}\right), (19) y i \displaystyle y_{i} = f l ​ ( ϕ ​ ( x i ​ W Q ) T ​ s i ϕ ​ ( x i ​ W Q ) T ​ z i + x i ) . \displaystyle=f_{l}\left(\frac{\phi\left(x_{i}W_{Q}\right)^{T}s_{i}}{\phi\left(x_{i}W_{Q}\right)^{T}z_{i}}+x_{i}\right). (20) In the above equations, x i x_{i} denotes the i i -th input and y i y_{i} the i i -th output for a specific transformer layer. Note that our formulation does not impose any constraint on the feature function and it can be used for representing any transformer model, in theory even the ones using softmax attention. This formulation is a first step towards better understanding the relationship between transformers and popular recurrent networks ( Hochreiter & Schmidhuber, 1997 ) and the processes used for storing and retrieving information.

## 4 Experiments

function forward( ϕ ⁡ ( Q ) \phi\left(Q\right) , ϕ ⁡ ( K ) \phi\left(K\right) , V V ) :

In this section, we analyze experimentally the performance of the proposed linear transformer . Initially, in § 4.1 , we evaluate the linearized attention in terms of computational cost, memory consumption and convergence on synthetic data. To further showcase the effectiveness of linear transformers , we evaluate our model on two real-world applications, image generation in § 4.2 and automatic speech recognition in § 4.3 . We show that our model achieves competitive performance with respect to the state-of-the-art transformer architectures, while requiring significantly less GPU memory and computation.

Throughout our experiments, we compare our model with two baselines, the full transformer with softmax attention and the Reformer ( Kitaev et al., 2020 ) , the latter being a state-of-the-art accelerated transformer architecture. For the Reformer, we use a PyTorch reimplementation of the published code and for the full transformer we use the default PyTorch implementation. Note that for Reformer, we do not use the reversible layers, however, this does not affect the results as we only measure the memory consumption with respect to the self attention layer. In all experiments, we use softmax ( Vaswani et al., 2017 ) to refer to the standard transformer architecture, linear for our proposed linear transformers and lsh-X for Reformer ( Kitaev et al., 2020 ) , where X denotes the hashing rounds.

For training the linear transformers , we use the feature map of equation 7 . Our PyTorch ( Paszke et al., 2019 ) code with documentation and examples can be found at https://linear-transformers.com/ . The constant memory gradient computation of equations 13 - 15 is implemented in approximately 200 lines of CUDA code.

### 4.1 Synthetic Tasks

#### 4.1.1 Convergence Analysis

To examine the convergence properties of linear transformers we train on an artifical copy task with causal masking. Namely, the transformers have to copy a series of symbols similar to the sequence duplication task of Kitaev et al. (2020) . We use a sequence of maximum length 128 with 10 different symbols separated by a dedicated separator symbol. For all three methods, we train a 4 layer transformer with 8 attention heads using a batch size of 64 and the RAdam optimizer ( Liu et al., 2019 ) with a learning rate of 10 − 3 10^{-3} which is reduced to 10 − 4 10^{-4} after 3000 updates. Figure 2 depicts the loss with respect to the number of gradient steps. We observe that linear converges smoothly and reaches a lower loss than lsh due to the lack of noise introduced by hashing. In particular, it reaches the same loss as softmax.

#### 4.1.2 Memory and Computational Requirements

In this subsection, we compare transformers with respect to their computational and memory requirements. We compute the attention and the gradients for a synthetic input with varying sequence lengths N ∈ { 2 9 , 2 10 , … , 2 16 } N\in\{2^{9},2^{10},\dots,2^{16}\} and measure the peak allocated GPU memory and required time for each variation of transformer. We scale the batch size inversely with the sequence length and report the time and memory per sample in the batch.

Every method is evaluated up to the maximum sequence length that fits the GPU memory. For this benchmark we use an NVidia GTX 1080 Ti with 11GB of memory. This results in a maximum sequence length of 4,096 elements for softmax and 16,384 for lsh-4 and lsh-8. As expected, softmax scales quadratically with respect to the sequence length. Our method is faster and requires less memory than the baselines for every configuration, as seen in figure 1 . We observe that both Reformer and linear attention scale linearly with the sequence length. Note that although the asymptotic complexity for Reformer is 𝒪 ⁡ ( N ​ log ⁡ N ) \mathcal{O}\left(N\log N\right) , log ⁡ N \log N is small enough and does not affect the computation time.

### 4.2 Image Generation

Transformers have shown great results on the task of conditional or unconditional autoregressive generation ( Radford et al., 2019 ; Child et al., 2019 ) , however, sampling from transformers is slow due to the task being inherently sequential and the memory scaling with the square of the sequence length. In this section, we train causally masked transformers to predict images pixel by pixel. Our achieved performance in terms of bits per dimension is on par with softmax attention while being able to generate images more than 1,000 times faster and with constant memory per image from the first to the last pixel. We refer the reader to our supplementary for comparisons in terms of training evolution, quality of generated images and time to generate a single image. In addition, we also compare with a faster softmax transformer that caches the keys and values during inference, in contrast to the PyTorch implementation.

#### 4.2.1 MNIST

First, we evaluate our model on image generation with autoregressive transformers on the widely used MNIST dataset ( LeCun et al., 2010 ) . The architecture for this experiment comprises 8 attention layers with 8 attention heads each. We set the embedding size to 256 which is 32 dimensions per head. Our feed forward dimensions are 4 times larger than our embedding size. We model the output with a mixture of 10 logistics as introduced by Salimans et al. (2017) . We use the RAdam optimizer with a learning rate of 10 − 4 10^{-4} and train all models for 250 epochs. For the reformer baseline, we use 1 and 4 hashing rounds. Furthermore, as suggested in Kitaev et al. (2020) , we use 64 buckets and chunks with approximately 32 elements. In particular, we divide the 783 long input sequence to 27 chunks of 29 elements each. Since the sequence length is realtively small, namely only 784 pixels, to remove differences due to different batch sizes we use a batch size of 10 for all methods.

Table 1 summarizes the results. We observe that linear transformers achieve almost the same performance, in terms of final perplexity, as softmax transformers while being able to generate images more than 300 times faster. This is achieved due to the low memory requirements of our model, which is able to simultaneously generate 10,000 MNIST images with a single GPU. In particular, the memory is constant with respect to the sequence length because the only thing that needs to be stored between pixels are the s i s_{i} and z i z_{i} values as described in equations 18 and 19 . On the other hand, both softmax and Reformer require memory that increases with the length of the sequence.

Image completions and unconditional samples from our MNIST model can be seen in figure 3 . We observe that our linear transformer generates very convincing samples with sharp boundaries and no noise. In the case of image completion, we also observe that the transformer learns to use the same stroke style and width as the original image effectively attending over long temporal distances. Note that as the achieved perplexity is more or less the same for all models, we do not observe qualitative differences between the generated samples from different models.

#### 4.2.2 CIFAR-10