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(n^2)O(n^2)$

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 $P^{n\times n}P^\{n\; \backslash times\; n\}$

We now have four submatrices: $A_P,B_P,F_P,A\_P,\; B\_P,\; F\_P,$

$$C_P=F_P{A}_P^{+}{B}_{P}C\_P\; =\; F\_P\; A\_P^+\; B\_P$$

Here, $++$ denotes the Moore-Penrose inverse (or pseudoinverse).
Thus, the Nyström approximation of $P,\widehat{P}P,\; \backslash hat\{P\}$

As shown in the second line, $\widehat{P}\backslash hat\{P\}$

Can we approximate self-attention with the Nyström method?

Our goal is to ultimately approximate the softmax matrix in standard self attention: S = softmax $\frac{Q{K}^T}{\sqrt{d}}\backslash frac\{QK^T\}\{\backslash sqrt\{d\}\}$

Here, $QQ$ and $KK$ denote the queries and keys respectively. Following the procedure discussed above, we would sample $mm$ rows and columns from $SS$, form four submatrices, and obtain $\widehat{S}\backslash hat\{S\}$

But, what does it mean to sample a column from $SS$? It means we select one element from each row. Recall how S is calculated: the final operation is a row-wise softmax. To find a single entry in a row, we must access all other entries (for the denominator in softmax). So, sampling one column requires us to know all other columns in the matrix. Therefore, we cannot directly apply the Nyström method to approximate the softmax matrix.

How can we adapt the Nyström method to approximate self-attention?

Instead of sampling from $SS$, the authors propose to sample landmarks (or Nyström points) from queries and keys. We denote the query landmarks and key landmarks as $\tilde{Q}\backslash tilde\{Q\}$

$$\tilde{F}=softmax(\frac{Q{\tilde{K}}^T}{\sqrt{d}})\tilde{A}=softmax(\frac{\tilde{Q}{\tilde{K}}^T}{\sqrt{d}}{)}^{+}\tilde{B}=softmax(\frac{\tilde{Q}{K}^T}{\sqrt{d}})\backslash tilde\{F\}\; =\; softmax(\backslash frac\{Q\backslash tilde\{K\}^T\}\{\backslash sqrt\{d\}\})\; \backslash hspace\{40pt\}\; \backslash tilde\{A\}\; =\; softmax(\backslash frac\{\backslash tilde\{Q\}\backslash tilde\{K\}^T\}\{\backslash sqrt\{d\}\})^+\; \backslash hspace\{40pt\}\; \backslash tilde\{B\}\; =\; softmax(\backslash frac\{\backslash tilde\{Q\}K^T\}\{\backslash sqrt\{d\}\})$$

The sizes of $\tilde{F}\backslash tilde\{F\}$

This is the Nyström approximation of the softmax matrix in the self-attention mechanism. We multiply this matrix with the values ( $VV$) to obtain a linear approximation of self-attention. Note that we never calculated the product $Q{K}^{T}QK^T$

How do we select landmarks?

Instead of sampling $mm$ rows from $QQ$ and $KK$, the authors propose to construct $\tilde{Q}\backslash tilde\{Q\}$

The overall algorithm is summarised by the following figure from the paper:

The three orange matrices above correspond to the three matrices we constructed using the key and query landmarks. Also, notice that there is a DConv box. This corresponds to a skip connection added to the values using a 1D depthwise convolution.

How is Nyströmformer implemented?

The original implementation of Nyströmformer can be found here and the HuggingFace implementation can be found here. Let’s take a look at a few lines of code (with some comments added) from the HuggingFace implementation. Note that some details such as normalization, attention masking, and depthwise convolution are avoided for simplicity.

key_layer = self.transpose_for_scores(self.key(hidden_states)) 
value_layer = self.transpose_for_scores(self.value(hidden_states)) 
query_layer = self.transpose_for_scores(mixed_query_layer) 

q_landmarks = query_layer.reshape(
    -1,
    self.num_attention_heads,
    self.num_landmarks,
    self.seq_len // self.num_landmarks,
    self.attention_head_size,
).mean(dim=-2) 

k_landmarks = key_layer.reshape(
    -1,
    self.num_attention_heads,
    self.num_landmarks,
    self.seq_len // self.num_landmarks,
    self.attention_head_size,
).mean(dim=-2) 

kernel_1 = torch.nn.functional.softmax(torch.matmul(query_layer, k_landmarks.transpose(-1, -2)), dim=-1) 
kernel_2 = torch.nn.functional.softmax(torch.matmul(q_landmarks, k_landmarks.transpose(-1, -2)), dim=-1) 

attention_scores = torch.matmul(q_landmarks, key_layer.transpose(-1, -2)) 

kernel_3 = nn.functional.softmax(attention_scores, dim=-1) 
attention_probs = torch.matmul(kernel_1, self.iterative_inv(kernel_2)) 
new_value_layer = torch.matmul(kernel_3, value_layer) 
context_layer = torch.matmul(attention_probs, new_value_layer) 

Using Nyströmformer with HuggingFace

Nyströmformer for Masked Language Modeling (MLM) is available on HuggingFace. Currently, there are 4 checkpoints, corresponding to various sequence lengths: nystromformer-512, nystromformer-1024, nystromformer-2048, and nystromformer-4096. The number of landmarks, $mm$, can be controlled using the num_landmarks parameter in the NystromformerConfig. Let’s take a look at a minimal example of Nyströmformer for MLM:

from transformers import AutoTokenizer, NystromformerForMaskedLM
import torch

tokenizer = AutoTokenizer.from_pretrained("uw-madison/nystromformer-512")
model = NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512")

inputs = tokenizer("Paris is the [MASK] of France.", return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits


mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]

predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
tokenizer.decode(predicted_token_id)

Output:
----------------------------------------------------------------------------------------------------
capital

Alternatively, we can use the pipeline API (which handles all the complexity for us):

from transformers import pipeline
unmasker = pipeline('fill-mask', model='uw-madison/nystromformer-512')
unmasker("Paris is the [MASK] of France.")

Output:
----------------------------------------------------------------------------------------------------
[{'score': 0.829957902431488,
  'token': 1030,
  'token_str': 'capital',
  'sequence': 'paris is the capital of france.'},
{'score': 0.022157637402415276,
  'token': 16081,
  'token_str': 'birthplace',
  'sequence': 'paris is the birthplace of france.'},
{'score': 0.01904447190463543,
  'token': 197,
  'token_str': 'name',
  'sequence': 'paris is the name of france.'},
{'score': 0.017583081498742104,
  'token': 1107,
  'token_str': 'kingdom',
  'sequence': 'paris is the kingdom of france.'},
{'score': 0.005948934704065323,
  'token': 148,
  'token_str': 'city',
  'sequence': 'paris is the city of france.'}]

Conclusion

Nyströmformer offers an efficient approximation to the standard self-attention mechanism, while outperforming other linear self-attention schemes. In this blog post, we went over a high-level overview of the Nyström method and how it can be leveraged for self-attention. Readers interested in deploying or fine-tuning Nyströmformer for downstream tasks can find the HuggingFace documentation here.