Kernel-Based Attention Approximation (Performers)

While standard self-attention provides strong sequence modeling capabilities, its quadratic computational and memory complexity ($O(N^2)$ for sequence length $N$) remains a significant bottleneck, especially when dealing with long sequences encountered in areas like high-resolution image generation, document summarization, or genomic data analysis. The Performer architecture, introduced by Choromanski et al. (2020), offers an effective solution by approximating the attention mechanism with linear complexity. It achieves this through a clever technique called Fast Attention Via positive Orthogonal Random Features (FAVOR+).

The Bottleneck: Explicit Attention Matrix Calculation

Recall the standard scaled dot-product attention:

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Here, $Q, K, V$ are matrices of query, key, and value vectors, respectively, each with $N$ rows (sequence length) and $d_k$ (or $d_v$) columns (embedding dimension). The primary computational cost lies in computing the $N \times N$ attention matrix $A = \text{softmax}(QK^T / \sqrt{d_k})$. This requires $O(N^2 d_k)$ operations and $O(N^2)$ memory, which quickly becomes prohibitive as $N$ increases.

Approximating Attention with Kernels and Random Features

The core idea behind Performers is to avoid the explicit computation of the $N \times N$ matrix $A$. Instead, it seeks to approximate the attention mechanism using kernel methods. Let's rewrite the $i$-th output vector of the attention mechanism (before normalization) as:

$$\text{Att}'(Q, K, V)_i = \sum_{j=1}^N \exp\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right) v_j$$

The standard attention output is then obtained by normalizing this sum using the denominator $D_i = \sum_{j=1}^N \exp(q_i^T k_j / \sqrt{d_k})$.

Performers leverage the concept that the function $\text{sim}(q_i, k_j) = \exp(q_i^T k_j / \sqrt{d_k})$ resembles a kernel function, specifically a Gaussian kernel after appropriate scaling and transformation. Kernel methods often allow calculating similarities implicitly using feature maps. Performers propose finding a feature map $\phi: \mathbb{R}^{d_k} \rightarrow \mathbb{R}^r$, where $r$ is the dimension of the random features (typically $r \ll N$), such that the kernel similarity can be approximated by an inner product in the feature space:

$$\exp\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right) \approx \phi(q_i)^T \phi(k_j)$$

If such a feature map $\phi$ exists, the attention computation can be cleverly reordered to avoid the $N \times N$ calculation. The unnormalized attention sum becomes:

$$\text{Att}'(Q, K, V)_i \approx \sum_{j=1}^N (\phi(q_i)^T \phi(k_j)) v_j = \phi(q_i)^T \sum_{j=1}^N (\phi(k_j) v_j^T)$$

Similarly, the denominator can be approximated as:

$$D_i \approx \sum_{j=1}^N \phi(q_i)^T \phi(k_j) = \phi(q_i)^T \sum_{j=1}^N \phi(k_j)$$

Notice the important change in the order of operations. We can first compute the sums $\sum_{j=1}^N (\phi(k_j) v_j^T)$ (an $r \times d_v$ matrix) and $\sum_{j=1}^N \phi(k_j)$ (an $r \times 1$ vector). These computations involve processing the keys and values once, taking roughly $O(N r d_k + N r d_v)$ time. Then, for each query $q_i$, we compute $\phi(q_i)$ and perform two matrix-vector multiplications, taking $O(r d_k + r d_v)$ time per query. The total time complexity becomes approximately $O(N r (d_k + d_v))$, which is linear in $N$ if $r$ is considered constant or grows much slower than $N$. The memory complexity also reduces to $O(N r + N d_k + N d_v)$, primarily for storing the mapped queries, keys, values, and intermediate sums.

Comparison of computational flow. Standard attention requires forming the expensive $N \times N$ matrix (red nodes). Performer uses feature maps $\phi$ to compute intermediate sums (yellow nodes) involving linear complexity operations $O(N)$, avoiding the quadratic bottleneck.

Constructing the Feature Map $\phi$

The effectiveness of Performers depends on finding a suitable feature map $\phi$. The FAVOR+ mechanism achieves this by constructing $\phi$ using random features, drawing inspiration from techniques used to approximate Gaussian kernels (like Random Fourier Features). An important contribution of Performers is the development of feature maps that guarantee positivity ($\phi(x)^T \phi(y) \ge 0$), which helps maintain stability and better approximates the properties of the softmax function (whose output components are non-negative).

Specifically, Performers define feature maps based on random projections and trigonometric functions, for example:

$$\phi(x) = \frac{h(x)}{\sqrt{m}} \left[ f_1(w_1^T x), \dots, f_m(w_m^T x), f'_1(w_1^T x), \dots, f'_m(w_m^T x) \right]^T$$

where $w_i$ are randomly sampled vectors, $h(x)$ is a function like $\exp(\|x\|^2 / 2)$, and $f_l, f'_l$ are pairs like $(\sin, \cos)$ or $(\exp, \exp)$. The dimension of the feature map $r$ is $2m$. The exact construction ensures unbiased or near-unbiased estimation of the softmax kernel with positive outputs.

Advantages

Advantages:

• Linear Complexity: Reduces attention time and space complexity from $O(N^2)$ to $O(N)$, enabling application to much longer sequences.
• Strong Theoretical Guarantees: Provides provable approximation bounds for the softmax kernel.
• Good Empirical Performance: Often achieves performance comparable to standard Transformers while being significantly more efficient on long-sequence tasks.
• Compatibility: Can be integrated as a replacement for the standard attention module with minimal changes to the overall Transformer architecture.

Considerations:

• Approximation Quality vs. Cost: The quality of the approximation depends on the random feature dimension $r$. A larger $r$ yields a better approximation but increases the constant factors in the linear complexity ($O(N r d)$). Choosing $r$ involves a trade-off between accuracy and computational cost. Typical values for $r$ might range from 64 to 256 or higher, depending on the specific task and resources.
• Randomness: The use of random features introduces stochasticity. However, the variance is typically low, especially with a reasonable choice of $r$, and Performers employ techniques like orthogonal random features to reduce this variance further. In practice, results are often stable across different random seeds.

By approximating the softmax kernel using positive random features, Performers provide a computationally efficient and theoretically grounded method to scale Transformer models to lengths previously considered infeasible, significantly broadening their applicability.