Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention

TL;DR

This work addresses the quadratic cost of self-attention in Transformers by introducing Nyströmformer, a landmark-based Nyström approximation that reduces attention complexity to $O(n)$. By selecting pre-softmax landmarks $\tilde{Q}$ and $\tilde{K}$ and forming a small $m\times m$ sample matrix, it derives a Nyström form $\hat{S}$ and computes the final output with a pseudoinverse, enabling linear-time attention when $m \ll n$. The method achieves competitive results on GLUE and IMDB, and demonstrates favorable efficiency and strong performance on Long Range Arena benchmarks compared to other efficient-attention methods. The approach provides a scalable and practical route to applying Transformer-like models to very long sequences, with publicly available code to facilitate adoption.

Abstract

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.

Figures (6)

• Figure 1: A key challenge of Nyström approximation. The orange block on the left shows a $n \times m$ sub-matrix of $S$ used by Nyström matrix approximation in \ref{['eq:gennys']}. Computing the sub-matrix, however, requires all entries in the $n\times n$ matrix before the softmax function ($QK^T$). Therefore, a direct application of Nyström approximation is problematic.
• Figure 2: Illustration of a Nyström approximation of softmax matrix in self-attention. The left image shows the true softmax matrix used in self-attention and the right images show its Nyström approximation. Our approximation is computed via multiplication of three matrices.
• Figure 3: An example of Nyström approximation vs. ground-truth self-attention. Top: standard self-attention computed by \ref{['eq:trueatten']}. Bottom: self-attention from our proposed Nyström approximation in \ref{['eq:approxatten']}. We see that the attention patterns are quite similar.
• Figure 4: The proposed architecture of efficient self-attention via Nyström approximation. Each box represents an input, output, or intermediate matrix. The variable name and the size of the matrix are inside box. $\times$ denotes matrix multiplication, and $+$ denotes matrix addition. The orange colored boxes are those matrices used in the Nyström approximation. The green boxes are the skip connection added in parrallel to the approximation. The dashed bounding box illustrates the three matrices of Nystroöm approximate softmax matrix in self-attention in Eq. \ref{['eq:approxsoftmax']}. sMEANS is the landmark selection using Segment-means (averaging $m$ segments of input sequence). pINV is the iterative Moore-Penrose pseudoinverse approximation. And DConv denotes depthwise convolution.
• Figure 5: Results on masked-language-modeling (MLM) and sentence-order-prediction (SOP). On BERT-small, our Nyström self-attention is competitive to standard self-attention, outperforming Linformer and other linear self-attentions.

Theorems & Definitions (3)

• Definition 1
• Lemma 1
• Lemma 2