Nyströmformer: Approximating self-attention in linear time and memory via the Nyström method

Self-attention is the reason transformers got good, and also the reason they get slow. Every token has to compare itself against every other token, so the compute and memory bill scales quadratically with sequence length. Feed it a document with 8,000 tokens and the attention matrix alone becomes a monster. Nyströmformer, detailed in a Hugging Face writeup, tackles this with an old numerical trick borrowed from matrix approximation theory rather than a new architecture gimmick.

The idea comes from the Nyström method, a decades-old technique for approximating a large matrix by sampling a handful of its rows and columns instead of computing every entry. Normally you'd sample directly from the attention matrix, but there's a catch: attention scores come out of a softmax, and softmax needs the whole row to normalize a single entry. You can't grab one column without effectively already knowing the rest. So sampling straight from the attention output is a dead end.

The workaround is to sample from the queries and keys instead, before the softmax ever happens. The authors pick a small set of 'landmark' queries and keys, typically just 32 or 64 of them, by averaging tokens into segments rather than randomly picking individual ones. Three smaller matrices get built from these landmarks, softmaxed separately, and multiplied together with a pseudo-inverse step to reconstruct an approximate version of the full attention matrix. Crucially, the expensive n-by-n query-key product never gets computed at all, which is what actually delivers the linear time and memory scaling.

What's notable is how well this holds up at scale. The paper reports competitive results even at sequence lengths of 4,096 and 8,192 tokens, using only 32 or 64 landmarks, a tiny fraction of the input. Hugging Face has already shipped four pretrained checkpoints, covering sequence lengths from 512 up to 4,096, and integrated the model into the Transformers library with a configurable num_landmarks parameter. A quick masked-language-modeling test with the 512-token checkpoint correctly filled in 'capital' for 'Paris is the [MASK] of France,' which at least confirms the approximation doesn't wreck basic language understanding.

The implementation itself is refreshingly compact, just a handful of matrix multiplications and softmaxes standing in for the usual attention block, plus a depthwise convolution skip connection the blog glosses over for simplicity. For anyone training on long-context data without a warehouse of GPUs, that's the real selling point here: efficient attention that doesn't ask you to babysit a novel architecture.