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

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

The benefits of our linear formulation increase as the sequence length increases. To showcase that, we train 16 layer transformers to generate CIFAR-10 images ( Krizhevsky et al., 2009 ) . For each layer we use the same configuration as in the previous experiment. For Reformer, we use again 64 buckets and 83 chunks of 37 elements, which is approximately 32, as suggested in the paper. Since the sequence length is almost 4 times larger than for the previous experiment, the full transformer can only be used with a batch size of 1 in the largest GPU that is available to us, namely an NVidia P40 with 24GB of memory. For both the linear transformer and reformer, we use a batch size of 4. All models are trained for 7 days. We report results in terms of bits per dimension and image generation throughput in table 2 . Note that although the main point of this experiment is not the final perplexity, it is evident that as the sequence length grows, the fast transformer models become increasingly more efficient per GPU hour, achieving better scores than their slower counterparts.

As the memory and time to generate a single pixel scales quadratically with the number of pixels for both Reformer and softmax attention, the increase in throughput for our linear transformer is even more pronounced. In particular, for every image generated by the softmax transformer, our method can generate 4,460 images . Image completions and unconditional samples from our model can be seen in figure 4 . We observe that our model generates images with spatial consistency and can complete images convincigly without significantly hindering the recognition of the image category. For instance, in figure 4(b) , all images have successfully completed the dog’s nose (first row) or the windshield of the truck (last row).

Unconditional samples

Image completion

Unconditional samples

Image completion

### 4.3 Automatic Speech Recognition

To show that our method can also be used for non-autoregressive tasks, we evaluate the performance of linear transformers in end-to-end automatic speech recognition using Connectionist Temporal Classification (CTC) loss ( Graves et al., 2006 ) . In this setup, we predict a distribution over phonemes for each input frame in a non autoregressive fashion. We use the 80 hour WSJ dataset ( Paul & Baker, 1992 ) with 40-dimensional mel-scale filterbanks without temporal differences as features. The dataset contains sequences with 800 frames on average and a maximum sequence length of 2,400 frames. For this task, we also compare with a bidirectional LSTM ( Hochreiter & Schmidhuber, 1997 ) with 3 layers of hidden size 320. We use the Adam optimizer ( Kingma & Ba, 2014 ) with a learning rate of 10 − 3 10^{-3} which is reduced when the validation error stops decreasing. For the transformer models, we use 9 layers with 6 heads with the same embedding dimensions as for the image experiments. As an optimizer, we use RAdam with an initial learning rate of 10 − 4 10^{-4} that is divided by 2 when the validation error stops decreasing.

All models are evaluated in terms of phoneme error rate (PER) and training time per epoch. We observe that linear outperforms the recurrent network baseline and Reformer both in terms of performance and speed by a large margin, as seen in table 3 . Note that the softmax transformer, achieves lower phone error rate in comparison to all baselines, but is significantly slower. In particular, linear transformer is more than 3 × 3\times faster per epoch. We provide training evolution plots in the supplementary.

## 5 Conclusions

In this work, we presented linear transformer , a model that significantly reduces the memory and computational cost of the original transformers. In particular, by exploiting the associativity property of matrix products we are able to compute the self-attention in time and memory that scales linearly with respect to the sequence length. We show that our model can be used with causal masking and still retain its linear asymptotic complexities. Finally, we express the transformer model as a recurrent neural network, which allows us to perform inference on autoregressive tasks thousands of time faster.

This property opens a multitude of directions for future research regarding the storage and retrieval of information in both RNNs and transformers. Another line of research to be explored is related to the choice of feature map for linear attention. For instance, approximating the RBF kernel with random Fourier features could allow us to use models pretrained with softmax attention.

## Acknowledgements

Angelos Katharopoulos was supported by the Swiss National Science Foundation under grant numbers FNS-30209 ”ISUL” and FNS-30224 ”CORTI”. Apoorv Vyas was supported by the Swiss National Science Foundation under grant number FNS-30213 ”SHISSM”. Nikolaos Pappas was supported by the Swiss National Science Foundation under grant number P400P2_183911 ”UNISON”.

## Appendix A Gradient Derivation

