Source-linked AI summary
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François Fleuret
TL;DR
Very long sequences make transformers prohibitively slow because self-attention has quadratic complexity in sequence length. The paper replaces softmax attention with kernel feature maps and exploits matrix associativity to obtain linear complexity, then uses an iterative formulation connecting transformers to RNNs. Linear transformers achieve comparable performance to vanilla transformers and are reported as up to thousands of times faster for autoregressive inference.
Problem
Quadratic memory and time complexity makes transformers slow for long sequences and limits practical context length.
Method
The paper uses kernel-based feature-map attention and matrix associativity, with causal masking and an RNN formulation for iterative autoregressive computation.
Results
Linear transformers reach transformer performance levels while being up to three orders of magnitude faster during inference, with thousands-fold speedups reported for autoregressive tasks.
Takeaways & Limitations
The formulation provides transformer self-attention with linear time and memory scaling and enables faster autoregressive inference.
Takeaways & Limitations
Naive gradient computation requires storing all intermediate states, increasing memory consumption by max(D, M) times.
Abstract
from arXiv · showhide
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 $\mathcal{O}\left(N^2\right)$ to $\mathcal{O}\left(N\right)$, where $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
Transformers deliver strong results across tasks but their global self-attention imposes quadratic cost, limiting speed and context length. The linear transformer addresses this with kernelized attention and matrix associativity, enabling linear scaling and faster autoregressive inference while retaining comparable performance.
- Transformers achieve impressive results across natural language, audio, and image tasks.
- Global self-attention creates quadratic memory and time complexity for sequences of length N.
- Quadratic complexity makes transformers slow to train, limits context, disrupts temporal coherence, and hinders long-term dependency capture.
- Sparse and locality-sensitive attention methods enable efficient training on long sequences but do not accelerate autoregressive inference.
- Kernel-based attention and matrix associativity give linear transformers reduced memory use, linear context scaling, and faster autoregressive inference.The formulation also supports causal masking with linear complexity and constant memory.
- Linear transformers reach transformer performance levels while being up to three orders of magnitude faster during inference on image generation and speech recognition.
2. Related Work
Related work improves transformer efficiency through parameter reduction, sparse or hashed attention, longer contexts, kernel formulations, and softmax linearization. Linear transformers differ by reducing both asymptotic memory and time complexity while supporting unconstrained queries and keys and autoregressive inference.
- Weight pruning, factorization, quantization, distillation, and improved pretraining objectives reduce computational requirements but retain quadratic sequence-length complexity.
- Transformer-XL and learned attention spans increase usable context, but both retain the vanilla model’s asymptotic complexity.
- Sparse attention and Reformer reduce complexity for long sequences, but Reformer uses locality-sensitive hashing and constrains keys to equal queries for attention.
- Kernel-based attention has been used to interpret attention, whereas this work uses kernelization to speed attention computation.
- The paper views autoregressive self-attention as a recurrent neural network and uses that connection to accelerate autoregressive transformer inference.
- Unlike concurrent linearized-attention work for object detection, this approach also provides linear-complexity, constant-memory autoregressive inference and training without constraining queries and keys.
3. Linear Transformers
Linear transformers replace softmax attention with kernel feature-map attention, using matrix associativity to achieve linear time and memory complexity. Causal masking becomes an iterative recurrence, connecting transformers to RNNs and accelerating autoregressive inference.
- Linearized Attention: The generalized attention formulation supports non-negative similarity functions, including kernel-based attention mechanisms.The paper notes that exact softmax linearization is infeasible with its infinite-dimensional exponential-kernel feature map, whereas finite-dimensional kernels can be used.
- Linearized Attention: Kernel feature-map attention and matrix associativity reduce transformer time and memory complexity from quadratic to linear in sequence length.The exact cost depends on feature-map dimensionality and query, key, and value dimensions.
- Feature Maps and Computational Cost: The experimental feature map φ(x) = elu(x) + 1 yields attention with O(NDM) computation and performs on par with the full transformer.The authors choose ELU over ReLU to avoid zero gradients for negative inputs.
- Causal Masking: Causal linear attention computes its attention memory and normalizer state from previous states in constant time, giving linear complexity over sequence length.A naive backward implementation would require storing all intermediate attention-memory states, increasing memory consumption by max(D, M) times; cumulative-sum gradients address this issue.
- Transformers are RNNs: During training, linear transformers retain parallelizable computations, while inference stores an internal state and updates it at each timestep like an RNN.This combines accelerator-friendly training with constant per-prediction time and memory during inference.
- Transformers are RNNs: Causal transformer layers can be written as recurrent neural networks with attention and normalizer memories, revealing a time-based relation between transformers and RNNs.The formulation does not constrain the feature function and can theoretically represent transformer models using softmax attention.
4. Experiments
The experiments evaluate linear transformers against softmax attention and Reformer on computational cost, convergence, image generation, and speech recognition. Across these settings, linear transformers retain comparable or competitive performance while substantially reducing memory use and inference or training time.
- The evaluation compares linear transformers with softmax attention and Reformer across synthetic convergence, image generation, and automatic speech recognition tasks.The experiments measure computational cost, memory consumption, convergence, and real-world application performance.
- Linear attention converges smoothly on the sequence duplication task and reaches the same loss as softmax, while avoiding the hashing noise observed with Reformer.The comparison uses causal masking and tracks loss over gradient steps.
- Linear attention scales linearly with sequence length in time and memory, while softmax attention scales quadratically; the linear method is faster and uses less memory in every tested configuration.The synthetic benchmark varies sequence lengths and evaluates peak GPU memory and computation time.
- Linear transformers achieve almost the same final perplexity as softmax on MNIST while generating images more than 300 times faster.Their memory remains constant with sequence length because only the recurrent state values are stored between pixels.
- On CIFAR-10, linear transformers generate 4,460 images for every image generated by softmax, with throughput advantages increasing as image sequence length grows.The experiment attributes the stronger throughput gain to quadratic memory and time costs for softmax and Reformer during pixel generation.
- In automatic speech recognition, linear transformers outperform the recurrent network baseline and Reformer in performance and speed, while training more than 3× faster per epoch than softmax.Softmax achieves a lower phoneme error rate but is significantly slower.
5. Conclusions
The paper presents linear transformers that reduce transformer memory and computational costs by scaling self-attention linearly with sequence length. Expressing transformers as recurrent neural networks enables much faster autoregressive inference, while the resulting formulation opens directions for studying information storage and retrieval and feature-map choices.
- Linear transformers compute self-attention with time and memory that scale linearly with sequence length by exploiting matrix associativity.The formulation also supports causal masking while retaining linear asymptotic complexity.
- Expressing transformers as recurrent neural networks enables autoregressive inference thousands of times faster.
- The recurrent formulation motivates future work on information storage and retrieval in RNNs and transformers, as well as alternative feature maps for linear attention.Random Fourier features are suggested as one route toward approximating the RBF kernel and using models pretrained with softmax attention.
Supplementary Material for
The supplementary material accompanies the paper Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention.
- The supplementary material is associated with the paper Transformers are RNNs.
- Its title identifies fast autoregressive transformers as a central topic.
- Its title identifies linear attention as the paper’s attention mechanism.
A. Gradient Derivation
The supplementary derivation develops gradients for causally masked linear transformers and shows how their cumulative computations mirror recurrent backpropagation patterns. It derives vectorized gradients for Q, K, and V while preserving linear time and constant memory.
- Gradient Derivation: Gradients for causally masked linear transformers can be computed in linear time and constant memory.
- Gradient Derivation: The derivation treats Q, K, and V as matrices with shapes Q ∈ R^N×D, K ∈ R^N×D, and V ∈ R^N×M.
- Gradient Derivation: The vectorized gradients for Q, K, and V are established through matrix-product formulations and partial derivatives.
- Gradient Derivation: Q affects only the corresponding output element, whereas each K_j affects all outputs with i ≥ j, requiring different gradient summations.
- Gradient Derivation: Cumulative sums for Q and K gradients run forward and backward, respectively, paralleling the forward pass and backpropagation through time in RNNs.
B. Training Evolution
Training evolution varies across tasks: linear transformers converge competitively with softmax in some settings and outperform Reformer, while softmax remains stronger for speech-recognition convergence. The comparisons use fixed training budgets or convergence-based stopping depending on the experiment.
- MNIST: 250 epochs: on MNIST, linear transformers converge on par with softmax attention and significantly outperform both Reformer variants.Because sequences are short, training time does not vary significantly across methods.
- CIFAR-10: 7 days: on CIFAR-10, lsh-1 and linear complete significantly more epochs than softmax and lsh-4 and achieve better performance.The passage expects this gap to increase with further increases in sequence length.
- Automatic Speech Recognition: 3× faster per epoch: in speech recognition, linear completes approximately four times as many epochs as softmax, although softmax converges better.Linear still significantly outperforms Reformer in convergence and final performance.
- Overall Training Evolution: Across the experiments, linear transformers converge consistently faster than Reformer and match softmax in the autoregressive experiments.MNIST uses 250 epochs, CIFAR uses 7 days, and speech-recognition methods train to convergence.
C.1. Stateful softmax attention
The stateful-softmax baseline saves and reuses keys and values, making it faster than vanilla transformers but retaining sequence-length-dependent state and quadratic complexity. The proposed linear formulation instead uses fixed-dimensional state and is more than 50× faster for CIFAR-10.
- Stateful softmax attention: Stateful-softmax saves keys and values and reuses them to predict each subsequent element.Its recurrent state therefore grows proportionally with sequence length.
- Stateful softmax attention: Stateful-softmax is significantly faster than vanilla transformers but remains quadratic in sequence length.The baseline improves speed without removing the underlying quadratic complexity.
- Stateful softmax attention: More than 50× faster CIFAR-10 generation is achieved by the proposed formulation than by stateful-softmax.The comparison concerns autoregressive image-generation throughput.
C.2. Equalizing the batch size
The latency experiment measures single-image generation with batch size 1 on both CPU and GPU. The linear transformer is faster than all evaluated methods and is almost 6.6× faster than softmax transformers on CIFAR-10.
- Equalizing the batch size: Batch size 1 latency is measured as the total time to generate one MNIST or CIFAR-10 image on CPU and GPU.Lower reported time is better.
- Equalizing the batch size: All methods underutilize the GPU and therefore achieve lower image-generation throughput than in the batch-throughput experiment.This motivates comparing single-image latency separately from throughput.
- Equalizing the batch size: Almost 6.6× faster CIFAR-10 image generation is achieved by the linear transformer than by softmax transformers.The comparison uses single-image generation latency with batch size 1.
- Equalizing the batch size: The linear autoregressive transformer is faster on CPU than GPU in every reported case.The passage attributes this to computing attention as an RNN.
D. Qualitative Results on Image Generation
Qualitative image-generation differences are generally small because all models have approximately the same perplexity, but Reformer samples show fewer variations. Image completion is easier than unconditional generation for all models.
- D. Qualitative Results on Image Generation: Approximately equal perplexity across models corresponds to no significant qualitative differences overall.The qualitative comparison concerns the image-generation experiments.
- D. Qualitative Results on Image Generation: Reformer models provide significantly fewer variations in their unconditional samples than the other evaluated models.This observation concerns unconditional image generation.
- D. Qualitative Results on Image Generation: Image completion is significantly easier than unconditional generation because all models perform significantly better on completion.The paper reports this pattern for both MNIST and CIFAR-10 qualitative experiments.