Low-Rank Projection Methods (Linformer)

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}$
• $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 main 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:

1. Projecting K: $E K$ takes $O(Nk d_k)$ time.
2. Projecting V: $F V$ takes $O(Nk d_v)$ time.
3. 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.
4. 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.