Linear Attention › The Associative Property Trick

To understand linear attention, we need to examine why standard softmax attention requires the $n \times n$ matrix in the first place. The culprit is the softmax normalization, which creates dependencies between all positions.

Standard attention computes:

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

where:

• $Q \in \mathbb{R}^{n \times d_k}$: the query matrix containing one row per token
• $K \in \mathbb{R}^{n \times d_k}$: the key matrix containing one row per token
• $V \in \mathbb{R}^{n \times d_v}$: the value matrix containing the information to aggregate
• $d_k$: the dimension of query and key vectors
• $n$: the sequence length (number of tokens)
• The softmax is applied row-wise, normalizing each query's attention distribution

For each query position $i$, the attention output is:

$$\text{output}_i = \frac{\sum_{j=1}^{n} \exp(q_i^T k_j / \sqrt{d_k}) \cdot v_j}{\sum_{j=1}^{n} \exp(q_i^T k_j / \sqrt{d_k})}$$

where:

• $q_i \in \mathbb{R}^{d_k}$: the query vector at position $i$
• $k_j \in \mathbb{R}^{d_k}$: the key vector at position $j$
• $v_j \in \mathbb{R}^{d_v}$: the value vector at position $j$
• $q_i^T k_j$: the dot product measuring similarity between query $i$ and key $j$
• $\exp(\cdot)$: the exponential function ensuring positive weights
• The denominator sums over all $n$ positions, creating the normalization constant

The denominator is the normalization term: to compute the attention weight for any single key position $j$, we need to know the scores for all positions, because the softmax divides by their sum. This creates a global dependency that forces us to compute all $n^2$ query-key products before we can produce any output.

Linear attention's solution is to replace the exponential-and-normalize pattern with a kernel function that decomposes into separate feature maps for queries and keys. Instead of:

$$\text{similarity}(q_i, k_j) = \exp(q_i^T k_j / \sqrt{d_k})$$

We use:

$$\text{similarity}(q_i, k_j) = \phi(q_i)^T \phi(k_j)$$

where:

• $\phi: \mathbb{R}^{d_k} \to \mathbb{R}^{D}$: a feature map that transforms vectors from the original $d_k$-dimensional space to a $D$-dimensional feature space
• $\phi(q_i) \in \mathbb{R}^D$: the transformed query vector
• $\phi(k_j) \in \mathbb{R}^D$: the transformed key vector
• $D$: the dimension of the feature space (often equal to $d_k$ for simple maps, or larger for random feature methods)

The key property is that the similarity decomposes into an inner product of transformed vectors, allowing us to reorder operations using matrix associativity.

With this formulation, the attention output becomes:

$$\text{output}_i = \frac{\sum_{j=1}^{n} \phi(q_i)^T \phi(k_j) \cdot v_j}{\sum_{j=1}^{n} \phi(q_i)^T \phi(k_j)}$$

Since $\phi(q_i)$ does not depend on $j$, we can factor it out of the sum:

$$\text{output}_i = \frac{\phi(q_i)^T \sum_{j=1}^{n} \phi(k_j) v_j^T}{\phi(q_i)^T \sum_{j=1}^{n} \phi(k_j)}$$

This factorization is the key to linear attention. The numerator contains $\phi(q_i)^T$ multiplied by a sum that doesn't depend on $i$. The denominator has the same structure. Both sums can be computed once and reused for every query position.

Let's define:

• $S = \sum_{j=1}^{n} \phi(k_j) v_j^T$: a matrix of shape $(D \times d_v)$
• $Z = \sum_{j=1}^{n} \phi(k_j)$: a vector of shape $(D,)$

Then the output for each position simplifies to:

$$\text{output}_i = \frac{\phi(q_i)^T S}{\phi(q_i)^T Z}$$

where:

• $\phi(q_i)^T S$: a matrix-vector product of shape $(1 \times D) \cdot (D \times d_v) = (1 \times d_v)$, giving the unnormalized output
• $\phi(q_i)^T Z$: a dot product of shape $(1 \times D) \cdot (D \times 1) = 1$, giving the normalization scalar

Computing $S$ requires $O(nd \cdot D)$ operations (summing $n$ outer products of dimension $D \times d_v$). Computing each output requires $O(D \cdot d_v)$ operations. With $n$ outputs, the total complexity is $O(nd^2)$ assuming $D$ is proportional to $d$.