Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

• Self-attention’s global receptive field causes quadratic time and memory costs, slowing training and limiting context and long-term dependency capture.Source: §1 · p. 1
• Sparse attention and locality-sensitive hashing support efficient processing of long sequences but do not speed up autoregressive inference.Source: §1 · p. 1
• Kernel-based self-attention uses matrix associativity for linear scaling, and causal masking is expressed with linear complexity and constant memory.Source: §1 · p. 1
• The formulation reveals a relationship between transformers and RNNs, enabling autoregressive inference orders of magnitude faster.Source: §1 · p. 1
• The linear transformer reaches transformer-level performance on image generation and automatic speech recognition while achieving up to three orders of magnitude faster inference.Source: §1 · p. 1

2.1. Efficient Transformers

• Pruning, factorization, quantization, distillation, replaced-token detection, and product-key attention improve efficiency, but these methods retain quadratic sequence-length complexity; this method reduces time and memory complexity.The paper says its complexity reduction is supported theoretically and empirically.Source: §2.1 · p. 2
• Transformer-XL and learned attention spans extend context but retain vanilla asymptotic complexity, whereas improving self-attention’s asymptotic complexity allows linear transformers to use significantly larger contexts.Transformer-XL maintains previous contexts in memory at significant additional computational cost; learned attention spans control memory footprint and computation time.Source: §2.1 · p. 2
• Reformer uses LSH for O(N log N) complexity but requires keys identical to queries, precluding decoding; linear transformers impose no such constraint and are three orders of magnitude faster with comparable validation perplexity.The speed comparison is for inference in autoregressive tasks.Source: §2.1 · p. 2

2.2. Understanding Self-Attention

• Tsai et al. (2019) formulated attention as a kernel smoother to explain its components and positional embeddings; this paper instead uses kernels to accelerate attention and reports normal convergence with positive query-key similarities.The cited kernel scores represent similarity between inputs.Source: §2.2 · p. 2
• Cordonnier et al. (2020) showed multi-head self-attention with enough heads can express any convolutional layer; this paper shows autoregressive self-attention can be viewed as an RNN to speed inference.The recurrent interpretation applies to a self-attention layer trained with an autoregressive objective.Source: §2.2 · p. 2

2.3. Linearized softmax

• Building on feature-map linearizations of softmax for sampled training and concurrent image-detection work, this paper develops linear-complexity autoregressive transformers with constant memory for training and inference.It also argues that every transformer can be viewed through kernels as a recurrent neural network.Source: §2.3 · p. 2

3.1. Transformers

• Softmax attention uses the exponential query–key dot product, while generalized attention weights may use any non-negative similarity function.Source: §3.1 · p. 3, §3.2 · p. 3

3.2. Linearized Attention

• Associativity lets feature-map attention be computed in O(N) time and memory, rather than the quadratic cost of softmax attention.Source: §3.2 · p. 3

3.2.1. FEATURE MAPS AND COMPUTATIONAL COST

• A degree-2 polynomial kernel has an exact finite feature map, whereas the exponential kernel’s infinite-dimensional map prevents exact softmax linearization; experiments use an ELU-based map with O(NDM) cost.Source: §3.2.1 · pp. 3–4

3.3. Causal Masking

• Causal masking makes each position depend only on earlier positions, and constant-time updates to S_i and Z_i yield linear complexity in sequence length.Source: §3.3 · p. 4

3.3.1. GRADIENT COMPUTATION

• Cumulative-sum gradients avoid storing all intermediate S_i values, enabling linear-time forward and backward passes with constant memory in sequence length.Source: §3.3.1 · p. 4

3.3.2. TRAINING AND INFERENCE

• Training remains parallelizable, while inference uses constant time and memory per prediction by updating an internal state at each step, making it thousands of times faster than other transformers.Source: §3.3.2 · p. 4

3.4. Transformers are RNNs

• Any causally masked transformer layer can be written as a time-recurrent RNN with attention-memory and normalizer-memory hidden states, without restricting the feature function.Source: §3.4 · p. 5

4.1.2. MEMORY AND COMPUTATIONAL REQUIREMENTS

• Linear attention is faster and uses less memory than the baselines at every tested configuration; both it and Reformer scale linearly, whereas softmax scales quadratically.The benchmark used an 11-GB GPU and tested sequence lengths up to 4,096 for softmax and 16,384 for Reformer.Source: Figure 1 · p. 6, §4.1.2 · p. 6

4.1.1. CONVERGENCE ANALYSIS

• Linear attention converges smoothly to the same loss as softmax and a lower loss than Reformer, whose hashing introduces noise.Source: Figure 2 · p. 6, §4.1.1 · p. 6

4.2.1. MNIST

• Linear transformers achieve nearly the same final perplexity as softmax while generating MNIST images more than 300× faster.Their memory remains constant with sequence length because only the si and zi values need to be stored between pixels; softmax and Reformer memory grows with sequence length.Source: Table 1 · p. 7, §4.2.1 · p. 7

4.2.2. CIFAR-10

• 4,460 images can be generated by the linear model for every image generated by softmax; after seven days of CIFAR-10 training, it completed 3× more epochs and achieved better perplexity.The model also produced spatially consistent images and convincing completions.Source: Table 2 · p. 7, §4.2.2 · p. 7

4.3. Automatic Speech Recognition