In the first section of our supplementary material, we derive in detail the gradients for causally masked linear transformers and show that they can be computed in linear time and constant memory. In particular, we derive the gradients of a scalar loss with respect to the numerator of the following equation, 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)}. (21) The gradient with respect to the denominator and the fraction are efficiently handled by autograd. Without loss of generality, we can assume that Q Q and K K already contain the vectors mapped by ϕ ⁡ ( ⋅ ) \phi\left(\cdot\right) , hence given the numerator V ¯ i = Q i T ​ ∑ j = 1 i K j ​ V j T , \bar{V}_{i}=Q_{i}^{T}\sum_{j=1}^{i}K_{j}V_{j}^{T}, (22) and ∇ V ¯ ℒ \nabla_{\bar{V}}\mathcal{L} we seek to compute ∇ Q ℒ \nabla_{Q}\mathcal{L} , ∇ K ℒ \nabla_{K}\mathcal{L} and ∇ V ℒ \nabla_{V}\mathcal{L} . Note that Q ∈ ℝ N × D Q\in\mathbb{R}^{N\times D} , K ∈ ℝ N × D K\in\mathbb{R}^{N\times D} and V ∈ ℝ N × M V\in\mathbb{R}^{N\times M} . To derive the gradients, we first express the above equation for a single element without using vector notation, V ¯ i ​ e = ∑ d = 1 D Q i ​ d ​ ∑ j = 1 i K j ​ d ​ V j ​ e = ∑ d = 1 D ∑ j = 1 i Q i ​ d ​ K j ​ d ​ V j ​ e . \bar{V}_{ie}=\sum_{d=1}^{D}Q_{id}\sum_{j=1}^{i}K_{jd}V_{je}=\sum_{d=1}^{D}\sum_{j=1}^{i}Q_{id}K_{jd}V_{je}. (23) Subsequently we can start deriving the gradients for Q Q by taking the partial derivative for any Q l ​ t Q_{lt} , as follows ∂ ℒ ∂ Q l ​ t = ∑ e = 1 M ∂ ℒ ∂ V ¯ l ​ e ​ ∂ V ¯ l ​ e ∂ Q l ​ t = ∑ e = 1 M ∂ ℒ ∂ V ¯ l ​ e ​ ( ∑ j = 1 l K j ​ t ​ V j ​ e ) . \frac{\partial\mathcal{L}}{\partial Q_{lt}}=\sum_{e=1}^{M}\frac{\partial\mathcal{L}}{\partial\bar{V}_{le}}\frac{\partial\bar{V}_{le}}{\partial Q_{lt}}=\sum_{e=1}^{M}\frac{\partial\mathcal{L}}{\partial\bar{V}_{le}}\left(\sum_{j=1}^{l}K_{jt}V_{je}\right). (24) If we write the above equation as a matrix product of gradients it becomes, ∇ Q i ℒ = ∇ V ¯ i ℒ ​ ( ∑ j = 1 i K j ​ V j T ) T , \nabla_{Q_{i}}\mathcal{L}=\nabla_{\bar{V}_{i}}\mathcal{L}\left(\sum_{j=1}^{i}K_{j}V_{j}^{T}\right)^{T}, (25) proving equation 13 from the main paper. In equation 24 we made use of the fact that Q l ​ t Q_{lt} only affects V ¯ l \bar{V}_{l} hence we do not need to sum over i i to compute the gradients. However, for K K and V V this is not the case. In particular, K j K_{j} affects all V ¯ i \bar{V}_{i} where i ≥ j i\geq j . Consequently, we can write the partial derivative of the loss with respect to K l ​ t K_{lt} as follows, ∂ ℒ ∂ K l ​ t \displaystyle\frac{\partial\mathcal{L}}{\partial K_{lt}} = ∑ e = 1 M ∑ i = l N ∂ ℒ ∂ V ¯ i ​ e ​ ∂ V ¯ i ​ e ∂ K l ​ t = ∑ e = 1 M ∑ i = l N ∂ ℒ ∂ V ¯ i ​ e ​ ∂ ( ∑ d = 1 D ∑ j = 1 i Q i ​ d ​ K j ​ d ​ V j ​ e ) ∂ K l ​ t \displaystyle=\sum_{e=1}^{M}\sum_{i=l}^{N}\frac{\partial\mathcal{L}}{\partial\bar{V}_{ie}}\frac{\partial\bar{V}_{ie}}{\partial K_{lt}}=\sum_{e=1}^{M}\sum_{i=l}^{N}\frac{\partial\mathcal{L}}{\partial\bar{V}_{ie}}\frac{\partial\left(\sum_{d=1}^{D}\sum_{j=1}^{i}Q_{id}K_{jd}V_{je}\right)}{\partial K_{lt}} (26) = ∑ e = 1 M ∑ i = l N ∂ ℒ ∂ V ¯ i ​ e ​ Q i ​ t ​ V l ​ e . \displaystyle=\sum_{e=1}^{M}\sum_{i=l}^{N}\frac{\partial\mathcal{L}}{\partial\bar{V}_{ie}}Q_{it}V_{le}. As for Q Q we can now write the gradient in vectorized form, ∇ K i ℒ = ( ∑ j = i N Q j ​ ( ∇ V ¯ j ℒ ) T ) ​ V i , \nabla_{K_{i}}\mathcal{L}=\left(\sum_{j=i}^{N}Q_{j}\left(\nabla_{\bar{V}_{j}}\mathcal{L}\right)^{T}\right)V_{i}, (27) proving equation 14 from the paper. Following the same reasoning, we can compute the partial derivative of the loss with respect to V l ​ t V_{lt} and prove equation 15. Note that the cumulative sum matrices for the gradient with respect to Q Q and K K have the same size, however one is computed in the forward direction (summing from 1 to N N ) similarly to the forward pass and the other is computed in the backwards direction (summing from N N to 1) similar to backpropagation through time done in RNNs.

