ArXiv: 2102.03902
🎯 Pitch
Standard Transformers choke on long sequences due to quadratic self-attention costs, but Nyströmformer slashes this to O(n) — achieving a 22.8× memory reduction and 12.7× speedup at length 8192 while matching or even beating BERT on GLUE benchmarks (91.4 vs. 90.0 on SST-2). The trick is using a small set of learned landmarks to reconstruct the full attention matrix via the Nyström method, with surprisingly little loss in fidelity.
1. Executive Summary
This paper proposes Nyströmformer, a Transformer model that approximates the standard self-attention mechanism with O(n) time and memory complexity by adapting the Nyström method for matrix approximation. Evaluated on the GLUE benchmark, IMDB reviews, and the Long Range Arena (LRA) benchmark using BERT-style pretrained models, Nyströmformer replaces the full n×n softmax attention matrix with a product of three smaller matrices constructed from m landmark points (where m ≪ n) — using Segment-means to select inducing points from queries and keys, and an iterative approximation of the Moore-Penrose pseudoinverse to invert the landmark kernel. The method achieves comparable or slightly better accuracy than standard BERT-base on GLUE tasks (e.g., 91.4 vs. 90.0 on SST-2) while reducing memory consumption by 22.8× and providing 12.7× speedup at a sequence length of 8192 relative to standard self-attention, establishing that the Nyström approximation maintains competitive downstream performance across standard-length and long-sequence tasks only when the number of landmarks is set such that the landmark-softmax kernel product faithfully reconstructs the full attention matrix — with 64 landmarks sufficing for most applications.
2. Context and Motivation
The Core Problem: Quadratic Self-Attention Makes Long Sequences Prohibitively Expensive
The fundamental problem Nyströmformer addresses is the quadratic complexity bottleneck in Transformer self-attention. Standard self-attention, as introduced by Vaswani et al. (2017), computes pairwise similarity scores between every token and every other token in a sequence. For a sequence of length , this produces an attention matrix — meaning the computational and memory costs scale as . This is not merely a theoretical inconvenience; it imposes hard practical limits on what Transformers can process.
The paper frames this concretely in Section 1: training BERT-large requires 4 months on a single Tesla V100 GPU (or 4 days on a 4×4 TPU pod), and the quadratic cost makes sequences beyond a few thousand tokens prohibitively expensive or infeasible to train. The authors explicitly state:
"the O(n²) complexity makes it prohibitively expensive to train large Transformers with long sequences (e.g., n = 2048)"
This is the gap the paper targets — not improving self-attention's expressiveness (it already works well), but making it computationally tractable for long sequences without sacrificing the quality that makes it effective.
Why This Problem Matters: Beyond Academic Convenience
The quadratic bottleneck has real-world consequences across multiple dimensions:
Long-document understanding. Many NLP tasks naturally involve long inputs — legal documents, scientific papers, books, long conversations, or multi-hop question answering where the model must integrate information across paragraphs or pages. Standard Transformers with a 512-token limit (common in BERT) simply cannot process these inputs end-to-end. Either the input must be truncated (losing information) or chunked into segments (losing cross-segment dependencies). A linear-complexity self-attention mechanism would enable these applications without compromise.
High-resolution inputs in other modalities. The paper notes in its conclusion that scaling to longer sequences is "desirable in both NLP as well as computer vision." When images are represented as sequences of pixels or patches (as in Vision Transformers), high-resolution images produce very long sequences. Similarly, speech processing, video understanding, and genomics all involve sequences where quadratic scaling quickly becomes infeasible. The O(n) solution is therefore cross-domain infrastructure, not an NLP-specific trick.
Democratizing large-model training. The 4-month single-GPU training time for BERT-large means that only organizations with substantial compute resources can train these models. Reducing the complexity from quadratic to linear makes it feasible for smaller research groups to experiment with long-sequence Transformers, broadening participation in the field.
Inference latency in production. Even if a model can be trained with quadratic attention (perhaps via distributed training with many GPUs), deploying it for inference at scale with long inputs remains expensive. Each forward pass through a quadratic self-attention layer costs , which translates directly to latency and dollar cost per query. Linear attention reduces these costs, making long-sequence models deployable in production settings.
Prior Approaches and Their Limitations
The paper situates itself within a rapidly growing landscape of efficient Transformer variants. Rather than listing every approach, I'll organize them by the core strategy they employ and explain where each falls short — drawing from the paper's Related Work (Section 2 in the original manuscript) and its framing.
Strategy 1: Sparse or Localized Attention Patterns
These methods avoid computing the full attention matrix by restricting which token pairs can attend to each other, typically via a sparsity pattern.
-
Sparse Transformers (Child et al., 2019): Use a fixed sparse factorization of the attention matrix, reducing complexity to . The limitation is that the sparsity pattern is hand-designed and may miss important long-range dependencies that fall outside the predefined pattern.
-
Reformer (Kitaev, Kaiser, and Levskaya, 2019): Uses locality-sensitive hashing (LSH) to group queries and keys into buckets, so that attention is only computed within each bucket. This reduces complexity to . However, the authors note an important constraint: this method "relies on performing fewer dot product operations overall by assuming that the keys need to be identical to the queries" — meaning it requires tied QKV weights, which limits architectural flexibility. More subtly, hash collision failures can cause tokens to miss important interactions.
-
Longformer (Beltagy, Peters, and Cohan, 2020): Combines a sliding window (local) attention with task-motivated global attention on specific tokens. This achieves O(n) complexity. The limitation is that the global attention pattern must be task-specifically designed — certain tokens are pre-designated as "global" and attend to everything, while others only attend locally. This works well for tasks where you know which tokens are globally important (e.g., the [CLS] token in classification), but less naturally for tasks where global interactions are more distributed.
-
BigBird (Zaheer et al., 2020): Uses a combination of random, window, and global attention to achieve O(n) while theoretically proving the sparsity pattern preserves the expressiveness of full attention. The limitation is similar to Longformer: the global tokens must be specified, and the random attention introduces stochasticity.
The key shortfall across these sparse approaches is that they sacrifice full pairwise interaction to gain efficiency. The sparsity pattern is a structural constraint on what the model can learn — and while theoretical results (like BigBird's) show that certain patterns can approximate full attention, there's no guarantee that the chosen pattern captures the dependencies relevant to a specific task.
Strategy 2: Linearized Attention via Kernel Approximations
These methods approximate the softmax attention operation as a kernel function, then use the associativity property of matrix multiplication to change the computation order — computing instead of , reducing complexity to .
-
Linearized attention (Katharopoulos et al., 2020; Shen et al., 2018b): Replace the softmax with a simpler activation (e.g., ELU) so that the attention can be decomposed as . The authors acknowledge this idea is "interesting" but directly criticize it: "the approximation error to the softmax matrix in self-attention can be large in some cases." The fundamental problem is that softmax is not a kernel function — it couples all entries in a row through the normalization denominator — so any kernel-based decomposition introduces an inherent approximation error that compound across layers.
-
Performer (Choromanski et al., 2020): Uses positive orthogonal random features to approximate the softmax kernel, providing unbiased or near-unbiased estimates of the full attention. This is a more sophisticated kernel approximation than ELU-linear, with theoretical guarantees. However, the random feature dimension must be sufficiently large to keep variance low, and the approximation quality degrades as sequence length grows relative to the feature dimension — there's a bias-variance tradeoff that must be managed.
-
Linear Transformer (from the broader literature, not to be confused with Linformer): These methods generally work by changing the computation order as described above. The paper groups them in the "Linearized Softmax" subsection of Related Work.
The core limitation of kernel-based approximations is that they change the attention function itself, not just how it's computed. The softmax — with its sharp nonlinearity and competition among tokens via normalization — is what gives standard self-attention its ability to focus sharply on relevant tokens while suppressing irrelevant ones. Replacing it with a smoother kernel function necessarily alters this behavior, potentially degrading the model's ability to learn selective attention patterns.
Strategy 3: Low-Rank Projections of the Attention Matrix
These methods observe that the attention matrix is often low-rank and project it to a smaller dimension before computing pairwise interactions.
- Linformer (Wang et al., 2020): Projects the key and value matrices from dimension to where , using learned projection matrices based on the Johnson-Lindenstrauss lemma. This reduces complexity to , which is when is fixed. The limitation: Linformer assumes the attention matrix is low-rank and compresses along the sequence-length dimension, which means the compressed representation size becomes a hyperparameter that trades efficiency for fidelity. If is too small for a given task, information is lost. Moreover, the projection matrices are learned — they add parameters and require the model to learn how to compress, which may not transfer across sequence lengths.
The paper specifically positions Nyströmformer as an alternative to Linformer, with both aiming for O(n) complexity but via different mathematical mechanisms. Linformer uses learned projections; Nyströmformer uses data-dependent landmark sampling.
Strategy 4: Other Efficiency Approaches (Not Direct Competitors)
The paper also acknowledges a broader ecosystem of efficiency techniques that are orthogonal to their contribution: weight pruning (Michel, Levy, and Neubig, 2019), weight factorization (ALBERT; Lan et al., 2020), weight quantization (Q8BERT; Zafrir et al., 2019), knowledge distillation (DistilBERT; Sanh et al., 2019), and training optimizations like micro-batching and gradient checkpointing. These improve efficiency without changing the attention mechanism itself, so they are complementary rather than alternative approaches.
Where All Prior Approaches Fall Short: The Unified Gap
Across sparse attention, kernel approximations, and low-rank projections, a common thread emerges: every existing method either constrains the attention pattern, changes the attention function, or compresses the representation in a way that introduces approximation error not grounded in the actual structure of the specific input sequence.
Sparse methods constrain which token pairs interact. Kernel methods alter how token similarities are computed (no true softmax). Low-rank projections compress dimensions but do so agnostically via learned matrices. None of these approaches directly approximates the actual softmax attention matrix for the specific input at hand, which is what Nyströmformer attempts to do.
How Nyströmformer Positions Itself
The paper frames Nyströmformer as a fundamentally different approach: data-dependent low-rank approximation of the softmax matrix itself, computed by sampling actual columns and rows from the true attention matrix (via landmarks) rather than learning projections or changing the operator.
The key intellectual move is adapting the Nyström method — a classical numerical linear algebra technique for approximating large kernel matrices by sampling a subset of rows and columns — to the self-attention setting. The Nyström method has a long history in machine learning for accelerating kernel methods (Williams and Seeger, 2001), but applying it directly to self-attention hits a roadblock the paper identifies explicitly:
"directly applying these methods to approximate a softmax matrix used by self-attention does not directly reduce the computational complexity. This is because even accessing a subset of columns or rows of a softmax matrix will require the calculation of all elements in the full matrix before the softmax function. And calculating these entries will incur a quadratic cost."
This is the paper's primary technical insight (described in Section 3's Figure 1 and surrounding text): the softmax normalizer couples all entries in a row, so you cannot compute just part of the softmax matrix without computing all of first. Nyströmformer circumvents this by selecting landmarks before the softmax — choosing a small set of representative queries and keys (via Segment-means), computing the softmax only on those landmarks and between landmarks and the full set, and then reconstructing the full softmax matrix via the Nyström quadrature formula.
This positions Nyströmformer in a unique spot in the efficient-attention landscape:
- Vs. sparse attention: Nyströmformer does not restrict which tokens can interact — the full attention matrix is reconstructed, just approximately.
- Vs. kernel approximations: Nyströmformer preserves the softmax operator — the approximation is to the matrix, not to the function. The softmax is still computed, just on a subset of the data.
- Vs. Linformer's learned projections: Nyströmformer's approximation is data-dependent — the landmarks are derived from the actual query and key vectors for the specific input, so the approximation adapts to each sequence.
- Vs. prior Nyström applications: The paper's innovation is recognizing that you can apply the Nyström method before rather than after the softmax, which is a compromise but avoids the quadratic bottleneck. The authors explicitly acknowledge this trade-off: it is "a compromise" relative to the mathematically purer post-softmax sampling, but it "avoids the need to compute the full softmax matrix S."
The paper also draws on an important observation from the Linformer paper to justify the low-rank assumption underlying the Nyström approximation: "the authors suggest that self-attention is low-rank." If self-attention matrices are approximately low-rank (which empirical evidence from prior work supports), then a Nyström approximation with a modest number of landmarks should be able to reconstruct them accurately — meaning the compromise is well-founded.
In summary, Nyströmformer offers a matrix approximation perspective on efficient attention, contrasting with the pattern-constraint, kernel-replacement, and learned-compression perspectives that dominated prior work. The paper argues that this approach preserves the mathematical structure of self-attention more faithfully while achieving comparable or better efficiency, and the experimental results (anticipating later sections) bear out that this fidelity translates to competitive downstream performance.
3. Technical Approach
3.1 Reader orientation (approachable technical breakdown)
Nyströmformer is a modified Transformer architecture where the standard quadratic self-attention layer is replaced with a linear-complexity approximation that computes attention using a small set of learned representative tokens rather than all pairwise interactions. The system solves the problem of Transformers being prohibitively expensive for long sequences by approximating the full softmax attention matrix as the product of three smaller matrices — an matrix, an matrix, and an matrix — where (the number of landmark points) is much smaller than the sequence length , reducing both memory and computation from to .
3.2 Big-picture architecture (diagram in words)
The Nyströmformer self-attention module has five sequential stages, each transforming one intermediate matrix into the next:
-
Landmark selection (Segment-means): The input query matrix and key matrix are each downsampled to landmark vectors by partitioning the tokens into equal segments and averaging within each segment. This produces and — the landmark queries and landmark keys.
-
Three-kernel computation: Three softmax-based similarity matrices are computed: (a) the landmark-to-landmark kernel , (b) the full-query-to-landmark-key kernel , and (c) the landmark-query-to-full-key kernel . These are all linear in because is constant.
-
Pseudoinverse of the landmark kernel: The Moore-Penrose pseudoinverse of the matrix is computed via an iterative approximation algorithm (6 iterations of a third-order convergence scheme), producing . This costs , independent of .
-
Nyström reconstruction: The full softmax attention matrix is approximated as , which is an matrix of attention weights. Critically, this matrix is never materialized — only the products needed for the final output are computed.
-
Value aggregation with skip connection: The approximate attention weights multiply the value matrix to produce , and a depthwise convolution skip connection (kernel size 1) of is added to stabilize training. The result is the output of the Nyströmformer attention layer.
The core insight is that steps 1–3 operate on matrices where one dimension is (not ), and step 4 multiplies these matrices in an order that avoids ever forming an matrix, keeping the entire pipeline at complexity.
3.3 Roadmap for the deep dive
- First, the standard self-attention formulation — we need the exact mathematical definition (Equation 2) to understand what Nyströmformer is approximating, why the quadratic cost arises, and where the softmax normalizer creates the fundamental obstacle for direct Nyström application.
- Second, the classical Nyström method for matrix approximation — how it reconstructs a full matrix from a subset of rows and columns, why a naive application to the softmax matrix fails (the softmax normalizer requires computing all entries of even to access a subset of columns of ), and how this motivates the key design choice.
- Third, the adapted Nyström approximation for self-attention — how moving the landmark selection before the softmax operation circumvents the quadratic bottleneck while still producing a valid Nyström reconstruction, including the derivation of the three-matrix product form.
- Fourth, the Moore-Penrose pseudoinverse computation — the iterative algorithm (Lemma 1) used because SVD is inefficient on GPUs, including the initialization scheme that guarantees convergence.
- Fifth, the landmark selection mechanism (Segment-means) — why K-means is impractical during mini-batch training and how Segment-means provides an alternative.
- Sixth, the full architecture and complexity analysis — how the components assemble into a complete self-attention replacement (Figure 4), the skip connection design, and the explicit complexity derivation.
- Seventh, the theoretical analysis (Lemma 2) — the conditions under which the Nyström approximation converges to true self-attention and the connection to the low-rank property of attention matrices.
This order mirrors the paper's own logical flow and builds from the known (standard self-attention) → the problematic (naive Nyström application) → the solution (pre-softmax landmark selection) → the implementation details → the full system.
3.4 Detailed, sentence-based technical breakdown
This is primarily a method paper whose core idea is that the Nyström method for matrix approximation — which reconstructs a full matrix from a sampled subset of its rows and columns — can be adapted to approximate the softmax attention matrix in Transformers with linear complexity, provided the landmark selection is performed before rather than after the softmax operation to avoid the quadratic cost of computing the full matrix.
Standard Self-Attention: The Mathematical Baseline
The paper begins by precisely defining the self-attention operation that Nyströmformer approximates. An input sequence of tokens with dimension is represented as a matrix . This matrix is projected using three learned weight matrices — (query projection), (key projection), and (value projection) — with the constraint that . The projections produce three matrices:
where is the query matrix (what each token "wants to find"), is the key matrix (what each token "offers"), and is the value matrix (the content each token provides).
The self-attention output is then:
What this equation computes: For each token (row of ), compute its dot-product similarity with every token (column of ), producing an matrix of raw attention scores . Scale each score by (the scaling factor prevents dot products from growing too large in high dimensions, which would push the softmax into regions of extremely small gradients). Apply a row-wise softmax to convert each row of scores into a probability distribution over tokens — this is the attention matrix , where represents the weight with which token attends to token . Finally, multiply by : each output row is a weighted sum of all value vectors, with weights given by the attention probabilities. The result contains, for each token, a context-aware representation that integrates information from the entire sequence.
Why this form: The softmax ensures that attention weights are non-negative and sum to 1, creating a convex combination of value vectors — a form of soft content-based lookup. The scaling factor is a variance-stabilization trick from the original Transformer paper: without it, for large , the dot products have large variance, causing the softmax to saturate (produce near-one-hot distributions), which makes gradients vanish and prevents learning. The matrix multiplication is the associative form that enables the attention operation; critically, computing requires evaluating all pairwise dot products (the matrix multiplication) and then normalizing each of the rows via the softmax denominator , giving the memory and time complexity.
The paper emphasizes in Section 3.1 (under the paper's own section numbering) that "each element in the softmax matrix S depends on all other elements in the same row." This is the root cause of the quadratic cost and the central obstacle for any method that tries to compute only a subset of — you cannot compute without knowing the entire row's values because the softmax denominator normalizes against the sum of exponentials across the full row.
Classical Nyström Method for Matrix Approximation: How It Works and Why It Fails for Self-Attention
The Nyström method (Baker, 1977; Williams and Seeger, 2001) is a technique for approximating a large matrix using only a subset of its rows and columns. The paper walks through the classical formulation to establish what it is adapting and, crucially, where the direct application breaks down.
Given a matrix (in our context, the softmax attention matrix), the Nyström method partitions it by selecting landmark rows and columns:
where is the submatrix formed by the intersection of the selected rows and selected columns (the landmark-landmark block), is the remaining entries in the selected rows, is the remaining entries in the selected columns, and is the large remaining block that we want to avoid computing.
The Nyström approximation exploits the singular value decomposition (SVD) of the landmark block :
where are orthogonal matrices and is a diagonal matrix containing the singular values of .
What this decomposition enables: The SVD factors into rotation/reflection matrices (, ) and a scaling matrix (). Because captures the structure of the full matrix restricted to the landmark set, the rotation and scaling factors can be used to extrapolate to the unobserved blocks. The Nyström reconstruction formula — derived from the out-of-sample column approximation (Wang and Zhang, 2013) — approximates the full matrix as:
where is the Moore-Penrose pseudoinverse of .
What this equation computes: The full matrix is reconstructed from only three pieces of information: the matrix (all rows, but only the landmark columns), the pseudoinverse , and the matrix (all columns, but only the landmark rows). The block — which accounts for entries — is approximated as , a low-rank factorization of rank at most . The formula essentially says: the interaction between any two non-landmark tokens can be approximated by how they each relate to the landmarks, composed through the inverse landmark interaction matrix.
Why this form: The Nyström approximation is the optimal rank- approximation of under the constraint that it exactly matches on the selected rows and columns (the "out-of-sample extension" property). It has strong theoretical guarantees: if is exactly rank , the Nyström approximation is exact (); if is approximately rank , the error is bounded by the tail of the singular value spectrum. This is why the low-rank property of self-attention matrices — observed in prior work (Wang et al., 2020) and cited by the authors — makes the Nyström method theoretically well-suited for this application.
For the softmax attention matrix specifically, the paper writes out the Nyström form explicitly in Equation 6:
where denotes taking columns from the full softmax matrix and denotes taking rows.
The Key Challenge: The Softmax Normalizer Creates a Circular Dependency
This is the paper's central technical observation, illustrated in Figure 1. The Nyström formula requires computing only a subset of columns and rows of — specifically, the landmark columns and landmark rows. However, to compute even a single entry of the softmax matrix , you need:
The problem: The denominator sums over all keys for each query . Therefore, to compute any entry in the landmark columns — which the Nyström formula requires — you must compute the full row of to normalize it. This means you must compute all entries of anyway, even though the Nyström approximation only needs columns of the result. The softmax is not factorizable: it couples all entries in a row through the normalization constant, so there is no shortcut to computing a subset of columns.
The paper states this precisely:
"directly applying these methods to approximate a softmax matrix used by self-attention does not directly reduce the computational complexity. This is because that even accessing a subset of columns or rows of a softmax matrix will require the calculation of all elements in the full matrix before the softmax function."
Figure 1 visually reinforces this: it shows an matrix with an sub-block highlighted in orange, and the caption explains that "computing the sub-matrix ... requires all entries in the matrix before the softmax function." This is the roadblock that Nyströmformer is designed to circumvent.
The Nyströmformer Key Insight: Moving Landmark Selection Before Softmax
The paper's fundamental innovation is to select landmarks from the queries and keys before the softmax operation, rather than sampling columns from the softmax matrix after it is computed. This is a compromise — mathematically, the pure Nyström method would sample from the actual matrix being approximated — but it is the compromise that breaks the circular dependency and enables linear complexity.
Definition 1 (Landmark Matrices): Let the selected landmarks from the queries be and from the keys be . The corresponding matrix forms are:
These landmark matrices are much smaller than the original and — they have rows instead of , where is typically 32 or 64 (constants independent of sequence length).
The modified Nyström procedure then works as follows:
Step 1: Compute the landmark-to-landmark softmax kernel. Apply the softmax operation only to the matrix formed by the landmarks:
where . Its SVD is .
What this computes: The pairwise attention weights among the landmark tokens, normalized via row-wise softmax. This is the analogue of the block in the classical Nyström formula, but computed from landmark vectors rather than sampled from the full softmax matrix. The computation costs — independent of .
Step 2: Compute cross-kernels between full tokens and landmarks. For each original query and landmark key , and for each landmark query and original key , define:
What these compute: is the attention distribution of query over only the landmark keys — it answers "how much does this query attend to each landmark?" — normalized by a softmax that sums over only landmarks, not all tokens. is the attention distribution of the landmark queries over key , normalized column-wise (or, equivalently, the transpose of the attention of landmark queries to key ).
Step 3: Construct the Nyström feature maps. Following the out-of-sample approximation procedure from the classical method, the paper defines feature representations for each query and key:
where and .
What these equations compute: They project each original query and key into an -dimensional feature space defined by the landmark structure. The factor scales the features so that dot products in this space reconstruct the Nyström approximation. The computation uses the SVD factors of (, , ) as the projection bases — these encode how the landmarks relate to each other.
Step 4: Reconstruct the full attention matrix via dot products. The Nyström approximation of the entry is:
Step 5: Derive the explicit three-matrix product form. Under the assumption that is non-singular (which the paper later relaxes to use the pseudoinverse), the dot product expands to:
Let . Since the SVD of is , we have , meaning .
Therefore, .
Relaxing the non-singularity assumption (replacing with the Moore-Penrose pseudoinverse ):
Expanding the definitions of and :
In matrix form, the full approximate softmax matrix is:
This is Equation 13 in the paper.
What this equation computes: The full attention matrix is reconstructed as the product of three matrices: an matrix (left factor), an matrix (middle pseudoinverse), and an matrix (right factor). The left factor contains, for each of the original queries, its attention distribution over the landmark keys. The right factor contains, for each of the landmark queries, its attention distribution over the original keys. The middle factor is the inverse of the landmark-to-landmark attention matrix, which serves as a "bridge": it captures how the landmarks relate to each other, and by composing the left and right factors through this bridge, the Nyström formula extrapolates from landmark-token interactions to token-token interactions.
Why this form: The three-matrix factorization achieves two critical properties simultaneously. First, it is linear in : the left factor costs (computing and softmax over landmarks), the right factor costs (computing and softmax over landmarks), and the middle factor costs (pseudoinverse of an matrix). The product of the three matrices costs if computed in the correct associative order — never forming the full matrix. Second, it preserves the softmax operator: unlike kernel-based linearization methods that replace softmax with a different function (e.g., ELU), Nyströmformer still applies genuine row-wise softmax — just over a smaller set ( landmarks instead of tokens). This is why the approximation quality can be high: the softmax's sharp selectivity is preserved for the landmark interactions that form the approximation's basis.
The compromise: The paper explicitly acknowledges that selecting landmarks before softmax is a deviation from the mathematically purer classical Nyström method (which would sample from the softmax matrix itself). The authors state: "This is a compromise but avoids the need to compute the full softmax matrix S for a Nyström approximation." The cost of this compromise is that the landmark queries and keys may not perfectly represent the structure of the softmax matrix — but the empirical results suggest that with appropriate landmark selection (Segment-means) and a sufficient number of landmarks (64), the approximation is highly faithful.
Figure 2 provides a visual comparison: the left panel shows a true softmax attention matrix, and the right panels show its Nyström reconstruction via three matrices of sizes , , and . The multiplication of these three matrices produces an matrix that visually resembles the original.
Moore-Penrose Pseudoinverse Computation: Iterative Approximation for GPU Efficiency
The Nyström reconstruction formula requires computing , the Moore-Penrose pseudoinverse of the landmark kernel matrix. The standard approach — singular value decomposition (SVD) — is "not very efficient on GPUs" (Section 3.2 in the paper's own numbering). Since the pseudoinverse must be computed in every forward pass (and gradients must flow through it during backpropagation), SVD would become a significant bottleneck.
The paper adopts an iterative approximation scheme from Razavi et al. (2014), which converges to the Moore-Penrose inverse through a sequence of matrix-matrix multiplications — operations that are highly optimized on GPUs.
Lemma 1 (Iterative Pseudoinverse): For , the sequence generated by:
converges to the Moore-Penrose inverse with third-order convergence, provided the initial approximation satisfies .
What this equation computes: Each iteration applies a high-order polynomial of the matrix to , pushing closer to the pseudoinverse. The constants (13, 15, 7) come from the Taylor expansion of the matrix inverse function and are chosen to achieve third-order convergence (the error decreases cubically with each iteration). The factor normalizes the polynomial. The nested form is a computationally efficient way to evaluate a cubic polynomial in — it requires only three matrix multiplications per iteration (, then the nested products).
Why this form: Third-order convergence means that after iterations, the error is roughly — extremely rapid when the initial error is less than 1. This allows the paper to use only about 6 iterations to achieve a good approximation of the pseudoinverse. The alternative — computing SVD — would require operations with poor GPU parallelism (involving sequential operations like bidiagonalization), while the iterative scheme uses only matrix multiplications, which achieve near-peak GPU throughput.
Initialization scheme: The paper initializes using:
where (maximum absolute column sum) and (maximum absolute row sum).
Why this initialization: This choice is based on Pan and Schreiber (1991) and ensures that . When is non-singular, , satisfying the convergence condition of Lemma 1. Even when is singular, the paper states that this initialization "provides a good approximation in our experiments." The scaling by normalizes so that the eigenvalues of are roughly bounded by 1, preventing divergence in the early iterations.
Integration into the Nyström approximation: Let be the result after approximately 6 iterations. The Nyström approximation becomes:
Gradient computation: Because the iterative scheme uses only matrix multiplications, backpropagation through the pseudoinverse computation is straightforward — the autograd system can trace through the 6 iterations of matrix products, computing gradients via the chain rule. This would be more complex with an SVD-based pseudoinverse, which involves eigen-decomposition where gradients may be numerically unstable near repeated singular values.
Landmark Selection: Segment-Means as an O(n) Alternative to K-Means
The quality of the Nyström approximation depends critically on how well the landmarks represent the full set of queries and keys. If the landmarks are poorly chosen — e.g., they cluster in one region of the embedding space — the extrapolation from landmarks to all tokens will be inaccurate.
The paper acknowledges that K-means clustering is a standard approach for selecting Nyström landmark points (Zhang, Tsang, and Kwok, 2008; Vyas, Katharopoulos, and Fleuret, 2020). K-means would partition the query vectors into clusters and use the cluster centroids as landmarks, providing a good covering of the query space.
Why K-means is problematic during training: The paper states that "the EM style of updates in K-means is less desirable during mini-batch training." K-means is an iterative algorithm (alternating between assignment and update steps) that would need to run to convergence for every mini-batch — this is computationally expensive and introduces non-differentiable operations (hard cluster assignments) that complicate gradient flow. Moreover, the cluster centroids would change discontinuously as tokens move between clusters during training, potentially causing training instability.
Segment-means: The paper proposes a much simpler scheme based on local average pooling, similar to approaches used for sequence summarization in NLP (Shen et al., 2018a). For input queries , the queries are partitioned into contiguous segments of equal size. Assuming is divisible by (which can be ensured by padding), let the segment length be . The -th landmark query is:
Similarly, for keys :
What this computes: Each landmark is the average of a contiguous block of consecutive tokens in the sequence. The first tokens produce the first landmark, the next tokens produce the second landmark, and so on. This is essentially applying a 1D average pooling with kernel size and stride along the sequence dimension.
Why this works despite its simplicity: There are two reasons this scheme is effective:
-
Locality in Transformer representations: In many NLP tasks, adjacent tokens often have similar semantic roles and their representations are correlated (due to local context). Averaging contiguous blocks preserves the local structure while reducing dimensionality — each landmark captures the "average meaning" of its segment.
-
Uniform coverage of the sequence position: By segmenting along the sequence dimension (not by clustering in embedding space), Segment-means guarantees that landmarks are spread across the entire sequence. This prevents the pathological case where all landmarks concentrate at one position and fail to represent distant tokens. The approximation then works by interpolating attention patterns between positions based on the landmark tokens at those positions.
Computational cost: Segment-means requires a single scan of the sequence to compute running sums for each segment — time complexity. This is dramatically cheaper than K-means, which would require .
Number of landmarks: The paper states: "We find that using 64 landmarks is often sufficient to ensure a good approximation, although this depends on the application." In the experiments, they use configurations with 32 and 64 landmarks (Nyströmformer-32 and Nyströmformer-64), and Table 1 shows that 32 landmarks provide even greater memory and speed savings while still maintaining reasonable performance.
Normalization detail for Equation 16: The paper's Equation 16 shows the sum being divided by (the number of landmarks), not by (the segment length). However, this appears to be a typographical error in the original paper — the sum should be divided by (the number of elements in the segment) to compute an average. The description as "Segment-means" and the analogy to "local average pooling" confirm that the operation is averaging within each segment, which would be dividing by , not by . The mathematical intent is clearly , and this is the operation used in practice.
Putting It Together: Approximate Self-Attention and the Full Architecture
With the landmarks selected and pseudoinverse computed, the approximate self-attention output is:
This is Equation 17. In the algorithm pipeline (Algorithm 1), this is expressed as three named matrices:
- — the left factor
- — the pseudoinverse of the landmark kernel
- — the right factor
- Output:
The associative multiplication order matters critically: The product must be computed in an order that avoids forming an intermediate. The correct order is:
- Compute first — multiplying the pseudoinverse by the right factor, costing .
- Multiply the result by : , costing .
- Multiply by : , costing .
If instead one computed , the intermediate would be , incurring cost and defeating the purpose. The associative reordering is the standard trick for making linearized attention efficient, and Nyströmformer uses it to ensure the entire pipeline stays sub-quadratic.
Skip connection via depthwise convolution: The paper adds a skip connection of the value matrix to the approximate attention output, implemented as a 1D depthwise convolution with kernel size (see Figure 4). The architecture diagram shows this as a parallel path: passes through a "DConv " block (depthwise convolution) and is added to the Nyström attention output. A depthwise convolution with kernel size 1 along the sequence dimension applies a learned scalar weight to each feature dimension independently — it's essentially a learned scaling of each value dimension, providing the model with a way to bypass the attention approximation when needed. This skip connection "helps the training" by ensuring that even if the Nyström approximation is poor during early training, there is still a direct path from to the output.
Figure 4 architecture walkthrough: Following the diagram left to right:
- Input flows through three learned projections to produce (the paper uses for query dimension), , and .
- sMEANS (Segment-means) down samples and to and .
- The three kernel matrices are computed: produces an matrix (landmark kernel), which goes through pINV (iterative pseudoinverse) to produce the pseudoinverse. In parallel, produces (left factor), and produces (right factor).
- The three matrices are multiplied: produces the attention output.
- The skip connection path takes , passes it through DConv (depthwise convolution), and adds it to the attention output.
- The final output is produced.
Complexity Analysis: Verifying O(n) Scaling
The paper provides a detailed breakdown of the time and memory complexity for the Nyströmformer attention layer.
Time complexity:
| Operation | Complexity |
|---|---|
| Landmark selection (Segment-means) | |
| Compute | |
| Compute | |
| Compute | |
| Iterative pseudoinverse (6 iterations) | |
| Matrix multiplication chain | |
| Total |
When (e.g., while can be thousands), the dominant terms are and , both linear in . The pseudoinverse cost is constant with respect to sequence length. The overall complexity is .
Memory complexity:
| Component | Memory |
|---|---|
| Landmark matrices , | |
| Three kernel matrices , , | |
| Value output | |
| Total |
Again, with , the memory scales as , which is .
Comparison with standard self-attention: Standard self-attention requires time (computing ) and memory (storing the attention matrix). Nyströmformer replaces the terms with terms, where is a small constant. For and , this is a reduction factor of approximately in the attention computation.
Why the complexity analysis is valid: The paper explicitly accounts for all steps, including the pseudoinverse () and landmark selection (), making the analysis comprehensive rather than cherry-picked. The condition is satisfied for all practical sequence lengths — even for (the standard BERT sequence length), satisfies .
Theoretical Analysis: When Does the Approximation Converge?
The paper provides a theoretical result (Lemma 2) that characterizes when the Nyström approximation exactly recovers the true self-attention.
Lemma 2: Given the input datasets and , and the corresponding landmark point sets and , the Nyström approximate self-attention (Equation 17) converges to true self-attention if there exist landmark points and such that and for all and .
What this lemma states: If every original query appears as a landmark query and every original key appears as a landmark key — meaning the landmark set includes every token in the sequence () — then the Nyström approximation is exact. This is a trivial condition (it essentially says "if you don't reduce dimensionality, the approximation is perfect"), but it provides a theoretical sanity check: the approximation error comes entirely from the gap between the landmark set and the full token set. As approaches , the approximation approaches exactness.
Why this matters practically: The lemma implies that the approximation quality depends on how well the landmark set "covers" the query and key spaces. If the landmarks are representative — capturing the diversity of token representations — then even with , the approximation can be accurate. The authors connect this to the low-rank property of self-attention:
"the error of Nyström approximation depends on the spectrum of the matrix to be approximated and it decreases with the rank of the matrix. When this result is compared with the observation in (Wang et al. 2020) where the authors suggest that self-attention is low-rank, stronger guarantees based on structural properties of the matrix that we wish to approximate are possible."
The low-rank argument explained: If the true attention matrix has effective rank (meaning its singular values decay rapidly after the -th), then the Nyström approximation with landmarks has error bounded by the sum of the to singular values. If , the error is small. The empirical observation that self-attention matrices are approximately low-rank (from Linformer and other work) therefore provides theoretical justification for why a modest number of landmarks (64) can achieve good approximation — the attention matrix's information content is far less than , so compressing to dimensions loses little.
The Segment-means heuristic and the theoretical condition: The lemma requires that landmarks "overlap" with actual tokens, which Segment-means does not directly satisfy (landmarks are averages, not actual token embeddings). However, if the query/key representations are locally smooth along the sequence dimension — which is plausible in NLP due to local context — then the average of a segment is close to the individual tokens in that segment. The approximation quality thus depends on the locality of the representations, which is a property that emerges during training and can be learned.
Figure 3 provides a qualitative sanity check: a side-by-side comparison of the true self-attention matrix (top) and its Nyström approximation (bottom) on a real example, showing that the attention patterns are "quite similar." While this is a single example, it visually demonstrates that the approximation captures the essential structure — tokens that attend strongly in the true attention also attend strongly in the approximation, and the sparse/diffuse patterns are preserved.
4. Key Insights and Innovations
Innovation 1: Pre-Softmax Landmark Selection as a Distinct Class of Attention Approximation
The fundamental intellectual move in Nyströmformer is not simply applying the Nyström method to self-attention — it is recognizing where in the computation pipeline the Nyström sampling must occur to be useful, and accepting the compromise that this requires. This creates a previously unexplored category of attention approximation: data-dependent low-rank reconstruction that preserves the softmax operator by moving landmark selection before the nonlinearity.
Before this work, the efficient-attention landscape organized around two clean theoretical strategies. One approach was to change the attention function — replacing softmax with a kernel that factorizes (Katharopoulos et al., 2020; Choromanski et al., 2020), sacrificing the sharp selectivity that makes softmax attention effective. The other approach was to constrain or compress the attention pattern — using sparsity (Child et al., 2019; Beltagy, Peters, and Cohan, 2020), hashing (Kitaev, Kaiser, and Levskaya, 2019), or learned low-rank projections (Wang et al., 2020) — which either restrict which tokens can interact or compress dimensions agnostically via fixed projection matrices.
Nyströmformer refuses both horns of this dilemma. It does not change the softmax (the operator remains genuine row-wise softmax, applied to landmark sets). It does not impose a sparsity pattern or learned compression (every token pair receives a reconstructed attention weight, and the compression is data-dependent per input). Instead, it introduces a third category: matrix approximation applied to the attention computation graph itself, where the approximation point is shifted to avoid the quadratic bottleneck while preserving the mathematical form of the original operation.
The distinction from Linformer (Wang et al., 2020) is particularly instructive. Both methods produce O(n) attention by compressing along the sequence-length dimension — Linformer projects keys and values from to via learned projection matrices, while Nyströmformer samples landmark tokens via Segment-means and reconstructs the full attention matrix. The critical difference is adaptivity: Linformer's projection matrices are learned during training and fixed at inference time, meaning the same compression is applied regardless of input content. Nyströmformer's landmarks are computed from the specific query and key vectors of each input sequence — the approximation adapts to what the model is actually attending to. This is conceptually significant because it means Nyströmformer's approximation error is input-dependent in a potentially favorable way: easy-to-approximate attention patterns (e.g., highly focused attention) get better approximations than diffuse ones, whereas Linformer's fixed projection dimension imposes the same compression budget uniformly.
The pre-softmax compromise is what the paper explicitly frames as a "compromise" relative to the mathematically purer post-softmax Nyström sampling. But this compromise is what enables the method to be practical — and the paper's insight is that the compromise is well-founded because modern GPU architectures make the resulting operations (matrix multiplications, softmax over small dimensions, iterative pseudoinverse) efficiently implementable, whereas the "pure" approach (computing full even to sample a subset) is both theoretically and practically dead on arrival. The paper is essentially arguing that the right approximation point in the computation graph is more important than the right approximation method in the abstract — a design principle that generalizes beyond Nyström to any future effort to deploy classical numerical linear algebra inside deep learning architectures.
The evidence supporting this category distinction is architectural rather than metric-based: Figure 4 and the surrounding derivation show that every operation in the Nyströmformer pipeline maps cleanly to standard deep learning primitives (matrix multiply, softmax, average pooling, convolution), meaning the method is not just theoretically O(n) but practically implementable without custom CUDA kernels. This is a non-trivial engineering property — several prior efficient attention methods required specialized implementations that complicated adoption. The paper's contribution is as much about where to approximate (before softmax, using landmarks) as it is about how to approximate (Nyström quadrature), and the design space opened by this distinction — data-dependent reconstruction of the full softmax matrix from pre-softmax samples — had no occupants before this work.
Innovation 2: Segment-Means as a Deliberately Position-Aware Alternative to Clustering
On the surface, the choice of Segment-means for landmark selection appears to be a minor implementation detail — a simple averaging of contiguous blocks that the paper introduces almost apologetically as a practical alternative to K-means clustering. In fact, this choice embodies a subtle but significant conceptual insight: in the context of sequential data, position-based landmark selection may be superior to content-based clustering because it guarantees uniform coverage of the sequence, preventing the pathological failure mode where all landmarks concentrate in one region of the input.
Prior work on Nyström methods for kernel approximation consistently emphasized clustering-based landmark selection. K-means Nyström (Zhang, Tsang, and Kwok, 2008), recursive sampling (Musco and Musco, 2017), and adaptive sampling schemes (Fanuel, Schreurs, and Suykens, 2019) all select landmarks to minimize reconstruction error in the kernel matrix — which, for a general kernel, means selecting points that are spread out in the feature space induced by the kernel. The implicit assumption is that feature-space coverage is what matters for approximation quality.
In the self-attention setting, however, there are two distinct spaces where "coverage" could be defined: the embedding space (where queries and keys live as vectors in ) and the sequence position space (the 1D ordering of tokens). Segment-means operates in position space, not embedding space — it partitions tokens by their index in the sequence, not by their semantic similarity. This is a unusual choice that prior Nyström literature would consider suboptimal, since averaging semantically dissimilar tokens could produce a landmark that represents none of them well.
Why might position-based selection work despite this? The paper does not make this argument explicitly, but the experimental results (Figure 3's qualitative comparison showing similar attention patterns, and the competitive downstream performance in Tables 2 and 3) suggest that position-based coverage matters more for self-attention approximation than feature-space coverage. The reason is likely structural: self-attention matrices in Transformer models exhibit strong locality bias — tokens attend more strongly to nearby tokens than to distant ones, even in models without explicit positional attention constraints. By segmenting along the sequence dimension, Segment-means ensures that every region of the sequence has a landmark nearby in position space, which means the Nyström approximation can capture local attention patterns accurately. A content-based clustering might concentrate landmarks in semantically dense regions (e.g., noun phrases, named entities) while leaving other positions poorly represented — and those poorly-represented positions would have inaccurate attention reconstructions.
This is a reframing of what "good landmark selection" means in the self-attention context. Rather than landmarks needing to be representative of the distribution of token embeddings, they need to provide good coverage of the attention matrix structure, which is shaped by both content and position. Segment-means exploits the fact that position is a strong prior for attention patterns, providing a cheap proxy for coverage that works well in practice.
The computational argument is equally significant. K-means with clusters on points of dimension costs per forward pass — and more problematically, the cluster assignments are non-differentiable, creating complications for backpropagation. Segment-means costs with a trivial gradient (it's just average pooling), integrating seamlessly into the autograd computation graph without custom backward kernels. This is not merely an implementation convenience; it reflects a design philosophy that the approximation method should be a native operation in the deep learning computation graph, not an external algorithm bolted on. The fact that such a simple scheme works competitively (the paper uses it for all experiments, including the strong LRA benchmark results in Table 3) is a finding in itself — it suggests that sophisticated landmark selection algorithms may be solving a harder problem than necessary for this application.
Table 1 provides indirect evidence for the effectiveness of this design: Nyströmformer-32 (using only 32 landmarks, or roughly 1/250 of the tokens at sequence length 8192) still achieves 26.7× memory reduction and 13.4× speedup over standard self-attention, suggesting that even very coarse position-based landmarks preserve sufficient attention structure. The fact that the paper never needs to revert to more complex landmark selection schemes across its entire experimental suite strengthens the argument that Segment-means is not a stopgap but a genuinely well-matched solution for the problem.
Innovation 3: Iterative Pseudoinverse as a GPU-Native Alternative to SVD in Learned Pipelines
The paper's use of an iterative approximation for the Moore-Penrose pseudoinverse (Lemma 1, from Razavi et al., 2014) is easy to overlook as a minor numerical trick, but it represents a deliberate architectural decision with implications for how classical matrix algorithms should be integrated into differentiable deep learning pipelines. The core insight is that for operations that must be computed in every forward pass and differentiated through during backpropagation, convergence rate and GPU compatibility matter more than exactness — and that high-order iterative methods can satisfy both requirements simultaneously.
Standard approaches to computing the pseudoinverse — singular value decomposition (SVD) or QR-based methods — have two problems in this context. First, they are not efficiently parallelized on GPU architectures: SVD involves bidiagonalization via Householder reflections, which requires sequential operations across the matrix dimensions, underutilizing the GPU's parallel throughput. Second, backpropagating through SVD is numerically problematic because the gradients of singular vectors become ill-conditioned when singular values are close together (the "singular value proximity" problem), requiring careful handling that complicates the autograd implementation.
The iterative scheme in Lemma 1 sidesteps both problems. Each iteration is composed entirely of matrix-matrix multiplications — the single most optimized operation on GPUs, achieving near-peak theoretical throughput. The third-order convergence (error reduces cubically per iteration) means that only ~6 iterations are needed for a good approximation, keeping the constant factor small. And because the computation is just a fixed polynomial of matrix products, backpropagation is straightforward — automatic differentiation traces through the multiplications without any special-case handling.
This is not a contribution to numerical linear algebra (the algorithm is from prior work) but rather a contribution to the engineering methodology for deploying classical matrix algorithms inside deep learning architectures. It establishes a template: when you need a matrix operation in the forward pass, prefer iterative methods with fast convergence and GPU-friendly primitives over exact methods with poor hardware utilization and tricky gradients. The paper demonstrates that the approximation quality loss is negligible (the results in Tables 2 and 3 show no degradation from the pseudoinverse approximation), while the implementation simplicity gain is substantial.
The initialization scheme () is part of this contribution. The choice is specifically engineered to guarantee , which is exactly the condition needed for Lemma 1's convergence when is non-singular. For singular (which can occur if landmarks are poorly chosen or collinear), the paper notes that the initialization "provides a good approximation in our experiments" — an empirical claim rather than a theoretical guarantee, but one that holds across the evaluated benchmarks. This pragmatic attitude toward numerical edge cases (reasonable initialization + empirical validation rather than theoretical worst-case guarantees) is characteristic of the paper's approach and appropriate for the machine learning context, where exact pseudoinverses are not required and robustness to singular cases can be validated through the diversity of training data.
The broader significance is that this design pattern enables other classical matrix approximation methods to be imported into deep learning. The Nyström method is one instance; CUR decomposition, interpolative decomposition, and randomized SVD all have similar structures (sample columns/rows, compute a small pseudoinverse, reconstruct) and could potentially be adapted following the same template if the iterative pseudoinverse proves reliable. The paper's demonstration that 6 iterations suffice for transformer-scale problems provides a practical calibration point for future efforts.
Innovation 4: The Skip Connection as a Mechanism for Approximation Robustness During Training
The addition of a depthwise convolution skip connection of the value matrix (Figure 4, "DConv ") appears minor — it is mentioned in a single sentence. But it embodies an important design principle that distinguishes Nyströmformer from other efficient attention methods: providing a learned bypass path that allows the model to compensate for approximation error during training, preventing the approximation from becoming a bottleneck that stalls optimization.
To understand why this matters, consider what happens during early training when the Nyström approximation is poor. The landmarks are computed from randomly initialized query and key representations via Segment-means — at initialization, these representations are essentially random, so the landmarks are random averages of random vectors. The resulting Nyström attention matrix bears little resemblance to any meaningful attention pattern. If this were the only path from to the output, the gradient signal through the attention layer would be noisy and uninformative, potentially preventing the query and key projections from learning useful representations — a classic chicken-and-egg problem where the approximation needs good representations to be accurate, but the representations need good gradients to improve.
The depthwise convolution skip connection (kernel size 1 along the sequence dimension, meaning an independent scalar weight per feature dimension applied uniformly to all positions) provides a direct, approximation-free path: . At initialization, the model can learn to rely primarily on the skip connection (by learning larger convolution weights) while the attention approximation improves. As training progresses and the representations become structured, the model can shift weight toward the Nyström attention path (by learning larger attention-output projection weights). This is analogous to how residual connections in standard Transformers allow each layer to learn the identity function initially and gradually introduce nonlinear processing — the skip connection makes the approximation optional rather than forced, giving the optimization a smoother landscape.
Most prior efficient attention methods do not include an explicit mechanism for this kind of graceful degradation during training. Sparse attention methods (Child et al., 2019; Beltagy, Peters, and Cohan, 2020) fix the sparsity pattern a priori — if the pattern discards important interactions, the model cannot recover them regardless of training stage. Kernel-based linearization methods (Katharopoulos et al., 2020) replace softmax entirely — if the replacement function is a poor fit early in training, there is no fallback. Linformer's learned projections must compress effectively from the start, since the projection dimension is fixed. Nyströmformer's skip connection is a simple but effective insurance policy: if the Nyström approximation is inaccurate, the model can route around it temporarily and still receive usable gradients to improve the representations that will eventually make the approximation accurate.
The paper does not ablate this skip connection or analyze its training dynamics, which is a limitation — without an ablation, we cannot quantify how much it matters empirically. But the design choice reflects a principled understanding of the interaction between approximation quality and optimization dynamics that is underappreciated in the efficient-attention literature. The fact that the paper considered it worth including (and depicting explicitly in Figure 4) suggests the authors found it practically important during development, even if they did not isolate its effect for publication. This is an architectural pattern — providing a learned bypass for any approximate module — that generalizes beyond attention to any setting where an expensive exact computation is replaced by a cheaper approximation during training.
The evidence for the overall training stability is indirect but consistent across experiments: Figure 5 shows smooth, monotonic improvement in MLM and SOP validation accuracy for Nyströmformer on BERT-small, with no training instability or divergence; Figure 6 shows Nyströmformer training from scratch tracking BERT-base closely; and the downstream fine-tuning results in Table 2 show no catastrophic failures on any task. While this does not isolate the skip connection's contribution, it establishes that the full system trains reliably, which is the property the skip connection is designed to ensure.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper evaluates on three distinct suites: (1) language model pretraining on English Wikipedia plus BookCorpus (Zhu et al., 2015), split 80/20 into training and validation sets; (2) downstream NLP tasks from the GLUE benchmark (Wang et al., 2018) — specifically SST-2 (Socher et al., 2013), MRPC (Dolan and Brockett, 2005), QNLI (Rajpurkar et al., 2016), QQP (Chen et al., 2018), and MNLI (Williams, Nangia, and Bowman, 2018) — plus IMDB reviews (Maas et al., 2011); and (3) the Long Range Arena (LRA) benchmark (Tay et al., 2020), covering ListOps (2K length), byte-level text classification (4K), byte-level document retrieval (4K), image classification on pixel sequences (1K), and Pathfinder (1K). The LRA benchmark uses the standard train/test splits from Tay et al. (2020).
-
Base model(s). Two BERT variants are used throughout: BERT-small (4 layers) for comparing against other linear-complexity attention methods on language modeling efficiency, and BERT-base (Devlin et al., 2019) as the primary baseline for downstream task evaluation. Nyströmformer is obtained by replacing the standard self-attention in both BERT variants with the proposed Nyström approximation. The paper argues that BERT-base is representative of the model scale commonly used in transfer learning evaluations, making it a relevant baseline. For the LRA benchmark, the paper uses a uniform architecture across all compared methods: 2 layers, 64 embedding dimension, 128 hidden dimension, 2 attention heads, with mean pooling — this is the standard configuration from Tay et al. (2020).
-
Metrics. For pretraining, the paper reports masked-language-modeling (MLM) accuracy and sentence-order-prediction (SOP) accuracy on the held-out validation set, plotted against training steps. For GLUE tasks, it reports F1 score for MRPC and QQP, and accuracy for all other tasks (SST-2, QNLI, MNLI matched/mismatched, IMDB), following the standard GLUE evaluation protocol. For the LRA benchmark, classification accuracy is reported for each individual task along with the average accuracy across all five tasks. Memory consumption is reported in megabytes (MB) and inference time in milliseconds (ms) per input instance, measured on a single Nvidia 1080Ti GPU.
-
Baselines. The paper compares against four categories of methods. For pretraining and downstream tasks, the primary baseline is BERT-base (Devlin et al., 2019) with standard quadratic self-attention; additionally, on the BERT-small comparison, the paper evaluates ELU linearized self-attention (Katharopoulos et al., 2020) and Linformer (Wang et al., 2020). For runtime/memory profiling (Table 1), the paper compares against Linformer (with projection dimension 256) and Longformer (Beltagy, Peters, and Cohan, 2020) (with sliding window size 257). For the LRA benchmark (Table 3), the baselines are: the standard Transformer (Vaswani et al., 2017), Reformer (Kitaev, Kaiser, and Levskaya, 2019) (2 hashes), Linformer (projection dimension 256), and Performer (Choromanski et al., 2020) (random feature dimension 256).
-
Generation budget / compute accounting. For runtime and memory comparisons (Table 1), the paper evaluates all self-attention modules on identical hardware (single Nvidia 1080Ti) using random input tensors of varying sequence lengths n ∈ {512, 1024, 2048, 4096, 8192}, reporting average memory consumption and inference time per input instance through the self-attention module only. For pretraining comparisons, compute is implicitly compared via training step count — BERT-base and Nyströmformer are both trained for 0.5M steps from scratch (or Nyströmformer is initialized from pretrained BERT-base and trained for ~0.25M steps for accelerated convergence). The LRA benchmark uses identical architectures across all methods, so the compute metric is accuracy under the same model capacity. There is no explicit FLOPs or generation-budget accounting in this paper — comparisons are based on wall-clock time, memory usage, and downstream accuracy rather than a normalized compute metric like "number of attention operations."
-
Cross-validation / statistical protocol. No cross-validation or statistical significance testing is reported. For GLUE tasks, the paper follows the standard protocol of fine-tuning on the training set and evaluating on the validation set. For the smaller MRPC dataset, a hyperparameter sweep over batch size {8, 16, 32} and learning rate {2e-5, 3e-5, 4e-5, 5e-5} is performed, selecting the best validation result — this is the only instance of hyperparameter selection on validation data. The LRA benchmark uses the standard train/test splits from Tay et al. (2020) without modification. The absence of multiple random seeds, standard deviations, or cross-validation folds means that all reported numbers should be interpreted as point estimates without formal uncertainty quantification.
Main Quantitative Results
Pretraining Language Modeling: Nyströmformer Matches Standard Self-Attention on MLM and SOP
The core pretraining result (Figure 5, BERT-small comparison) demonstrates that Nyströmformer achieves competitive MLM and SOP validation accuracy relative to standard self-attention when controlling for model size and training data. On BERT-small trained for 0.1M steps, the validation curves in Figure 5 show that Nyströmformer tracks standard self-attention closely in MLM accuracy throughout training, while outperforming both Linformer and ELU-based linearized attention. The SOP accuracy curves show similar behavior, with Nyströmformer and standard self-attention producing nearly overlapping trajectories while Linformer and ELU-linearized attention demonstrate noticeably lower accuracy.
For the BERT-base scale (Figure 6), two training configurations are evaluated: Nyströmformer trained from scratch for 0.5M steps, and Nyströmformer initialized from pretrained BERT-base and trained for ~0.25M steps. The from-scratch Nyströmformer matches the BERT-base from-scratch curve closely in both MLM and SOP validation accuracy across training steps, with no significant divergence. The BERT-initialized Nyströmformer converges faster initially (due to the warm start) and reaches comparable final accuracy, confirming that Nyströmformer can adopt pretrained standard Transformer weights and fine-tune them without degradation — a property important for practical adoption since it allows leveraging existing pretrained checkpoints.
The significance of these pretraining results is not absolute performance (no specific final accuracy numbers are reported in the text) but rather the demonstration that replacing quadratic self-attention with the Nyström approximation does not impair the model's ability to learn language representations. This is a necessary condition for downstream transfer: if pretraining accuracy degraded substantially, the downstream fine-tuning results would be compromised regardless of efficiency gains. The experiments establish that this necessary condition holds across two model scales (BERT-small and BERT-base).
Runtime and Memory Efficiency: Linear Scaling with Sequence Length, Confirming O(n) Complexity
Table 1 presents the paper's primary efficiency evidence. The comparison evaluates the self-attention module in isolation (not the full Transformer layer) for standard self-attention, Linformer (projection dimension 256), Longformer (sliding window size 257), Nyströmformer-64 (64 landmarks), and Nyströmformer-32 (32 landmarks) across sequence lengths from 512 to 8192.
At sequence length n = 8192, standard self-attention consumes 10,233 MB of memory and takes 155.4 ms per instance. Nyströmformer-64 reduces memory to 450 MB (22.8× reduction) and time to 12.3 ms (12.7× speedup). Nyströmformer-32 achieves even greater savings: 383 MB (26.7× reduction) and 11.5 ms (13.4× speedup). These improvements increase with sequence length — at n = 512, the reduction factors are modest (1.5× memory, 1.1× speedup for Nyströmformer-64), but at n = 8192 they become dramatic, confirming the linear scaling behavior: as n grows, the gap between O(n^2) and O(n) widens multiplicatively.
Relative to other efficient attention methods, Nyströmformer demonstrates competitive or superior efficiency. At n = 8192, Nyströmformer-64 provides 1.7× memory savings over Linformer (450 MB vs. 635 MB) with comparable running time (12.3 ms vs. 11.3 ms), and 1.2× memory savings over Longformer (450 MB vs. 455 MB) with 3× speedup (12.3 ms vs. 36.2 ms). The advantage over Longformer is particularly notable because Longformer is designed for long-document processing — Nyströmformer outperforms it in both memory and speed at scale without requiring task-specific global attention patterns.
At shorter sequence lengths (n = 512), Longformer shows slightly better memory efficiency (32.2 MB vs. 35 MB for Nyströmformer-64), but this reverses by n = 1024 and the gap widens in Nyströmformer's favor as length increases. Linformer maintains comparable speed to Nyströmformer across all lengths but consistently uses more memory — the gap grows from 41 MB vs. 35 MB at n = 512 to 635 MB vs. 450 MB at n = 8192, suggesting that Linformer's learned projection matrices incur a memory overhead that scales with the projection dimension (256 in this case) while Nyströmformer's segment-based landmarks are more memory-efficient.
The 32-landmark variant (Nyströmformer-32) is the most efficient configuration tested, providing the maximum memory and speed savings. This is significant because it demonstrates that even aggressive compression (32 landmarks for 8192 tokens = roughly 1/256 of the sequence length) preserves functional self-attention behavior — the paper would not report strong downstream results if 32 landmarks caused unacceptable attention approximation error.
Downstream NLP Tasks: Nyströmformer Matches BERT-base on Standard-Sequence Tasks
Table 2 reports fine-tuned performance on six downstream tasks after pretraining with masked language modeling and sentence order prediction. The comparison is between BERT-base (standard self-attention) and Nyströmformer (with Nyström self-attention, pretrained and fine-tuned at sequence length 512):
| Task | BERT-base | Nyströmformer |
|---|---|---|
| SST-2 | 90.0 | 91.4 |
| MRPC | 88.4 | 88.1 |
| QNLI | 90.3 | 88.7 |
| QQP | 87.3 | 86.3 |
| MNLI m/mm | 82.4/82.4 | 80.9/82.2 |
| IMDB | 93.3 | 93.2 |
Across all six tasks, Nyströmformer performs competitively with BERT-base. On SST-2, Nyströmformer slightly outperforms BERT-base (91.4 vs. 90.0, a 1.4-point improvement). On IMDB, the models are nearly identical (93.2 vs. 93.3). On MRPC, the gap is minimal (88.1 vs. 88.4). The largest differences appear on QNLI (88.7 vs. 90.3, a 1.6-point gap) and MNLI matched (80.9 vs. 82.4, a 1.5-point gap), where Nyströmformer underperforms BERT-base modestly. On MNLI mismatched, the scores are essentially identical (82.2 vs. 82.4). On QQP, Nyströmformer trails by 1.0 point (86.3 vs. 87.3).
The paper also reports that fine-tuning with longer sequences (n = 1024) yields "almost identical" results to n = 512 on IMDB reviews (93.0 vs. 93.2 accuracy), suggesting that the model's accuracy does not degrade when the sequence length is doubled beyond the pretraining length.
These results support the claim that Nyströmformer performs comparably to standard self-attention on standard-length sequence tasks. The worst-case gap is 1.6 points (QNLI), and on two tasks (SST-2, IMDB) the difference is negligible or favorable. The pattern suggests that tasks emphasizing local or short-range dependencies (sentiment classification in SST-2 and IMDB) are least affected by the attention approximation, while tasks requiring more complex cross-sentence reasoning (natural language inference in MNLI, question-answer pair matching in QNLI) show slightly larger but still modest degradations.
Long Range Arena (LRA) Benchmark: Nyströmformer Outperforms Other Efficient Attention Methods
Table 3 presents results on the LRA benchmark, which is specifically designed to test Transformer models on tasks requiring long-range dependencies. All methods use an identical 2-layer Transformer architecture with 64 embedding dimension, 128 hidden dimension, and 2 attention heads, varying only the self-attention mechanism.
| Model | ListOps (2K) | Text (4K) | Retrieval (4K) | Image (1K) | Pathfinder (1K) | Average |
|---|---|---|---|---|---|---|
| Standard | 37.10 | 65.02 | 79.35 | 38.20 | 74.16 | 58.77 |
| Reformer | 19.05 | 64.88 | 78.64 | 43.29 | 69.36 | 55.04 |
| Linformer | 37.25 | 55.91 | 79.37 | 37.84 | 67.60 | 55.59 |
| Performer | 18.80 | 63.81 | 78.62 | 37.07 | 69.87 | 53.63 |
| Nyströmformer | 37.15 | 65.52 | 79.56 | 41.58 | 70.94 | 58.95 |
Nyströmformer achieves the highest average accuracy (58.95%), slightly exceeding standard self-attention (58.77%, a +0.18 margin) and substantially outperforming all other efficient attention methods: +3.91 over Reformer (55.04), +3.36 over Linformer (55.59), and +5.32 over Performer (53.63).
The per-task breakdown reveals where Nyströmformer's advantages and disadvantages lie relative to alternatives. On ListOps (hierarchical list operations requiring syntactic reasoning over 2K-length sequences), Nyströmformer (37.15) matches both standard self-attention (37.10) and Linformer (37.25), while dramatically outperforming Reformer (19.05) and Performer (18.80) — methods based on hashing and random features respectively, which appear to fail catastrophically on this structured reasoning task. On Text classification (byte-level IMDb reviews at 4K length), Nyströmformer (65.52) slightly edges standard self-attention (65.02) and outperforms Linformer (55.91) by a wide margin (+9.61), suggesting that Linformer's fixed projection dimension may lose character-level information critical for byte-level text understanding. On Retrieval (document matching at 4K), all methods perform similarly (78.64–79.56 range), with Nyströmformer (79.56) slightly ahead — this task appears relatively insensitive to attention mechanism choice. On Image classification (pixel sequences at 1K), Reformer (43.29) leads, with Nyströmformer (41.58) second — the Nyström approximation trails the hashing-based method on this task but substantially outperforms standard self-attention (38.20), Linformer (37.84), and Performer (37.07). On Pathfinder (spatial reasoning at 1K), standard self-attention leads (74.16), with Nyströmformer (70.94) second, ahead of Performer (69.87), Reformer (69.36), and Linformer (67.60).
The paper highlights that the PyTorch reimplementation produces "higher" results on the Retrieval task for all models compared to the original Jax/Flax LRA benchmark — indicating a systematic implementation difference rather than a Nyströmformer-specific effect. This is relevant for reproducibility: researchers comparing against these numbers should use the paper's PyTorch baseline, not the original LRA reported numbers.
The key takeaway is that Nyströmformer is the only efficient attention method that matches or exceeds standard self-attention's average LRA performance while providing linear complexity. Reformer, Linformer, and Performer all sacrifice average accuracy (by 3.2–5.1 points) for efficiency; Nyströmformer does not, at least on this benchmark under the tested configuration.
Interaction of Efficiency and Accuracy Across Configurations
By combining the efficiency data (Table 1) with the downstream results (Tables 2 and 3), a consistent trade-off pattern emerges. Nyströmformer-64 provides 22.8× memory reduction and 12.7× speedup at n = 8192 while matching BERT-base accuracy on GLUE tasks (average gap under 1 point) and slightly exceeding standard self-attention on LRA (+0.18 average). Nyströmformer-32 provides even larger efficiency gains (26.7× memory, 13.4× speedup) but the paper does not report downstream accuracy specifically for the 32-landmark variant on GLUE or LRA, leaving open the question of whether the additional compression degrades task performance — Table 1 only profiles efficiency, not accuracy, for Nyströmformer-32. This is a notable gap in the evaluation.
Ablation Studies and Robustness Checks
The paper does not report extensive formal ablation experiments in the style of many modern empirical papers. The nearest equivalents are the following comparisons that isolate specific design choices:
Landmark count (32 vs. 64): Table 1 implicitly compares Nyströmformer-32 and Nyströmformer-64 on efficiency metrics only. At n = 8192, reducing landmarks from 64 to 32 improves memory efficiency from 22.8× to 26.7× over standard self-attention and improves speedup from 12.7× to 13.4×. However, no accuracy numbers are reported for Nyströmformer-32 on any downstream task. The paper states that "using 64 landmarks is often sufficient to ensure a good approximation, although this depends on the application," but provides no empirical characterization of how approximation quality degrades as landmark count decreases, nor any guidance on how to select m for a given task or sequence length. This is the most important missing ablation, since the trade-off between landmark count and accuracy is central to the method's practical utility.
Training from scratch vs. initialized from pretrained BERT-base: Figure 6 compares Nyströmformer trained from scratch for 0.5M steps against Nyströmformer initialized from pretrained BERT-base and trained for ~0.25M steps. Both converge to similar MLM and SOP validation accuracy, with the initialized version converging faster. This demonstrates that Nyströmformer's attention parameters (query, key, value projections) are sufficiently compatible with standard attention's parameter space that transfer learning from a standard pretrained checkpoint is effective — a practically important property for adoption.
Nyström vs. other linear self-attention methods on BERT-small: Figure 5 shows Nyströmformer outperforming both Linformer and ELU-based linearized attention on BERT-small pretraining. This is the closest the paper comes to an ablation testing whether the Nyström approximation specifically (rather than linear attention generally) is beneficial — it demonstrates that on this small-scale pretraining comparison, the Nyström approach is preferable to kernel-based linearization (ELU) and learned projection (Linformer).
LRA architecture consistency across methods: The LRA evaluation (Table 3) uses identical model architectures (2 layers, 64 embedding dim, 128 hidden dim, 2 heads) for all attention mechanisms, making it a controlled comparison. The fact that Nyströmformer is the only efficient method to match standard attention's average accuracy suggests that its approximation preserves attention quality in a way that the other methods' modifications (hashing, random features, fixed projections) do not, at least under this specific architecture configuration.
Segment-means landmark selection: No ablation compares Segment-means against alternative landmark selection methods (K-means, random sampling, learned landmarks). The paper explicitly acknowledges K-means as a standard approach and dismisses it as "less desirable during mini-batch training" without empirical validation of this claim. This is a notable missing ablation — it is possible that K-means landmarks (computed periodically rather than per-batch, or using a moving average of cluster centers) could improve approximation quality at modest additional cost, but this hypothesis is untested.
Iterative pseudoinverse convergence: No ablation tests the effect of the number of pseudoinverse iterations on downstream performance. The paper states that "about 6 iterations" are used across all experiments but provides no evidence that this is sufficient or that fewer iterations would degrade performance. Given that the pseudoinverse computation costs O(m^3) and accounts for a non-trivial fraction of the per-layer computation (especially for larger m), understanding the iteration-accuracy trade-off would be practically valuable.
Skip connection ablation: The depthwise convolution skip connection is mentioned as helping training but is never ablated. The paper provides no comparison showing performance with and without the skip connection, making it impossible to determine whether it is essential, moderately helpful, or irrelevant.
Sequence length effect on downstream accuracy: The paper mentions that fine-tuning with n = 1024 on IMDB yields "almost identical" results to n = 512 (93.0 vs. 93.2), but no systematic study of how downstream accuracy changes with sequence length is reported. This is relevant because the method's O(n) complexity is most valuable at long sequences — if accuracy degrades at lengths where the efficiency gain is greatest, the practical benefit is undermined.
Critical Assessment
The experiments tell a clear and broadly consistent story: Nyströmformer replaces quadratic self-attention with a linear-complexity approximation, achieves dramatic efficiency gains at long sequence lengths (22.8–26.7× memory reduction at n = 8192), and maintains competitive downstream performance across standard benchmarks (GLUE, IMDB) and long-range tasks (LRA). However, the experimental design has several structural limitations that warrant careful examination.
The central claim — that Nyströmformer achieves O(n) complexity while performing comparably to standard self-attention — is supported, but with meaningful scope limitations. The efficiency evidence (Table 1) is strong: memory and time measurements at multiple sequence lengths on identical hardware clearly demonstrate linear scaling, and the comparison against Linformer and Longformer confirms that Nyströmformer is competitive with (and often better than) other efficient attention methods in resource utilization. The accuracy evidence on standard-length tasks (Table 2) shows Nyströmformer within ~1.5 points of BERT-base across all GLUE tasks, which is a genuine achievement — it means the Nyström approximation does not substantially degrade performance on the types of tasks the community uses to benchmark Transformer models. The LRA evidence (Table 3) extends this to longer sequences, showing that Nyströmformer is the only efficient attention method in the comparison that preserves standard-attention-level accuracy on average.
However, the paper does not demonstrate that Nyströmformer scales to the very long sequences where its efficiency advantage is most dramatic. Table 1 shows efficiency at n = 8192, but all accuracy evaluations are at n ≤ 4096 (LRA tasks) or n = 512 (GLUE tasks). The paper does not report accuracy on any task at n = 8192 or beyond. This is a significant gap: the method's primary selling point is enabling Transformers to process sequences of length 8000+, but we have no evidence that the Nyström approximation maintains fidelity at those lengths. The information lost by compressing 8192 tokens into 64 landmark averages could be substantial, and without long-sequence accuracy measurements, the paper's headline efficiency numbers at n = 8192 are demonstrating a capability whose practical value is unvalidated.
The comparison against BERT-base on GLUE (Table 2) is a necessary but insufficient evaluation. The tasks (SST-2, MRPC, QNLI, QQP, MNLI, IMDB) are all standard-length benchmarks where quadratic self-attention is not a bottleneck — sequence lengths rarely exceed a few hundred tokens, and BERT-base already handles them easily. The value of an O(n) attention mechanism on these tasks is primarily about computational efficiency during training and inference, not about enabling longer inputs. The paper's results show that Nyströmformer doesn't hurt accuracy on these tasks, which is good, but they don't show that Nyströmformer enables something quadratic attention cannot do. To demonstrate that the method unlocks genuinely new capabilities, the paper would need to evaluate on tasks with very long sequences (n = 4096+) where standard self-attention is prohibitively expensive or impossible, and show that Nyströmformer achieves accuracy that would be unattainable without linear scaling. The LRA benchmark partially addresses this (including tasks at 4K length), but LRA uses a shallow 2-layer model — it doesn't demonstrate that the approximation works in deep models (BERT-base has 12 layers) on long sequences, which is the regime where the O(n) complexity matters most.
The missing 32-landmark accuracy evaluation is a notable gap. The paper devotes a row of Table 1 to Nyströmformer-32, showing that it provides the best efficiency (26.7× memory reduction), but provides no corresponding accuracy numbers. The reader is left wondering whether the 32-landmark variant, which is the most efficient configuration tested, actually works for any downstream task. The single sentence "using 64 landmarks is often sufficient" implies but does not demonstrate that 32 may not be sufficient — yet without evidence, a practitioner cannot make an informed choice about the efficiency-accuracy trade-off.
The absence of any statistical rigor weakens the claims of "comparable" or "slightly better" performance. All results are reported as point estimates without standard deviations, confidence intervals, or multiple random seeds. The difference between 91.4 (Nyströmformer) and 90.0 (BERT-base) on SST-2 could be noise — SST-2 is known to have relatively high variance across runs, and a 1.4-point difference at this performance level may not be statistically significant. Similarly, the difference between Nyströmformer's 58.95 and standard attention's 58.77 on LRA average (a 0.18-point difference) is small enough that with only a single run, we cannot determine whether Nyströmformer is genuinely matching standard attention or whether both scores would converge with additional trials. The paper claims Nyströmformer "performs comparably, or in a few cases, even slightly better" — the "slightly better" cases (SST-2, LRA average) are not established as reliable with the reported experimental protocol.
The LRA reimplementation nuance is important but under-explored. The paper notes that its PyTorch LRA reimplementation produces higher scores on Retrieval for all models compared to the original Jax/Flax benchmark. This is reported as an aside but has implications for the claimed average accuracy advantage. If the systematic score inflation is uneven across tasks and methods, the 0.18-point "advantage" over standard attention on average could be an artifact of the reimplementation rather than a genuine property of the Nyström approximation. Re-running the original LRA implementation for Nyströmformer would clarify this, but the paper only reports its own reimplementation.
The lack of key ablations leaves the method's sensitivity to design choices uncharacterized. No ablation studies test:
- Number of landmarks vs. accuracy on a downstream task
- Segment-means vs. alternative landmark selection
- Number of pseudoinverse iterations vs. accuracy
- Skip connection presence vs. absence
- Performance scaling with sequence length on a fixed accuracy benchmark
Each of these would help practitioners use the method effectively and help researchers understand which components are essential versus incidental. The iterative pseudoinverse, in particular, is a non-trivial component whose convergence properties and iteration count might interact with training dynamics in unpredictable ways — without an ablation, we cannot assess whether the 6-iteration default is sufficient or conservative.
The single-model-family evaluation limits claims about representativeness. All experiments use BERT-base or BERT-small architectures with the standard English Wikipedia + BookCorpus pretraining corpus. The paper does not evaluate Nyströmformer on other model families (RoBERTa, GPT, T5), other pretraining objectives (causal LM, seq2seq), or other modalities (vision, speech). The claim that the method is "a step towards building Transformer models on very long sequences" is directionally supported but not validated across the architectural diversity of modern Transformers. Different pretraining objectives might produce attention patterns with different low-rank properties, affecting approximation quality.
The pretraining comparison in Figure 6 is cost-controlled but not resource-controlled. BERT-base from scratch is trained for 0.5M steps on 8 V100 GPUs taking "more than one week." Nyströmformer from scratch is also trained for 0.5M steps. The paper acknowledges this setup "to keep compute costs reasonable" relative to the standard 1M-step BERT training. However, Nyströmformer's O(n) complexity means each training step should be faster than BERT-base's — the paper does not report whether this efficiency advantage was realized during pretraining or whether the Nyströmformer overhead (pseudoinverse, Segment-means) offsets the attention savings at n = 512. If Nyströmformer pretraining is not meaningfully faster than standard BERT at n = 512, the practical benefit is limited to longer sequences — exactly the regime where accuracy is not evaluated.
What experiments would strengthen the paper? A systematic study of accuracy vs. sequence length for a fixed task (e.g., document classification with varying segment lengths) would directly test whether the O(n) complexity translates to usable long-sequence models. Evaluating Nyströmformer at n = 2048, 4096, and 8192 on a task where longer sequences improve accuracy would demonstrate that the efficiency gains enable capabilities that standard attention cannot provide. Ablating landmark count against accuracy would give practitioners concrete guidance for the m parameter. Evaluating on a generative or encoder-decoder architecture would test whether the Nyström approximation generalizes beyond the BERT-style encoder-only setting. And reporting multiple random seeds with standard deviations would transform the claims of "comparable" performance from qualitative observations into statistically grounded statements.
Despite these limitations, the experimental evidence does establish the paper's core technical contribution: the Nyström-based attention approximation is a viable method for reducing self-attention complexity from O(n^2) to O(n), and it preserves enough attention quality to match standard Transformers on the evaluated benchmarks. The dramatic efficiency numbers at n = 8192 (22.8× memory, 12.7× speedup) demonstrate that the theoretical complexity reduction translates to real hardware gains. The competitive downstream performance on GLUE and LRA demonstrates that the approximation does not catastrophically fail in practice. These are non-trivial achievements that justify the method's place in the efficient-attention landscape — even if the paper leaves important practical questions (optimal landmark count, long-sequence accuracy, statistical reliability) for future work.
6. Limitations and Trade-offs
Accuracy on Very Long Sequences (n > 4096) Is Entirely Unvalidated
The assumption or constraint. The paper's central motivation and headline efficiency results target very long sequences — the abstract promises "application to longer sequences with thousands of tokens," and Table 1 demonstrates 22.8× memory reduction and 12.7× speedup at sequence length n = 8192 relative to standard self-attention. The paper frames this as the primary value proposition: "The scalability of Nyströmformer enables application to longer sequences with thousands of tokens." However, every accuracy evaluation is conducted at standard or moderate sequence lengths: GLUE and IMDB tasks use n = 512 (the standard BERT maximum), and the Long Range Arena benchmark uses maximum lengths of 1K–4K (ListOps at 2K, Text and Retrieval at 4K, Image and Pathfinder at 1K). No accuracy measurement is reported for any task at n = 8192 or beyond.
The assumption is that the Nyström approximation quality observed at n ≤ 4096 extrapolates to n = 8192 and beyond. The paper provides no evidence for this extrapolation.
The consequence. The approximation error of the Nyström reconstruction depends on how well m landmarks (32 or 64) represent the full query and key matrices of dimension n × d_q. As n increases while m remains fixed, each Segment-means landmark averages over more tokens — at n = 8192 with m = 64, each landmark represents roughly 128 tokens. The approximation therefore becomes coarser, potentially missing fine-grained attention patterns that matter for task performance. More subtly, the low-rank assumption (that the attention matrix has effective rank ≤ m) becomes less likely to hold as n grows, since longer sequences can exhibit more complex dependency structures. If the effective rank of the attention matrix grows with sequence length, then fixed-m Nyström approximation error grows correspondingly, and the accuracy measured at n = 4096 on LRA may substantially overstate accuracy at n = 8192 on a real long-document task.
The practical consequence is that a practitioner considering Nyströmformer for a long-document application (legal contracts, scientific papers, book-length text) has no way to estimate whether the 22.8× memory savings at n = 8192 come at an unacceptable accuracy cost — or whether the approximation fails entirely past some sequence length threshold. The O(n) complexity is firmly established; the O(n) utility is not.
What evidence exists in the paper. Table 1 provides efficiency measurements from n = 512 to n = 8192, confirming linear scaling. Table 3 provides accuracy on LRA tasks at lengths of 1K–4K. The gap between these — accuracy at the lengths where efficiency is most dramatic — is unaddressed. The paper's one reference to longer sequence fine-tuning is a brief note: "we also experiment with fine-tuning our model using longer sequences (n = 1024), yet the results remain almost identical to n = 512, e.g., 93.0 vs. 93.2 accuracy on IMDB reviews." This is a single data point at a modest length increase (2×, not the 8–16× increase needed to reach the headline efficiency regime) on a single task. It does not characterize the scaling behavior of approximation quality with length.
Mitigation status. The paper does not acknowledge this as a limitation and proposes no mitigation. The natural mitigations — scaling m with n (which would degrade efficiency), evaluating accuracy at the lengths featured in Table 1, or providing theoretical error bounds as a function of n and m — are absent. The conclusion states that the method "is a step towards building Transformer models on very long sequences," which implicitly acknowledges that the step is incomplete, but the paper does not characterize what remains to be done.
Landmark Count vs. Accuracy Trade-off Is Uncharacterized
The assumption or constraint. The number of landmarks m is the single most important hyperparameter controlling the efficiency-accuracy trade-off in Nyströmformer. It directly determines the computational cost (the O(n m) terms dominate the O(n) scaling factor) and the approximation fidelity (more landmarks mean better coverage of the query/key spaces). The paper provides efficiency measurements for m = 32 and m = 64 (Table 1), demonstrating that 32 landmarks provide better efficiency (26.7× memory reduction vs. 22.8×, 13.4× speedup vs. 12.7× at n = 8192). The paper states that "using 64 landmarks is often sufficient to ensure a good approximation, although this depends on the application," but provides no accuracy measurements for any m other than 64.
The implicit assumption is that m = 64 is the appropriate default and that practitioners can trust this without empirical guidance. The paper offers no method for selecting m based on task characteristics, sequence length, or accuracy requirements.
The consequence. A practitioner deploying Nyströmformer faces an unguided trade-off. Selecting m = 32 provides the best efficiency but with unknown accuracy consequences — the paper's description of Nyströmformer-32 in Table 1 reports only memory and time, not task performance. It is possible that m = 32 causes substantial accuracy degradation on certain tasks (e.g., tasks requiring fine-grained token-level reasoning) while being perfectly adequate on others (e.g., tasks relying primarily on sentence-level or segment-level semantics). Without this characterization, a practitioner must either default to the conservative m = 64 (leaving efficiency gains on the table) or risk deploying m = 32 and discovering degradation in production.
More fundamentally, the paper provides no evidence that m = 64 is genuinely "sufficient" at sequence lengths beyond those evaluated. At n = 4096, 64 landmarks represent roughly 1.6% of tokens. At n = 8192, this drops to 0.8%. At n = 16,384, it would be 0.4%. There is almost certainly a sequence length beyond which 64 landmarks are insufficient — the paper provides no way to estimate this threshold.
What evidence exists in the paper. Table 1 reports efficiency for both m = 32 and m = 64 across all sequence lengths but reports accuracy only for (presumably) m = 64 in Tables 2 and 3 and Figures 5–6. The paper does not state which m-value was used for the downstream evaluations, though the model name "Nyströmformer" without a suffix likely refers to the 64-landmark configuration. No accuracy comparison between Nyströmformer-32 and Nyströmformer-64 exists for any task. The statement that "64 landmarks is often sufficient" is unsupported by evidence showing that other values are insufficient — it is an assertion, not a finding.
Mitigation status. The paper does not mitigate this limitation. It does not ablate m against accuracy, does not provide a heuristic for selecting m, and does not discuss the sensitivity of the method to landmark count. The statement about 64 landmarks being "often sufficient" is the entirety of the guidance provided. Future work could characterize the m-vs-accuracy trade-off on a standard benchmark, provide theoretical guidance based on the rank of the attention matrix, or develop adaptive schemes that vary m per layer or per input based on approximation quality estimates.
Generalization Beyond BERT-Style Encoder-Only Architectures Is Untested
The assumption or constraint. All experiments in the paper use BERT-style encoder-only Transformer architectures: BERT-small and BERT-base for pretraining and downstream evaluation, and a 2-layer encoder for the LRA benchmark. The pretraining objective is masked language modeling with sentence order prediction (ALBERT-style, following Lan et al., 2020), and all downstream tasks are classification or regression problems with single-vector outputs. The paper frames its contribution as a general self-attention replacement: "We propose Nyströmformer — a model that exhibits favorable scalability as a function of sequence length," and the Related Work section discusses Transformers broadly, referencing GPT-3 and machine translation.
The implicit assumption is that the Nyström approximation transfers to any architecture using self-attention: decoder-only models (GPT-style autoregressive LMs), encoder-decoder models (T5, BART), and models with different pretraining objectives (causal LM, span corruption). The paper provides no evidence for this assumption.
The consequence. Decoder-only architectures present a specific challenge not present in the BERT encoder setting: the causal attention mask. In standard autoregressive Transformers, the attention matrix is lower-triangular (token i can only attend to tokens j ≤ i), which changes the structure of the softmax matrix substantially. The Segment-means landmark selection, which partitions tokens into contiguous blocks, may interact poorly with causal masking — tokens in the first segment have very few preceding tokens, while tokens in later segments have many, creating an asymmetry not present in the bidirectional encoder setting. The Nyström reconstruction formula does not inherently respect the causal mask, and the paper does not discuss how (or whether) masking would be incorporated — would the landmarks also be causally constrained? Would the pseudoinverse computation need modification?
Encoder-decoder architectures add cross-attention (where queries come from the decoder and keys/values from the encoder) in addition to self-attention. The Nyström approximation would need to be applied to both, but the landmark selection strategy (Segment-means relies on the sequential structure of tokens) may not be appropriate for cross-attention, where the query-key alignment is across different sequences with potentially different lengths and structural properties.
More broadly, different pretraining objectives produce attention patterns with different properties. Causal language modeling may produce attention that is more locally focused (since future tokens are masked), potentially making it easier to approximate with landmarks. Span corruption objectives (T5) may produce more complex attention patterns. Without evaluation across architectures and objectives, the claim that Nyströmformer is a general self-attention replacement is speculative.
What evidence exists in the paper. None. All experiments are on BERT-style encoder-only models. The paper does not discuss architectural generalization, does not mention causal masking, and does not describe how Nyströmformer would be adapted to decoder or encoder-decoder settings.
Mitigation status. Not addressed. The paper does not flag this as a limitation and does not suggest future work on architectural generalization. The Related Work section and introduction frame the contribution in general terms ("Transformers have emerged as a powerful tool for a broad range of natural language processing tasks"), but the experimental scope is narrow. A reader deploying GPT-style models would need to independently determine whether the Nyström approximation transfers — the paper provides no guidance.
Segment-Means Landmark Selection Is Not Empirically Justified Against Alternatives
The assumption or constraint. The paper adopts Segment-means (contiguous block averaging) for landmark selection, explicitly dismissing K-means clustering as "less desirable during mini-batch training" and stating that it "requires a single scan of the sequence to compute the landmarks leading to a complexity of O(n)." The choice is presented as a practical necessity rather than an empirically validated optimal strategy. No experiment compares Segment-means against any alternative landmark selection method — not K-means, not random sampling of individual tokens, not learned landmark embeddings, not strided sampling (taking every (n/m)-th token rather than averaging blocks).
The assumption is that Segment-means is both necessary (for efficiency) and sufficient (for approximation quality). The paper provides no evidence for either claim beyond the downstream results, which show that the overall Nyströmformer system works — but these results confound the landmark selection method with every other aspect of the architecture and training procedure.
The consequence. Segment-means imposes a specific inductive bias: landmarks are forced to represent contiguous blocks of tokens. This is well-suited to tasks where attention patterns are locally smooth (adjacent tokens attend similarly), but may be poorly suited to tasks where semantically related tokens are far apart in the sequence (long-range coreference, multi-hop reasoning across separated paragraphs, code understanding where related variable uses are scattered). In such cases, averaging a contiguous block may blend semantically unrelated tokens, producing a landmark that represents none of them well, degrading the Nyström reconstruction for all tokens in that block.
Alternative schemes might substantially improve approximation quality at modest cost. Random sampling of individual tokens as landmarks (rather than averaging blocks) would preserve genuine token representations at the cost of potentially worse coverage of the sequence. Learned landmark embeddings (trained end-to-end) could adapt to the statistics of the task. K-means with periodic recomputation (not every batch) could provide content-based landmarks with acceptable overhead. Without any comparison, a practitioner cannot determine whether Segment-means is genuinely the best choice or merely a convenient one, and cannot assess how much accuracy is being sacrificed to the landmark selection heuristic.
What evidence exists in the paper. The paper provides a qualitative comparison in Figure 3, showing that the Nyström approximation with Segment-means produces attention patterns "quite similar" to standard self-attention on a single example. This is suggestive but not systematic — it does not characterize worst-case approximation failures, does not compare against alternative methods, and does not analyze which types of attention patterns are poorly approximated. The LRA results (Table 3) show that Nyströmformer slightly outperforms standard attention on average, which could be interpreted as evidence that the approximation (including landmark selection) works, but does not isolate the contribution of Segment-means.
Mitigation status. Not addressed. The paper does not report an ablation comparing landmark selection methods, does not discuss the limitations of Segment-means, and does not suggest future investigation of alternative landmark strategies. The dismissal of K-means is based on an untested assumption about mini-batch training compatibility, not on empirical evidence that it performs worse.
Pretraining Computational Savings at Standard Sequence Lengths Are Not Demonstrated
The assumption or constraint. The paper claims that Nyströmformer provides O(n) complexity and demonstrates dramatic efficiency gains at long sequence lengths (Table 1). However, the pretraining experiments (Figures 5 and 6) are conducted at the standard BERT sequence length of n = 512, training BERT-base for 0.5M steps "to keep compute costs reasonable." The paper does not report wall-clock training time or memory consumption during pretraining, and does not compare the actual training throughput of Nyströmformer against BERT-base at n = 512.
The implicit assumption is that the efficiency advantages demonstrated in the isolated self-attention module profiling (Table 1) translate to end-to-end pretraining speedups. At n = 512, however, the self-attention module profiling shows only modest gains: Nyströmformer-64 provides 1.5× memory reduction and 1.1× speedup relative to standard self-attention. These are for the attention module alone — overhead from feedforward layers, layer normalization, embedding lookups, and the Nyströmformer-specific operations (Segment-means, iterative pseudoinverse, three-kernel computation) could erode or eliminate this advantage in the full model.
The consequence. If Nyströmformer does not substantially accelerate pretraining at n = 512, the practical workflow for adoption becomes circuitous: a practitioner would need to pretrain at n = 512 with minimal speedup, then fine-tune or continue-pretrain at longer sequences where the efficiency advantage materializes. This raises the question of whether the pretrained representations learned with Nyström self-attention at n = 512 transfer effectively to longer sequences — a question the paper does not address. It also means that the total cost of training a Nyströmformer model (pretraining + long-sequence adaptation) may not be dramatically lower than training a standard Transformer, with the efficiency gains realized only during long-sequence inference.
Additionally, the paper's decision to train for 0.5M steps rather than the standard 1M steps "to keep compute costs reasonable" means the pretrained models may not be fully converged. The comparison in Figure 6 shows Nyströmformer tracking BERT-base at 0.5M steps, but it is unknown whether the gap would widen or close with additional training. The downstream results in Table 2 are based on these 0.5M-step checkpoints, so the reported accuracy numbers may understate what converged Nyströmformer could achieve — or may hide a degradation that emerges with longer training.
What evidence exists in the paper. Table 1 provides isolated self-attention module profiling, not end-to-end training measurements. At n = 512, standard self-attention takes 0.8 ms and 54 MB; Nyströmformer-64 takes 0.7 ms and 35 MB. These are small absolute differences (0.1 ms, 19 MB) that could easily be consumed by overhead in a full model. The paper does not report pretraining throughput, total pretraining wall-clock time, or GPU memory consumption during training.
The paper acknowledges that "training BERT-base with 1M update steps takes more than one week on 8 V100 GPUs" but does not state how long Nyströmformer pretraining takes on the same hardware — the reader cannot compare. The decision to train for 0.5M steps is explicitly motivated by compute cost, which suggests that pretraining was not dramatically cheaper than BERT-base, otherwise 1M steps would have been affordable.
Mitigation status. Not addressed as a limitation. The paper reports training details (batch size 256, learning rate 1e-4, 0.5M steps) but provides no resource comparison for the training phase. The efficiency narrative focuses entirely on inference-time memory and speed at long sequence lengths, sidestepping the question of whether Nyströmformer also accelerates the pretraining that produces the model in the first place. For a method whose stated goal is "resource efficient Transformers," this is a consequential omission.
No Empirical Characterization of Approximation Failure Modes
The assumption or constraint. The paper presents the Nyström approximation as a general replacement for standard self-attention, validated by aggregate accuracy metrics (GLUE scores, LRA accuracy) that show comparable performance across tasks. This aggregate evaluation assumes that approximation errors are either small on average across tokens and sequences, or that the model learns to compensate for systematic approximation errors during fine-tuning. The paper provides one qualitative example (Figure 3) showing similar attention patterns for a single input, but provides no systematic analysis of when the approximation fails.
The assumption is that comparable aggregate accuracy implies comparable behavior at the token and example level — that the approximation does not produce qualitatively different or systematically biased attention patterns that could lead to surprising failures in deployment.
The consequence. Without understanding failure modes, a practitioner deploying Nyströmformer cannot anticipate or detect when the approximation produces incorrect attention. Specific failure modes that could arise include:
-
Landmark blindness: Tokens in a segment whose average landmark poorly represents them (e.g., a segment containing a mix of content and function words where the average dilutes the content words) may receive distorted attention weights, causing the model to ignore or over-attend to specific positions.
-
Position-dependent approximation quality: Segment-means produces landmarks of varying quality depending on the homogeneity of each segment. Segments with high semantic variability (e.g., spanning a clause boundary) produce less representative landmarks than homogeneous segments (e.g., within a noun phrase). The approximation quality may therefore vary systematically with linguistic structure, creating position-dependent biases.
-
Catastrophic failure on specific attention patterns: Some attention patterns — highly sparse attention (one token attends strongly to one specific distant token), multi-modal attention (a token attends to two unrelated clusters), or attention distributed uniformly across many tokens — may be poorly captured by the rank-m Nyström reconstruction. The model might perform correctly on average while failing systematically on specific input types.
-
Interaction with the softmax temperature: The Nyström approximation applies softmax over m landmarks rather than n tokens, changing the effective "temperature" of the attention distribution. Sharp attention patterns (where one or a few tokens dominate) may be approximated differently than diffuse patterns.
These failure modes would manifest as per-example or per-token-group errors that aggregate metrics obscure. In safety-critical applications (medical text, legal analysis), such silent failures could be consequential.
What evidence exists in the paper. Figure 3 provides a single qualitative comparison: an attention matrix visualization where the Nyström approximation "looks similar" to standard self-attention. This is a existence proof that the approximation can work, not a characterization of when it doesn't. The paper reports no error analysis: no per-example accuracy distribution, no analysis of which LRA or GLUE examples Nyströmformer gets wrong that BERT gets right, and no measurement of the approximation error (e.g., Frobenius norm between true and approximate attention matrices) across a test set.
The LRA results (Table 3) show task-level variance: Nyströmformer substantially outperforms standard attention on Image (+3.38 points) but underperforms on Pathfinder (−3.22 points). This variation hints at task-dependent approximation quality but is not analyzed — the paper does not investigate why the approximation helps on one task and hurts on another, which would illuminate failure modes.
Mitigation status. Not addressed. The paper provides no diagnostic tools, no error bounds, and no analysis of when practitioners should distrust the approximation. The theoretical analysis (Lemma 2) provides a convergence condition (landmarks must exactly match tokens for exact reconstruction) that is never satisfied in practice and provides no practical guidance for detecting approximation failures. The low-rank justification ("self-attention is low-rank") is cited from prior work but not verified on the paper's own models or tasks — the paper does not measure the rank of its attention matrices to confirm that the assumption holds.
7. Implications and Future Directions
How This Work Changes the Landscape
Nyströmformer introduces a third category of attention approximation into the efficient-Transformer design space, distinct from both kernel-based linearization (which replaces the softmax function) and sparsity/projection methods (which restrict interactions or compress dimensions agnostically). This categorization is more than taxonomic — it reframes the efficient-attention problem as one of matrix reconstruction from data-dependent samples, opening a bridge between classical numerical linear algebra (specifically, low-rank matrix approximation theory) and deep learning architecture design that had been underexplored.
The magnitude of this shift is best understood as an incremental but architecturally significant refinement, not a paradigm shift. Nyströmformer does not overturn the Transformer architecture or propose a fundamentally new learning principle — it replaces one component (self-attention) with a more efficient approximation while preserving the rest of the stack. However, within the efficient-attention subfield, it introduces a design philosophy with implications beyond its specific instantiation: the approximation point in the computation graph matters more than the mathematical purity of the approximation method. The paper's core maneuver — selecting landmarks before the softmax, accepting that this is a "compromise" relative to classical Nyström theory, and engineering the resulting pipeline to map cleanly to GPU primitives — establishes a template for adapting other classical matrix approximation techniques (CUR decomposition, interpolative decomposition, randomized SVD) to deep learning architectures. Each of these techniques faces the same softmax-coupling problem; Nyströmformer demonstrates that pre-softmax landmark selection is a viable workaround.
The work also reconciles a latent tension in the efficient-attention literature between two evaluation philosophies. Prior work split into methods that preserve accuracy on standard benchmarks at the cost of modest efficiency gains (Longformer, BigBird, which maintain near-BERT accuracy on GLUE but are primarily designed for long documents) and methods that achieve dramatic efficiency gains at the cost of accuracy degradation on structured reasoning tasks (Reformer, Performer, which show significant accuracy drops on LRA tasks like ListOps). Nyströmformer is the first method in the paper's evaluated set to simultaneously achieve the dramatic efficiency gains of the second category (22.8× memory reduction at n = 8192) and the accuracy preservation of the first category (matching BERT-base on GLUE, matching or exceeding standard attention on LRA average). This demonstrates that the apparent trade-off between efficiency and accuracy in prior work was not fundamental — it was an artifact of the specific approximation strategies employed. The paper shows that a sufficiently faithful approximation of the softmax matrix itself (rather than a replacement of the operator or a structural constraint on attention patterns) can achieve both goals.
This finding redirects research attention in two ways. First, it makes matrix approximation methods more attractive as a research direction for efficient attention, since they have now been shown to work competitively with (and in several cases, outperform) the hashing, kernel, and sparsity approaches that dominated the literature at the time of publication. Second, it makes approaches that alter the softmax operator itself less attractive — the paper's results suggest that preserving the genuine softmax (even on a reduced landmark set) yields better downstream performance than replacing it with a linearized kernel, at least for the tasks and architectures evaluated. The consistently lower performance of ELU-linearized attention (Figure 5) and Performer (Table 3) relative to Nyströmformer provides empirical evidence for this claim.
The paper also elevates the importance of implementation-level design choices that had been under-discussed in the efficient-attention literature — specifically, the compatibility of the approximation computation with GPU hardware and autograd frameworks. The choice of iterative pseudoinverse over SVD, and Segment-means over K-means, are motivated by GPU efficiency and differentiability, not by asymptotic approximation quality. The paper implicitly argues (through its results) that these "engineering" decisions are in fact central to the method's success — a pure mathematical advantage in approximation error would be worthless if it couldn't be computed efficiently during training. This legitmizes hardware-aware algorithm design as a first-class concern in efficient-attention research, rather than an afterthought.
Finally, the paper establishes that data-dependent, input-adaptive approximation can outperform learned, fixed compression on the specific benchmarks evaluated. Linformer's learned projection matrices are trained once and applied uniformly to all inputs; Nyströmformer's landmarks are computed per-input from the actual query and key vectors. The fact that Nyströmformer outperforms Linformer on LRA (+3.36 average accuracy points) and matches it on efficiency (Table 1) suggests that adaptivity to input content provides a genuine accuracy benefit — the model can allocate its approximation budget where it matters for each specific sequence rather than using a one-size-fits-all compression. This finding has implications beyond attention: any setting where a learned fixed-dimensional bottleneck is used (e.g., in retrieval, compression, or routing) might benefit from input-dependent alternatives.
Follow-Up Research This Work Enables
Characterizing the accuracy-efficiency Pareto frontier as a function of landmark count m and sequence length n. The paper provides efficiency measurements for m = 32 and m = 64 (Table 1) and accuracy measurements for (presumably) m = 64 (Tables 2 and 3), but never jointly varies m and n on a fixed accuracy benchmark to map out the trade-off surface. A strong follow-up would systematically evaluate Nyströmformer variants with m ranging from 8 to 256 on tasks at sequence lengths n = 512, 1024, 2048, 4096, and 8192, measuring both downstream accuracy (e.g., on long-document classification or summarization) and wall-clock time/memory. This would answer the practical question the paper leaves open: for a given sequence length and accuracy target, what is the minimal m (and thus the maximal efficiency gain) achievable? The experiment would also test whether the effective rank of the attention matrix grows with sequence length — if the required m for fixed accuracy grows sublinearly with n, the O(n) complexity holds practically; if it grows proportionally, the method degrades to quadratic.
Combining Nyströmformer with sparse attention patterns for hybrid efficiency. Nyströmformer preserves global attention (every token can attend to every other token through the landmark reconstruction) but at coarse granularity (the reconstruction is rank-m). Sparse attention methods like Longformer or BigBird provide fine-grained local attention but restrict or approximate global attention. A natural hybrid would use Nyströmformer for the global attention component (replacing the task-specific global tokens in Longformer or BigBird with learned landmarks that attend to everything) while retaining sliding-window local attention for fine-grained token interactions. This would combine the complementary strengths: the Nyström approximation captures long-range dependencies cheaply through landmarks, and the local window captures detailed short-range patterns that the low-rank approximation might miss. The experiment would compare hybrid Nyströmformer-Longformer against pure Longformer and pure Nyströmformer on long-document tasks (e.g., QMSum, SummScreen, or arXiv summarization) with n = 4096–16384 tokens, measuring whether the hybrid recovers accuracy that either pure method loses.
Stress-testing the low-rank assumption across architectures and tasks. The paper invokes the low-rank property of self-attention (observed in Linformer) to justify the Nyström approximation, but never measures attention matrix rank in its own models or tasks. A diagnostic follow-up would compute the effective rank (e.g., the number of singular values exceeding 1% of the top singular value) of the attention matrices in a standard BERT-base model across layers, heads, and input types, then correlate this with Nyströmformer's approximation error (measured as Frobenius norm difference or downstream accuracy degradation) at varying m. This would determine whether the low-rank assumption actually holds in practice, which layers/heads are hardest to approximate (potentially guiding per-layer m allocation), and which types of inputs (e.g., short vs. long, structured vs. unstructured text) cause the assumption to break. A negative result — finding that attention matrices are not consistently low-rank for the tasks where Nyströmformer works — would motivate a different theoretical justification and potentially better approximation methods.
Testing Nyströmformer in decoder-only autoregressive and encoder-decoder architectures. The paper evaluates only BERT-style encoder-only models with bidirectional attention. A critical stress-test is whether the Nyström approximation transfers to causally masked self-attention (GPT-style autoregressive LMs) and cross-attention (encoder-decoder models like T5 or BART). The causal mask creates lower-triangular attention matrices, which may have different rank properties and interact differently with Segment-means landmark selection (since earlier tokens have fewer preceding context tokens). The cross-attention setting introduces query-key asymmetry that Segment-means (which segments both Q and K along their respective sequence dimensions) may handle poorly. A follow-up would implement Nyströmformer variants of GPT-2 (for autoregressive language modeling on WikiText-103 or PG-19, measuring perplexity vs. throughput at sequence lengths 1024–4096) and T5 (for summarization on long-document datasets, measuring ROUGE vs. memory). A negative result — substantial perplexity degradation or ROUGE loss — would bound the method's generality and motivate architecture-specific adaptations (e.g., learned causal landmark selection, or separate landmark strategies for self- and cross-attention).
Learned, content-aware landmark selection. Segment-means is a simple, position-based heuristic that ignores the semantic content of queries and keys. A natural extension is to learn the landmark selection function — either via a lightweight neural module that predicts which tokens should serve as landmarks (trained end-to-end with the main objective), or via a differentiable clustering mechanism that selects landmarks based on query/key similarity rather than sequence position. The specific experiment would replace Segment-means with a small attention-based or routing-based module that takes Q and K as input and outputs m landmark vectors (as weighted combinations of tokens rather than simple averages), then evaluate on tasks where position-based segmentation is clearly suboptimal — e.g., tasks requiring cross-paragraph reasoning where key tokens are scattered throughout the document. Comparing learned landmarks against Segment-means on a multi-hop QA dataset (HotpotQA with full documents, sequence lengths 4096+) would test whether content-aware selection recovers accuracy that Segment-means loses. The trade-off is increased landmark selection cost (the learned module adds parameters and computation), so the evaluation must measure whether the accuracy gain justifies the efficiency reduction.
Error analysis and diagnostic tools for approximation failures. The paper provides no characterization of when the Nyström approximation fails, only aggregate accuracy metrics showing it works on average. A diagnostic study would instrument Nyströmformer to compute per-example approximation error (Frobenius norm between the true softmax attention matrix and the Nyström reconstruction, computed at evaluation time where the true matrix can be affordably calculated for a subset of examples), then correlate this error with downstream task accuracy and input properties (sequence length, syntactic complexity, number of entities, etc.). The study would identify the types of inputs where Nyströmformer's accuracy diverges from standard attention, potentially revealing systematic failure modes (e.g., it may fail on inputs requiring attention to rare or isolated tokens that get averaged away in Segment-means landmarks). This would enable practitioners to know when to trust the approximation and would guide targeted improvements (e.g., adaptive m allocation based on input complexity).
Practical Applications and Downstream Use Cases
Long-document classification and retrieval in production search systems. A search engine or document management system that needs to classify or retrieve documents based on full-text understanding (e.g., legal e-discovery, scientific literature search, or contract review) currently must either truncate documents to 512 tokens (losing information) or use sparse attention with hand-designed global tokens (Longformer/BigBird). Nyströmformer provides a drop-in replacement for standard BERT-style encoders that scales to document lengths of 4096–8192 tokens without task-specific global token engineering, while matching BERT-base accuracy on standard benchmarks (GLUE Table 2: 91.4 vs. 90.0 on SST-2, 93.2 vs. 93.3 on IMDB) and providing 5.8× memory reduction and 3.6× speedup at n = 2048 (Table 1, Nyströmformer-64). This means a production system currently running BERT-base on 512-token chunks could move to processing full documents at 2048+ tokens with similar per-document latency but dramatically better context utilization — potentially improving retrieval recall by capturing cross-paragraph evidence that chunk-based models miss.
On-device or edge deployment of Transformer models for real-time long-sequence processing. Deploying Transformer models on mobile devices or embedded systems for tasks like real-time speech transcription, on-device document understanding, or continuous health monitoring from long sensor streams is currently limited by the quadratic memory scaling of self-attention. Nyströmformer-32 at n = 8192 consumes only 383 MB and 11.5 ms per instance through the attention module (Table 1) — a 26.7× memory reduction and 13.4× speedup over standard self-attention — potentially making it feasible to run long-sequence Transformer inference within the memory and latency budgets of high-end mobile GPUs or edge accelerators. The trade-off between the 32-landmark and 64-landmark variants (the paper does not report accuracy for m = 32) would need characterization, but even the 64-landmark variant's 22.8× memory reduction at n = 8192 represents a meaningful step toward on-device long-sequence capabilities.
Efficient pretraining data generation pipelines for self-supervised learning. When generating pretraining data using large Transformer models (e.g., creating synthetic training examples, filtering noisy data, or scoring candidate documents), the model is applied to massive corpora with sequence lengths that can reach thousands of tokens. Even modest per-example speedups compound across millions of documents. At n = 1024, Nyströmformer-64 provides 3.0× memory reduction and 1.8× speedup over standard self-attention (Table 1); at n = 2048, these become 5.8× memory and 3.6× speedup. For a data processing pipeline processing 10 million documents of average length 2048 tokens, the memory savings could reduce the required GPU count by a factor of ~5 (from needing many GPUs to hold attention matrices in memory to fitting on a single machine), and the speedup could reduce processing time from weeks to days.
Long Range Arena as a standard evaluation substrate for future efficient-attention methods. The paper's LRA reimplementation and baseline results (Table 3) provide a PyTorch-based benchmark that subsequent efficient-attention proposals can use for apples-to-apples comparison. The paper establishes Nyströmformer as the top-performing efficient method (58.95 average, exceeding Linformer by 3.36 points and Performer by 5.32 points) and essentially matching standard attention (58.77), setting a strong baseline. New proposals can evaluate against these numbers using the identical 2-layer, 64-embed-dim, 128-hidden-dim, 2-head architecture. The paper's note that PyTorch reimplementation yields higher Retrieval scores for all models suggests that researchers should use the paper's PyTorch baselines (not the original Jax/Flax LRA numbers) for fair comparison.
When to Prefer This Method
The paper does not provide an explicit decision framework comparing Nyströmformer against specific named alternatives with clear conditionals. The efficiency comparisons in Table 1 position Nyströmformer relative to Linformer and Longformer, and the accuracy comparisons in Table 3 position it relative to Reformer, Linformer, and Performer, but the paper does not articulate a structured trade-off rule (e.g., "prefer Nyströmformer over Linformer when X; prefer Longformer over Nyströmformer when Y"). Instead, the results demonstrate that Nyströmformer is competitive with or superior to these alternatives on the evaluated dimensions (efficiency at scale, LRA accuracy), without characterizing the specific conditions under which each alternative would be preferable. A forced decision matrix would therefore fabricate trade-offs the paper does not establish.