[Paper Review] Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention

Yunyang Xiong, Zhanpeng Zeng+5 more|arXiv (Cornell University)|Feb 7, 2021

Topic Modeling60 references29 citations

TL;DR

Nyströmformer introduces a Nyström-based approximation of self-attention that achieves O(n) time and memory, enabling efficient Transformers for long sequences while maintaining competitive performance.

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.

Motivation & Objective

• Motivate the need for scalable self-attention in Transformer models for long sequences.
• Propose a Nyström-based approximation to replace the full softmax self-attention with linear complexity.
• Develop landmark (Nyström) strategies and an iterative pseudoinverse method to approximate the softmax matrix.
• Evaluate the approach on standard NLP benchmarks and long-sequence tasks to compare accuracy and efficiency against baselines.

Proposed method

• Represent self-attention as S = softmax(QK^T / sqrt(d_q)) and aim to approximate S with O(n) complexity using Nyström methods.
• Introduce landmarks Q- and K-based to form a small m x m matrix A_S via softmax on landmarks; compute its Moore-Penrose inverse.
• Derive a Nyström-form approximation of S: hat{S} = softmax(Q ~K^T / sqrt(d_q)) A_S^+ softmax(~Q K^T / sqrt(d_q)).
• Adapt landmark selection with Segment-means to compute Q- and K-landmarks efficiently in O(n).
• Compute the Moore-Penrose inverse Z^* of A_S using an iterative method to avoid costly SVD, stabilizing around 6 iterations in practice.
• Approximate self-attention output as hat{S} V, yielding a linear-time self-attention module integrated with a skip connection via a depthwise conv on V.

Experimental results

Research questions

• RQ1 Can Nyström-based approximations of the softmax self-attention achieve comparable accuracy to standard attention on NLP benchmarks?
• RQ2 Does the Nyströmformer maintain linear time and memory complexity with respect to sequence length while preserving performance?
• RQ3 How does Nyströmformer perform on long-sequence tasks compared to other efficient-self-attention methods?
• RQ4 What are the trade-offs in landmark selection (Segment-means) and pseudoinverse approximation for practical training?

Key findings

• Nyströmformer achieves comparable or slightly better performance than baseline BERT on GLUE and IMDB with substantial efficiency gains.
• On long-range tasks (Long Range Arena), Nyströmformer performs favorably relative to other efficient self-attention approaches (e.g., Reformer, Linformer, Performer).
• The method scales linearly with sequence length n in both time and memory when the number of landmarks m is much smaller than n.
• Segment-means landmark selection enables O(n) computation with minimal overhead.
• Empirical results show Nyströmformer can outperform several linear-attention variants while maintaining accuracy close to standard self-attention.