## Appendix B Training Evolution

In figure 5 we present the training evolution of all transformer models in our experiments. For the MNIST experiment (Fig. 5(a) ) we train all methods for 250 epochs. The sequence length is small enough so that the training time does not vary significantly for all methods. We observe that our method converges on par with softmax attention outperforming significantly both reformer variants.

On the other hand, for CIFAR-10 (Fig. 5(b) ) we train all methods for a fixed amount of time, namely 7 days. We observe that lsh-1 and linear complete significantly more epochs than softmax and lsh-4 and achieve better performance. This gap is expected to increase with a further increase in sequence length.

Finally, in our last experiment on automatic speech recognition (Fig. 5(c) ), softmax outperforms significantly both Reformer and linear in terms of convergence. Note that linear is 3 × 3\times faster per epoch which means it has completed approximately 4 times more epochs in comparison to softmax. Even though softmax attention is better in this task, we observe that linear transformers significantly outperform Reformer both in terms of convergence and final performance.

## Appendix C Image Generation Throughput Discussion

### C.1 Stateful softmax attention

In § 4.2 of the main paper, we report the image generation throughput and we compare with softmax transformer and lsh . In this section we create another baseline, denoted as stateful-softmax , that implements a softmax autoregressive transformer as a recurrent model. Namely, all the keys and values are saved and then passed to the model again when predicting the next element of the sequence. The state of this recurrent model is the set of keys and values which has size proportional to the sequence length. This is qualitatively different to our proposed model that has a state with fixed dimensions and computing the i i -th state given the previous one has fixed computational cost regardless of i i .

Table 4 summarizes the results. We observe that stateful-softmax is significantly faster than vanilla transformers. However, its complexity is still quadratic with respect to the sequence length and our forumlation is more than 50 × \times faster for CIFAR-10. Moreover, we would like to point out that implementing a similar stateful attention for Reformer is not a trivial task as the sorting and chunking operations need to be performed each time a new input is provided.

### C.2 Equalizing the batch size

In the previous sections we evaluate the throughput of all transformer variants for the task of autoregressive image generation. However, another important factor to consider is latency, namely the total time required to produce a single image. To this end, we use a batch size of 1 and measure the time required by all methods to generate a single image. In addition to running the inference on the GPU, we also evaluate the time required on CPU. The results are reported in table 5 .

We observe that all methods underutilize the GPU and achieve significantly smaller image generation throughput than the one shown in table 4 . The proposed linear transformer is faster than all the methods and in particular it is almost 6.6 × \times faster than softmax transformers for generating an image on CIFAR-10. Note that our linear autoregressive transformer is the only method that is faster on the CPU than on the GPU in every case. This is due to the fact that computing the attention as an RNN has such a low cost that the main computational bottleneck becomes the inevitable outer loop over the sequence.

## Appendix D Qualitative Results on Image Generation

In this section we provide qualitative results for our image generation experiments. Since the perplexity of all models is approximately the same, as expected, the qualitative differences are not significant. A rather interesting observation however is that the Reformer models provide significantly fewer variations in their unconditional samples. Moreover, we observe that image completion is a significantly easier task than unconditional generation as all models perform significantly better.

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .