Linear Attention › Kernel Feature Maps

The choice of feature map $\phi$ determines the properties of linear attention. Different feature maps approximate different similarity functions, each with trade-offs between computational efficiency and how well they mimic softmax behavior.

The connection between feature maps and attention kernels comes from the theory of reproducing kernel Hilbert spaces (RKHS). A kernel function $\kappa(x, y)$ measures similarity between two vectors, and Mercer's theorem tells us that any positive semi-definite kernel can be written as:

$$\kappa(x, y) = \langle \phi(x), \phi(y) \rangle$$

where:

• $\kappa: \mathbb{R}^d \times \mathbb{R}^d \to \mathbb{R}$: a kernel function that takes two vectors and returns a scalar similarity score
• $x, y \in \mathbb{R}^d$: input vectors to compare
• $\phi: \mathbb{R}^d \to \mathbb{R}^D$: a feature map that lifts inputs to a (possibly infinite-dimensional) feature space
• $\langle \cdot, \cdot \rangle$: the inner product in the feature space

This result, known as Mercer's theorem, tells us that for any valid kernel function, there exists a corresponding feature map. Choosing a feature map implicitly defines a similarity kernel, and vice versa.

Standard softmax attention uses an exponential kernel (ignoring the scaling factor):

$$\kappa_{\text{softmax}}(q, k) = \exp(q^T k)$$

where $q^T k$ is the dot product between query and key vectors. The exponential function ensures the similarity is always positive and amplifies differences between scores.

To approximate this with linear attention, we need a feature map $\phi$ such that:

$$\phi(q)^T \phi(k) \approx \exp(q^T k)$$

Finding such a feature map is the central challenge of linear attention design. The closer the approximation, the more closely linear attention mimics softmax behavior.

Several feature maps have been proposed, each with different properties:

1. ELU + 1 (Efficient Linear Units)

$$\phi(x) = \text{ELU}(x) + 1 = \begin{cases} x + 1 & \text{if } x > 0 \\ e^x & \text{if } x \leq 0 \end{cases}$$

This simple feature map ensures all features are positive (required for the normalization to work properly). It doesn't approximate the exponential kernel well but is computationally cheap.

2. ReLU

$$\phi(x) = \max(0, x)$$

Even simpler than ELU, but can produce zero features, potentially causing numerical issues in the denominator.

3. Random Fourier Features (Performer)

$$\phi(x) = \frac{1}{\sqrt{m}} \left[ \sin(\omega_1^T x), \cos(\omega_1^T x), \ldots, \sin(\omega_m^T x), \cos(\omega_m^T x) \right]$$

where:

• $x \in \mathbb{R}^{d_k}$: the input vector (query or key)
• $m$: the number of random frequency vectors
• $\omega_i \in \mathbb{R}^{d_k}$: random frequency vectors sampled from a standard normal distribution $\mathcal{N}(0, I)$
• $\omega_i^T x$: the projection of input onto the $i$-th random direction
• $\frac{1}{\sqrt{m}}$: normalization factor to control the variance of the approximation

The output has dimension $2m$ (sine and cosine for each frequency). This provides an unbiased approximation to the Gaussian kernel. With additional scaling, it can approximate the softmax kernel.

4. Positive Random Features (FAVOR+)

To ensure non-negativity while approximating the exponential kernel, Performer uses:

$$\phi(x) = \frac{\exp(-\|x\|^2/2)}{\sqrt{m}} \left[ \exp(\omega_1^T x), \ldots, \exp(\omega_m^T x) \right]$$

where:

• $x \in \mathbb{R}^{d_k}$: the input vector
• $\|x\|^2 = \sum_{i=1}^{d_k} x_i^2$: the squared Euclidean norm of the input
• $\exp(-\|x\|^2/2)$: a normalization factor that depends on the input magnitude
• $\omega_i \sim \mathcal{N}(0, I)$: random projection vectors sampled from a standard normal distribution
• $\exp(\omega_i^T x)$: ensures all features are strictly positive
• $m$: the number of random features

This construction gives $\mathbb{E}[\phi(q)^T \phi(k)] = \exp(q^T k)$, an unbiased estimator of the softmax kernel. The key insight is that by using exponentials instead of sines and cosines, all features remain positive, which is necessary for the normalization in attention to work correctly.