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

Introduction

Transformers have exhibited remarkable performance on various Natural Language Processing and Computer Vision tasks. Their success can be attributed to the self-attention mechanism, which captures the pairwise interactions between all the tokens in an input. However, the standard self-attention mechanism has a time and memory complexity of O(n2)O(n^2) (where nn is the length of the input sequence), making it expensive to train on long input sequences.

The Nyströmformer is one of many efficient Transformer models that approximates standard self-attention with O(n)O(n) complexity. Nyströmformer exhibits competitive performance on various downstream NLP and CV tasks while improving upon the efficiency of standard self-attention. The aim of this blog post is to give readers an overview of the Nyström method and how it can be adapted to approximate self-attention.

Nyström method for matrix approximation

At the heart of Nyströmformer is the Nyström method for matrix approximation. It allows us to approximate a matrix by sampling some of its rows and columns. Let's consider a matrix Pn×nP^{n \times n}, which is expensive to compute in its entirety. So, instead, we approximate it using the Nyström method. We start by sampling mm rows and columns from PP. We can then arrange the sampled rows and columns as follows:

We now have four submatrices: AP,BP,FP,A_P, B_P, F_P, and CPC_P, with sizes m×m,m×(n−m),(n−m)×mm \times m, m \times (n - m), (n - m) \times m and (n−m)×(n−m)(n - m) \times (n - m) respectively. The mm sampled columns are contained in APA_P and FPF_P, whereas the mm sampled rows are contained in APA_P and BPB_P. So, the entries of AP,BP,A_P, B_P, and FPF_P are known to us, and we will estimate CPC_P. According to the Nyström method, CPC_P is given by:

CP=FPAP+BPC_P = F_P A_P^+ B_P

Here, ++ denotes the Moore-Penrose inverse (or pseudoinverse). Thus, the Nyström approximation of P,P^P, \hat{P} can be written as: