Linear Attention Explained: The Core Concept and the Math, Without the Fluff

A self-contained walkthrough of the key idea from "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention" (Katharopoulos et al., ICML 2020, arXiv:2006.16236)

This article is based on a 2020 paper by Katharopoulos et al. which shows that by replacing the softmax in standard attention with a kernel dot product, you can reduce the complexity of self-attention from $O(N^2)$ to $O(N)$, and that this automatically makes the transformer equivalent to a recurrent neural network at inference time.

The paper is clean, but it mixes a lot of background and experimental detail with the core idea. This article cuts the fluff, gets straight to the core concept and math, and explains every variable as it appears. Nothing is left undefined.

1. What Standard Attention Does

Let $\mathbf{x} \in \mathbb{R}^{N \times F}$ be a sequence of $N$ feature vectors, each of dimension $F$. A transformer layer projects this input into three matrices:

where $W_Q, W_K \in \mathbb{R}^{F \times D}$ and $W_V \in \mathbb{R}^{F \times M}$ are learned projection matrices. $D$ is the key/query dimension and $M$ is the value dimension.

For a general non-negative similarity function $\text{sim}(q, k)$, the attention output at position $i$ is:

$Q_i$ denotes the $i$-th row of $Q$ (a $D$-dimensional vector), same for $K_j$ and $V_j$. The output $V'_i \in \mathbb{R}^M$ is a weighted average of all value vectors.

For softmax attention, the similarity is:

The full attention matrix $\text{softmax}(QK^T / \sqrt{D})$ has shape $N \times N$. Computing and storing it costs $O(N^2)$ in both time and memory. This is the bottleneck.

2. The Key Idea: Factorize sim with a Feature Map

Instead of computing similarity directly, express it as an inner product of explicit feature representations:

where $\phi : \mathbb{R}^D \to \mathbb{R}^C$ is a feature map that produces non-negative outputs, and $C$ is the feature map output dimension.

Substituting into equation (1):

Why does this matter? The two bracketed terms in equation (3) are the same for every query $i$:

• $\sum_j \phi(K_j) V_j^T \in \mathbb{R}^{C \times M}$ — compute once
• $\sum_j \phi(K_j) \in \mathbb{R}^C$ — compute once

Then for each query, apply $\phi(Q_i)$ via a dot product. Total cost: $O(NCM)$ instead of $O(N^2 M)$.

3. Adding Causal Masking (For Autoregressive Use)

For autoregressive generation, position $i$ must only attend to positions $j \leq i$. Equation (1) becomes:

Define two cumulative state quantities:

Equation (4) then simplifies to:

$S_i$ and $z_i$ update in $O(1)$ from the previous step:

The full causal attention pass therefore costs $O(NCM)$ in time, linear in $N$.

4. This Is an RNN

The recurrence above is, by definition, a recurrent neural network. The full transformer layer (including the per-position feedforward sublayer $f_l$) can be written as:

Variable reference:

• $x_i$: input at timestep $i$, shape $1 \times F$
• $W_K, W_V, W_Q$: learned projection matrices (same as before)
• $s_i$: attention memory state, shape $C \times M$
• $z_i$: normalizer memory state, shape $C$
• $f_l$: per-position feedforward sublayer (e.g. a two-layer MLP)
• $y_i$: output at timestep $i$
• $\phi$: feature map applied to the key and query projections

At training time, the cumulative sums can be parallelized (like a prefix scan), so GPU efficiency is preserved. At inference time, you maintain $s_i$ and $z_i$ as a fixed-size state and update one token at a time in $O(1)$ per step. No growing KV cache.

5. The Feature Map

For the kernel decomposition in equation (2) to work, $\phi$ must produce non-negative outputs. The paper uses:

where $\text{elu}$ is the exponential linear unit:

Adding 1 shifts the output range to always be positive. ELU is preferred over ReLU because ReLU sets gradients to zero for negative inputs, which can stall training.

This feature map cannot exactly recover softmax attention (the exact feature map for the exponential kernel is infinite-dimensional). But empirically, it converges to similar performance.

6. Complexity Summary

$C$ = feature map dimension, $D$ = key/query dimension, $M$ = value dimension, $N$ = sequence length.

Closing

The core insight here is not a new architecture so much as a change of perspective. Self-attention was always computing a weighted combination of values, and that weighting was always a similarity function. By choosing a similarity function that factors as a dot product of feature maps, matrix associativity does the rest, turning a quadratic computation into a linear one, and making the causal version a literal RNN.

The tradeoff is that you lose the ability to exactly replicate softmax attention, and the feature map choice matters for performance. But the paper shows this is a small price to pay for orders-of-magnitude faster inference.

Full paper: arXiv:2006.16236 | Code: linear-transformers.com