Approximating Attention: Linear Transformers

As highlighted in the chapter introduction, the standard self-attention mechanism's primary computational bottleneck is its quadratic complexity concerning the input sequence length $N$. Calculating the attention score matrix via $A = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)$ involves computing the $N \times N$ matrix $QK^T$. This requires $N^2$ dot products between query and key vectors (each of dimension $d_k$), followed by scaling, softmax normalization, and finally multiplication by the value matrix $V$. The overall complexity is typically dominated by the $QK^T$ calculation, resulting in $O(N^2 d_k)$ time and $O(N^2)$ space complexity for the attention matrix itself. This scaling behavior makes it computationally prohibitive to apply standard Transformers to very long sequences (e.g., high-resolution images treated as sequences of patches, long documents, or genomic data).

To overcome this limitation, a significant line of research has focused on developing Linear Transformers, variants that approximate the self-attention mechanism with linear, $O(N)$, time and memory complexity. The core idea is to avoid the explicit computation and storage of the full $N \times N$ attention matrix.

The Challenge of Softmax

If we consider the attention computation without the softmax normalization, $O_{unnorm} = (QK^T)V$, matrix multiplication associativity allows us to reorder the computation as $O_{unnorm} = Q(K^TV)$. Here, $K^T$ is $d_k \times N$ and $V$ is $N \times d_v$. Computing $K^TV$ takes $O(N d_k d_v)$ operations. Multiplying the result by $Q$ ($N \times d_k$) then takes another $O(N d_k d_v)$ operations. If the embedding dimension $d = d_k = d_v$ is considered fixed or much smaller than $N$, the overall complexity appears linear in $N$.

However, the row-wise softmax function, defined as: $\text{softmax}(Z)_i = \frac{\exp(Z_i)}{\sum_j \exp(Z_j)}$ where $Z_i$ represents the $i$-th row of the input matrix (in our case, the $i$-th row of $\frac{QK^T}{\sqrt{d_k}}$), prevents this simple reordering. The normalization term $\sum_j \exp(\dots)$ couples all elements within a row, seemingly requiring the computation of all $N$ dot products for a given query before normalization and multiplication by values.

Strategies for Linearization

Linear Transformer variants employ different strategies to approximate the attention mechanism and bypass the softmax bottleneck, effectively enabling a rearrangement similar to $Q(K^TV)$:

1. Kernel Approximation: This approach approximates the dot-product similarity followed by the exponential function inherent in softmax, often viewing $\exp(q_i^T k_j / \sqrt{d_k})$ as a kernel function. If this kernel can be approximated by an inner product of feature maps $\phi(\cdot)$, such that $K(q_i, k_j) \approx \phi(q_i)^T \phi(k_j)$, the attention sum $\sum_j K(q_i, k_j) v_j$ can be rewritten: $\sum_j \left( \phi(q_i)^T \phi(k_j) \right) v_j = \phi(q_i)^T \left( \sum_j \phi(k_j) v_j^T \right)$ (Note: Vector/matrix dimensions need careful handling here, this shows the reordering). The term $\sum_j \phi(k_j) v_j^T$ computes an aggregate representation based on keys and values. This aggregate can be computed once (linearly in $N$) and then multiplied by the feature map of each query $\phi(q_i)$ (again, linearly in $N$). The challenge lies in finding suitable feature maps $\phi$ that accurately approximate the exponential kernel while allowing this decomposition. Models like the Performer use techniques based on random features (specifically, Fast Attention via Positive Orthogonal Random Features or FAVOR+) to construct these mappings with theoretical guarantees on the approximation quality. These methods aim for $O(N)$ complexity. We will examine kernel-based methods more closely in the section "Kernel-Based Attention Approximation (Performers)".

2. Low-Rank Projection: This strategy is based on the observation or assumption that the attention matrix $A = \text{softmax}(QK^T / \sqrt{d_k})$ is often low-rank or can be well-approximated by a low-rank matrix. The Linformer model exemplifies this by introducing learnable projection matrices $E$ ($k \times N$) and $F$ ($k \times N$), where $k$ is a fixed dimension significantly smaller than $N$ ($k \ll N$). It computes projected key and value matrices $K' = EK$ and $V' = FV$. These projected matrices have dimensions $k \times d_k$ and $k \times d_v$, respectively. The attention mechanism then operates using these projected matrices, for instance, by computing attention weights based on $Q(K')^T$ ($N \times k$) and applying them to $V'$ ($k \times d_v$). The complexity reduces from $O(N^2)$ to $O(Nk)$, achieving linear scaling if $k$ is constant or grows sub-linearly with $N$. This approach will be detailed in the section "Low-Rank Projection Methods (Linformer)".

Implications and Trade-offs

Linear attention mechanisms offer a compelling advantage in terms of computational efficiency, particularly for long sequences.

Feature	Standard Attention	Linear Attention (Approx.)
Time Complexity	$O(N^2 d)$	$O(N d r)$ (where $r$ depends on method, e.g., feature dim or projection size $k$)
Memory (Attn Matrix)	$O(N^2)$	$O(Nr)$ or less
Computation	Explicit $N \times N$ matrix	Avoids explicit $N \times N$ matrix
Exactness	Exact	Approximate
Suitability	Short to moderate sequences	Long sequences

The primary trade-off is that these methods compute an approximation of the full self-attention mechanism. While empirical results often show that these approximations perform remarkably well, sometimes even outperforming standard attention on specific tasks (potentially due to implicit regularization effects), they might not capture exactly the same relational information as the original formulation. The quality of the approximation and its impact on downstream task performance must be evaluated empirically.

By reducing the computational burden from quadratic to linear, these methods significantly expand the applicability of Transformer architectures to domains previously hindered by sequence length constraints, enabling models to process entire documents, high-resolution images, or long time-series data more effectively.