• Linear attention beats the LSTM and Reformer in phoneme error rate and speed, trains more than 3× faster per epoch, while softmax has lower error but is significantly slower.Source: Table 3 · p. 8, §4.3 · p. 8

• Matrix associativity lets linear transformers compute self-attention with time and memory that scale linearly with sequence length.The model also retains linear asymptotic complexity with causal masking.Source: §5 · p. 8
• Expressing the transformer as a recurrent neural network enables autoregressive inference thousands of times faster.Source: §5 · p. 8
• Future research could explore information storage and retrieval in RNNs and transformers, and feature maps such as random Fourier features for using softmax-pretrained models.Approximating the RBF kernel with random Fourier features is offered as one possible approach.Source: §5 · p. 8

A. Gradient Derivation

• The supplement derives gradients for causally masked linear transformers that can be computed in linear time and constant memory.It focuses on the numerator of the attention expression.Source: A. Gradient Derivation · p. 11
• Autograd efficiently handles gradients of the denominator and fraction; the derivation assumes Q and K already contain vectors mapped by φ.Source: A. Gradient Derivation · p. 11
• For numerator gradients, the derivation seeks ∇QL, ∇KL, and ∇VL for Q ∈ R^{N×D}, K ∈ R^{N×D}, and V ∈ R^{N×M}.It first rewrites the expression for a single element without vector notation.Source: A. Gradient Derivation · p. 11
• Each Q_lt affects only V̄_l, so computing its gradient requires no sum over i.The supplement derives the partial derivative for Q_lt and writes the result as a matrix product of gradients, proving equation 13.Source: A. Gradient Derivation · p. 11
• Unlike Q_lt, K_j affects every V̄_i for i ≥ j, and the supplement derives its partial derivative and vectorized gradient.The vectorized result proves equation 14; the same reasoning gives the V_lt gradient and proves equation 15.Source: A. Gradient Derivation · pp. 11–12
• The cumulative-sum matrices for Q and K gradients have the same size but run in opposite directions, paralleling the forward pass and RNN backpropagation through time.Source: A. Gradient Derivation · p. 12

B. Training Evolution

• On MNIST, linear converges on par with softmax and significantly outperforms both Reformer variants.All methods are trained for 250 epochs; the short sequence length means training time does not vary significantly across methods.Source: B. Training Evolution · p. 12, Figure 5 · p. 12
• On CIFAR-10, linear and lsh-1 complete significantly more epochs than softmax and lsh-4 in seven days and achieve better performance.The passage expects this performance gap to increase as sequence length grows.Source: B. Training Evolution · p. 12, Figure 5 · p. 12
• The CIFAR-10 comparison fixes training time at seven days, during which linear and lsh-1 complete significantly more epochs than softmax and lsh-4.The passage anticipates that the performance gap will grow with longer sequences.Source: B. Training Evolution · p. 12
• On speech recognition, softmax significantly outperforms linear and Reformer in convergence, although linear is 3× faster per epoch and completes approximately 4 times more epochs.Despite softmax's advantage over linear, linear significantly outperforms Reformer in convergence and final performance.Source: B. Training Evolution · p. 12, Figure 5 · p. 12
• The speech-recognition experiment trains all methods to convergence, and linear's 3× per-epoch speed advantage lets it complete approximately 4 times more epochs than softmax.Softmax nevertheless converges significantly better than linear in this task.Source: B. Training Evolution · p. 12, Figure 5 · p. 12
• Figure 5 compares training evolution across MNIST, CIFAR-10, and speech recognition under different training schedules: 250 epochs, seven days, and convergence, respectively.The experiments' details are referenced in § 4.2.1, § 4.2.2, and § 4.3.Source: Figure 5 · p. 12

C.1. Stateful softmax attention

• Stateful-softmax saves and reuses all keys and values for each next-element prediction, so its recurrent state grows with sequence length; the proposed model instead has fixed-dimensional state and fixed-cost updates.The baseline passes the saved keys and values to the model again when predicting the next sequence element.Source: §C.1 · p. 12
• A similar stateful attention implementation for Reformer is nontrivial because sorting and chunking must be repeated for each new input.Source: §C.1 · p. 12
• More than 50× faster on CIFAR-10, the proposed formulation outperforms stateful-softmax, which is significantly faster than vanilla transformers but remains quadratic in sequence length.Table 4 summarizes autoregressive image-generation throughput for MNIST and CIFAR-10.Source: §C.1 · p. 12, Table 4 · p. 13

C.2. Equalizing the batch size

• The linear transformer is fastest across methods and nearly 6.6× faster than softmax transformers for single-image CIFAR-10 generation.Table 5 compares single-image generation times on MNIST and CIFAR-10; lower is better.Source: Table 5 · p. 13, §C.2 · p. 13
• The linear autoregressive transformer is faster on CPU than GPU in every case because its low-cost recurrent attention leaves the sequence outer loop as the main bottleneck.Source: §C.2 · p. 13
• All methods underutilize the GPU and deliver substantially lower image-generation throughput than in the preceding throughput evaluation.Source: §C.2 · p. 13

• Perplexity was approximately the same across models, and their qualitative differences were not significant.Source: D. Qualitative Results on Image… · p. 13
• Reformer models provided significantly fewer variations in their unconditional samples.Source: D. Qualitative Results on Image… · p. 13
• All models performed significantly better on image completion, which was a significantly easier task than unconditional generation.Source: D. Qualitative Results on Image… · p. 13