[Paper Review] Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

Angelos Katharopoulos, Apoorv Vyas+2 more|arXiv (Cornell University)|Jun 29, 2020

Neural Networks and Applications35 references338 citations

TL;DR

The paper introduces linear transformers that replace softmax attention with a kernel-based linear attention, enabling O(N) time/memory and autoregressive inference that is thousands of times faster while retaining comparable performance to vanilla transformers.

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 $\\mathcal{O}\\left(N^2\ ight)$ to $\\mathcal{O}\\left(N\ ight)$, 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.

Motivation & Objective

• Motivate and address the quadratic cost of self-attention in transformers for long sequences.
• Propose a linear attention formulation using kernel feature maps to reduce complexity.
• Show causal masking with linear attention and derive an RNN-like recurrence for efficient autoregressive inference.
• Empirically validate on image generation and speech recognition against baselines.

Proposed method

• Reformulate self-attention as a kernelized dot-product with a feature map phi so that V' = (phi(Q) (phi(K)^T V)) / (phi(Q)^T sum_j phi(K_j)).
• Use associative property to compute sums S = sum_j phi(K_j) V_j^T and Z = sum_j phi(K_j) for all queries, yielding O(N) time for forward pass.
• Introduce causal masking with incremental updates to S_i and Z_i to achieve linear time and constant memory for autoregressive generation.
• Derive an RNN-like view of a transformer layer with causal masking, featuring attention memory s and normalizer memory z and a two-step update.

Experimental results

Research questions

• RQ1 Can self-attention in transformers be reformulated to achieve linear (O(N)) time and memory without sacrificing performance?
• RQ2 Does a kernel-based linear attention with a positive similarity function enable effective autoregressive inference comparable to softmax attention?
• RQ3 Can transformer layers with causal masking be interpreted as RNNs, enabling efficient sequential generation?
• RQ4 What are the empirical trade-offs in accuracy and speed for long-sequence tasks (image generation, speech recognition) using linear attention vs. standard transformers and Reformer?
• RQ5 How does linear attention perform under training vs inference conditions for autoregressive tasks?

Key findings

• Linear transformers attain similar performance to vanilla transformers while being up to 4000x faster on autoregressive prediction of long sequences.
• Memory and time scale linearly with sequence length for linear attention, unlike quadratic scaling in softmax attention.
• In MNIST, CIFAR-10, and WSJ experiments, linear transformers achieve competitive bits/dim or PER with substantially higher throughput and lower memory.
• Autoregressive inference with linear attention can be performed with constant per-step cost by maintaining and updating internal state (S_i, Z_i).
• The work establishes a theoretical link between transformers with causal masking and recurrent neural networks, via an RNN-like recurrence for attention memory.