As discussed earlier in this chapter, the quadratic complexity of standard self-attention, $O(N^2 d)$ , where $N$ is the sequence length and $d$ is the model dimension, represents a significant bottleneck for processing long sequences. While powerful, calculating attention scores between every pair of tokens becomes computationally prohibitive as $N$ increases.

Several approaches aim to alleviate this computational burden. One prominent direction involves approximating the attention mechanism using low-rank projections, as exemplified by the Linformer model.

The Low-Rank Hypothesis

The core idea behind Linformer is based on the hypothesis that the self-attention mechanism, despite its potential to model complex interactions across the entire sequence, can often be approximated by a low-rank matrix. In essence, the $N \times N$ attention matrix $P = Softmax(\frac{QK^T}{\sqrt{d_k}})$ might be structurally redundant. This means the mapping from input context (represented by Values, $V$ ) to output context (the weighted sum $PV$ ) doesn't necessarily require the full expressive power of an arbitrary $N \times N$ matrix. If this hypothesis holds, we can potentially achieve significant efficiency gains by working with a compressed representation.

Linformer: Projecting Keys and Values

Linformer (Linear Transformer) proposes a clever way to achieve linear complexity by introducing learned projection matrices $E_i$ and $F_i$ . Instead of computing the full $N \times N$ attention matrix, Linformer projects the Key ( $K$ ) and Value ( $V$ ) matrices along the sequence length dimension before the attention computation.

Let the input sequence length be $N$ and the head dimension be $d_k$ (for Keys/Queries) or $d_v$ (for Values). The original matrices are:

Query $Q \in \mathbb{R}^{N \times d_k}$
Key $K \in \mathbb{R}^{N \times d_k}$
Value $V \in \mathbb{R}^{N \times d_v}$

Linformer introduces two projection matrices, $E \in \mathbb{R}^{k \times N}$ and $F \in \mathbb{R}^{k \times N}$ , where $k$ is a projected dimension significantly smaller than $N$ ( $k \ll N$ ). These matrices are used to create projected Key and Value matrices:

K_{proj} = E K \quad (\text{dimension } k \times d_k)

V_{proj} = F V \quad (\text{dimension } k \times d_v)

Notice how the projection reduces the sequence dimension from $N$ to $k$ . The crucial step is that the attention scores are now computed between the original Query matrix $Q$ and the projected Key matrix $K_{proj}$ :

P_{proj} = Softmax\left(\frac{Q K_{proj}^T}{\sqrt{d_k}}\right) \quad (\text{dimension } N \times k)

The final output is then obtained by multiplying this projected attention score matrix $P_{proj}$ with the projected Value matrix $V_{proj}$ :

Attention_{Linformer}(Q, K, V) = P_{proj} V_{proj} \quad (\text{dimension } N \times d_v)

Complexity Analysis

Let's analyze the computational complexity. The original attention calculation is dominated by the $Q K^T$ matrix multiplication, which takes $O(N^2 d_k)$ time, and the subsequent multiplication by $V$ , which takes $O(N^2 d_v)$ time.

In Linformer:

Projecting K: $E K$ takes $O(Nk d_k)$ time.
Projecting V: $F V$ takes $O(Nk d_v)$ time.
Computing $Q K_{proj}^T$ : This involves multiplying an $N \times d_k$ matrix by a $d_k \times k$ matrix, resulting in $O(Nk d_k)$ time.
Computing $P_{proj} V_{proj}$ : This involves multiplying an $N \times k$ matrix by a $k \times d_v$ matrix, resulting in $O(Nk d_v)$ time.

Since $k$ is chosen such that $k \ll N$ , the overall complexity of Linformer's attention mechanism becomes $O(Nk)$ , which is linear with respect to the sequence length $N$ . This is a substantial improvement over the standard $O(N^2)$ complexity.

Comparison of computational flow in standard self-attention versus Linformer's projected attention. Linformer introduces projection matrices (E, F) to reduce the dimensionality of Keys and Values along the sequence length axis before calculating attention scores.

Implementation Considerations

Projection Matrices: The projection matrices $E$ and $F$ are typically learned during training. They can be shared across different attention heads or even layers to reduce the parameter count further. A common implementation uses simple linear layers acting on the Key and Value matrices transposed ( $K^T$ , $V^T$ ) to perform the projection efficiently.
Choice of k: The projected dimension $k$ is a hyperparameter. Values like 128, 256, or 512 are often used, significantly smaller than typical sequence lengths (thousands or tens of thousands) where Linformer becomes advantageous. The choice impacts the trade-off between computational efficiency and model accuracy. A smaller $k$ is faster but might lead to a larger approximation error.
Theoretical Guarantees: The Linformer paper provides theoretical analysis showing that under certain assumptions, the self-attention matrix is indeed low-rank and can be well-approximated by this projection method.

Advantages and Disadvantages

Advantages:

Linear Time Complexity: Reduces attention time complexity from $O(N^2)$ to $O(Nk)$ .
Linear Space Complexity: Reduces memory usage associated with storing the attention matrix.
Scalability: Enables processing much longer sequences than standard Transformers.
Simple Modification: Relatively straightforward to implement by adding projection layers.

Disadvantages:

Approximation Error: As an approximation, it may not capture all the nuances of the full attention matrix, potentially leading to a performance decrease on certain tasks compared to standard attention (especially for shorter sequences where $N^2$ is manageable).
Hyperparameter Tuning: Requires tuning the projection dimension $k$ .
Low-Rank Assumption: Its effectiveness relies on the underlying assumption that the attention matrix can be well approximated by a low-rank structure, which might not hold equally well for all data or tasks.

Linformer represents a significant step towards building more efficient Transformer models. By challenging the necessity of the full quadratic attention computation and leveraging the concept of low-rank approximation, it offers a practical way to scale Transformers to longer sequences, opening up possibilities for applications involving extensive documents, high-resolution images, or long-form audio.