Source-linked AI summary
Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, Vikas Singh
TL;DR
The paper addresses the quadratic memory and time cost that limits self-attention on long sequences. It adapts the Nyström method to approximate self-attention in O(n), achieving comparable performance to standard self-attention and favorable results against efficient alternatives on longer-sequence tasks.
Problem
Self-attention has O(n^2) memory and time complexity, limiting its application to longer sequences.
Method
Nyströmformer uses landmark points and an out-of-sample approximation to reconstruct the softmax attention matrix in O(n) without computing the full n × n matrix.
Results
Nyströmformer performs comparably to vanilla self-attention and favorably relative to Reformer, Linformer, and Performer on Long Range Arena, with an approximately 3.4% average-accuracy margin.
Takeaways & Limitations
The design provides a resource-efficient Transformer approach for longer sequences while retaining comparable performance on the reported tasks.
Takeaways & Limitations
The Nyström approximation requires design modifications because directly applying standard Nyström methods to softmax attention still requires computing all QK^T entries.
Abstract
from arXiv · showhide
Transformers have emerged as a powerful tool for a broad range of natural language processing tasks. A key component that drives the impressive performance of Transformers is the self-attention mechanism that encodes the influence or dependence of other tokens on each specific token. While beneficial, the quadratic complexity of self-attention on the input sequence length has limited its application to longer sequences -- a topic being actively studied in the community. To address this limitation, we propose Nyströmformer -- a model that exhibits favorable scalability as a function of sequence length. Our idea is based on adapting the Nyström method to approximate standard self-attention with $O(n)$ complexity. The scalability of Nyströmformer enables application to longer sequences with thousands of tokens. We perform evaluations on multiple downstream tasks on the GLUE benchmark and IMDB reviews with standard sequence length, and find that our Nyströmformer performs comparably, or in a few cases, even slightly better, than standard self-attention. On longer sequence tasks in the Long Range Arena (LRA) benchmark, Nyströmformer performs favorably relative to other efficient self-attention methods. Our code is available at https://github.com/mlpen/Nystromformer.
Introduction
Transformers use self-attention to model full-sequence token interactions, but its O(n^2) memory and time complexity limits longer sequences. Nyströmformer adapts the Nyström method to provide an O(n) approximation, with evaluations spanning standard and long-sequence settings.
- Self-attention computes each token representation from all other tokens, enabling interactions across the full sequence.
- O(n^2) memory and time complexity makes self-attention expensive for large models and long sequences such as n = 2048.
- Nyströmformer approximates self-attention in O(n) memory and time by using landmark points to reconstruct the softmax matrix without computing the n × n matrix.
- The evaluation uses transfer learning with pretraining on English Wikipedia and BookCorpus followed by finetuning on target tasks.
- On Long Range Arena tasks, Nyströmformer performs well relative to efficient self-attention methods, with an approximately 3.4% average-accuracy margin.
Related Work
The paper situates its approach among efficient Transformer architectures, linearized Softmax methods, and Nyström-like matrix approximations. It emphasizes that general matrix-approximation methods do not directly reduce self-attention complexity when applied to softmax matrices.
- Efficient Transformers: Efficient Transformer strategies include pruning, factorization, quantization, distillation, sparse attention, locality-sensitive hashing, and random projections.The cited approaches target memory efficiency or reduce attention complexity to O(n√n), O(n log n), or O(n).
- Efficient Transformers: Reformer reduces complexity to O(n log n) through locality-sensitive hashing, assuming keys need to be identical to queries.
- Efficient Transformers: Linformer reduces complexity to O(n) using random projections based on the JL lemma and a linear projection step.
- Linearized Softmax: Linearized Softmax methods approximate softmax-related operations with sampled classes, random Fourier features, or linear dot products.These methods are presented as approaches for efficient sampling or computation involving approximate softmax distributions.
- Nyström-like methods: Nyström-like methods approximate matrices by sampling columns, with the Nyström method originally developed for discretizing integral equations.Variants include methods using k-means and randomized sampling.
- Nyström-like methods: General matrix-approximation methods that sample rows and columns do not directly reduce self-attention complexity because accessing softmax-matrix subsets still requires calculation.
Nystr¨om-Based Linear Transformers
Nyströmformer adapts Nyström matrix approximation to self-attention by selecting landmarks before softmax, avoiding the full n × n softmax computation. With far fewer landmarks than tokens, the approximation has linear time and memory complexity while retaining similar attention patterns under supported conditions.
- Motivation: Self-attention computes n^2 pairwise similarity scores, giving O(n^2) time and memory complexity that limits use on long sequences.Each softmax entry also depends on every element in its row, reinforcing the quadratic computation.
- Challenge: Directly applying Nyström approximation remains computationally expensive because sampled softmax entries still require computing the full pre-softmax QK^T matrix.The row-wise softmax denominator sums exponentials across all elements in the same row.
- Approximation quality: An example shows Nyström approximate self-attention with attention patterns quite similar to ground-truth self-attention.The figure compares standard self-attention with the proposed approximation computed by multiplying three matrices.
- Construction: Nyströmformer selects landmark queries and keys before softmax, forms three approximation matrices, and combines them through an out-of-sample reconstruction.Segment-means produces the landmarks, while the small landmark matrix supports the Nyström construction without accessing the full softmax matrix.
- Approximation quality: Segment-means computes landmarks in O(n), and the authors report that 64 landmarks are often sufficient for a good approximation, depending on the application.The approximation is expected to improve when landmark points overlap sufficiently with the original data points; the required condition is problem dependent.
- Complexity: When m ≪ n, the proposed approximation has O(n) time and memory complexity with respect to sequence length.The analysis accounts for landmark selection, pseudoinverse approximation, and matrix multiplications before deriving linear scaling in n.
Experiments
The experiments evaluate Nyströmformer in pretraining, downstream NLP, and long-sequence settings. Across these settings, it remains competitive with standard self-attention while improving efficiency and outperforming several efficient-attention baselines.
- Experimental setup: The experiments cover language-model pretraining, GLUE and IMDB fine-tuning, and long-range tasks from the LRA benchmark.The pretraining stage uses MLM and SOP objectives, while downstream evaluation includes several natural-language-understanding datasets.
- Experimental setup: Nyströmformer replaces self-attention in BERT-small and BERT-base with the proposed Nyström approximation.BERT-small is used for comparisons with linear Transformers, while BERT-base is used as the downstream baseline.
- Efficiency: At sequence length 8192, Nyström self-attention provides 1.2× memory saving and 3× speed-up over Longformer, plus 1.7× memory saving over Linformer with similar running time.The efficiency measurements report average memory consumption and running time per input instance.
- Pretraining: Nyström self-attention is competitive with standard self-attention in MLM and SOP, while outperforming Linformer and other linear self-attention variants.The training-curve results also report competitiveness with BERT-base and faster training when initialized from pretrained BERT-base.
- Downstream NLP: Nyströmformer performs competitively with BERT-base across downstream tasks, with nearly identical IMDB accuracy at sequence lengths 512 and 1024: 93.0 vs. 93.2.The paper interprets these results as evidence that the model scales linearly with input length.
- Long Range Arena (LRA): On LRA, Nyströmformer matches vanilla self-attention in average accuracy (+0.18%) and exceeds Reformer, Linformer, and Performer by +3.91%, +3.36%, and +5.32%, respectively.The experiments use a common two-layer Transformer configuration and mean pooling across tasks.
Conclusion
The paper adapts the Nyström method to approximate self-attention with linear scaling in sequence length. Its design preserves competitive performance while offering resource-utilization benefits for longer sequences.
- Conclusion: Nyströmformer adapts the Nyström matrix-approximation method within a deep Transformer to approximate self-attention efficiently.The implementation maps the key operations to popular deep-learning libraries.
- Conclusion: Nyströmformer maintains a performance profile comparable to other self-attention approximations while providing additional resource-utilization benefits.The authors characterize the method as a step toward Transformer models on very long sequences.