Linear Attention

What You’ll Learn

You’ll learn why standard attention’s $O(n^2)$ cost makes long sequences expensive, and how the associativity trick can avoid the $n \times n$ bottleneck. You’ll also see why softmax blocks that shortcut, how kernel functions repair it, and what linear attention trades for speed.

The Wall: Quadratic Attention

Standard self-attention computes:

$\text{Attention}(Q, K, V) = \text{softmax}(QK^\top)V$

The problem is the $n \times n$ matrix $QK^\top$ , where $n$ is the sequence length. We showed that this costs $O(n^2 d)$ in time and $O(n^2)$ in memory.

For short sequences, this is fine. For long ones, it’s a wall. A 100K-token context needs a 100K × 100K attention matrix: 10 billion entries. Can we avoid ever building this matrix?

The Insight: Associativity

Step through the visualization below. It shows how the same computation can be reordered to avoid the $n \times n$ bottleneck.

From Quadratic to Linear

Standard Attention

Q × Kᵀ n×n O(n²d)

softmax n×n O(n²)

× V n×d O(n²d)

Time O(n²d)

Memory O(n²)

The n×n attention matrix must be fully materialized.

Standard attention computes Q×Kᵀ first, producing an n×n matrix. Every query must score against every key before any output is produced.

What to look for: In standard attention, the first operation produces an $n \times n$ matrix (expensive). In linear attention, the first operation produces a $d \times d$ matrix (cheap and fixed). The trick is associativity of matrix multiplication: $(QK^\top)V = Q(K^\top V)$ . Same result, different cost.

A Concrete Example

Consider 3 tokens with $d = 2$ :

$Q = \begin{bmatrix} 0.2 & 0.8 \\ 0.5 & 0.3 \\ 0.9 & 0.1 \end{bmatrix}, \quad K = \begin{bmatrix} 0.1 & 0.7 \\ 0.6 & 0.4 \\ 0.3 & 0.9 \end{bmatrix}, \quad V = \begin{bmatrix} 1.0 & 0.5 \\ 0.2 & 0.8 \\ 0.7 & 0.3 \end{bmatrix}$

Standard order: $QK^\top$ produces a 3×3 matrix (9 entries), then multiply by $V$ .

Reversed order: $K^\top V$ produces a 2×2 matrix (4 entries), then multiply by $Q$ .

With 3 tokens and $d = 2$ , the savings are small. But with 10,000 tokens and $d = 128$ , standard attention builds a 10,000 × 10,000 matrix (100 million entries), while the reversed order builds a 128 × 128 matrix (16,384 entries). That’s 6,100x fewer entries.

The Catch: Softmax Gets in the Way

If attention were just $(QK^\top)V$ , we’d be done. Reorder to $Q(K^\top V)$ and enjoy linear scaling. But standard attention has softmax in the middle:

$\text{softmax}(QK^\top) \cdot V$

Softmax normalizes each row by the sum of all entries in that row. That means every entry in a row depends on every other entry in that row. You can’t compute the output for query $i$ without first computing all $n$ dot products $q_i \cdot k_j$ . Softmax couples everything together and blocks the reordering.

The Fix: Kernel Functions

What if we replaced softmax with something that doesn’t couple the entries? A kernel function $\phi$ applied to Q and K independently:

$\text{softmax}(q_i^\top k_j) \approx \phi(q_i)^\top \phi(k_j)$

If this holds, then the attention output for query $i$ becomes:

$\text{output}_i = \frac{\sum_j \phi(q_i)^\top \phi(k_j) \cdot v_j}{\sum_j \phi(q_i)^\top \phi(k_j)}$

The numerator factors into $\phi(q_i)^\top \left(\sum_j \phi(k_j) v_j^\top \right)$ . That inner sum $\sum_j \phi(k_j) v_j^\top$ is a $d' \times d$ matrix that doesn’t depend on $i$ . Compute it once, reuse it for every query. No $n \times n$ matrix needed.

Standard attention for query $i$ :

$o_i = \frac{\sum_{j=1}^{n} e^{q_i^\top k_j} \cdot v_j}{\sum_{j=1}^{n} e^{q_i^\top k_j}}$

Replace $e^{q_i^\top k_j}$ with $\phi(q_i)^\top \phi(k_j)$ :

$o_i = \frac{\sum_{j=1}^{n} \phi(q_i)^\top \phi(k_j) \cdot v_j}{\sum_{j=1}^{n} \phi(q_i)^\top \phi(k_j)}$

Factor out $\phi(q_i)^\top$ (it doesn’t depend on $j$ ):

$o_i = \frac{\phi(q_i)^\top \sum_{j=1}^{n} \phi(k_j) v_j^\top}{\phi(q_i)^\top \sum_{j=1}^{n} \phi(k_j)}$

Define:

$S = \sum_{j=1}^{n} \phi(k_j) v_j^\top$ (shape: $d' \times d$ , computed once)
$z = \sum_{j=1}^{n} \phi(k_j)$ (shape: $d'$ , computed once)

Then for each query:

$o_i = \frac{\phi(q_i)^\top S}{\phi(q_i)^\top z}$

Computing $S$ costs $O(nd'd)$ . Each query lookup costs $O(d'd)$ , and there are $n$ queries, so the total is $O(nd'd)$ . If $d' = d$ , the total is $O(nd^2)$ : linear in sequence length.

Quick Check

Why can't we just reorder standard attention from (QK⊤)V to Q(K⊤V) and call it linear?

Your turn

Choose an answer to check your understanding.

Attempts: 0

How Much Faster?

Click through different sequence lengths. Watch how the gap between standard ( $O(n^2 d)$ ) and linear ( $O(nd^2)$ ) attention grows as sequences get longer.

Scaling: Standard vs. Linear Attention

Sequence length (n):

d = 128 (typical head dimension)

Standard
O(n²d)

134.2M

Linear
O(nd²)

16.8M

Linear is 8.0x cheaper

At 1024 tokens, linear attention is modestly faster. The real gains come at longer sequences.

What to look for: At $n = d$ (128 tokens), both approaches cost the same. Below that, standard attention is actually cheaper. The crossover happens at $n = d$ , and beyond that, linear attention pulls away fast. At 16K tokens, linear is 128x cheaper.

This reveals an important detail: linear attention isn’t universally faster. It swaps $n^2 d$ for $nd^2$ . When $n < d$ , the standard approach wins. Linear attention only pays off for sequences longer than the model dimension, which is exactly the regime where standard attention becomes painful.

The Trade-off: What Gets Lost

Linear attention isn’t a free upgrade. The kernel $\phi$ approximates softmax, and that approximation has consequences.

What softmax does well: It produces sharp, peaked attention patterns. When the model needs to focus hard on one or two tokens (like resolving a pronoun), softmax concentrates nearly all weight on the relevant entries and suppresses everything else.

What kernels struggle with: Most kernel functions produce smoother, more diffuse attention patterns. The model can still identify relevant tokens, but it can’t suppress irrelevant ones as aggressively. This hurts tasks that require precise, selective attention.

	Standard (softmax)	Linear (kernel)
Time	$O(n^2 d)$	$O(nd^2)$
Memory	$O(n^2)$	$O(nd + d^2)$
Attention pattern	Sharp, selective	Smoother, more diffuse
Best for	Precision tasks, short/medium sequences	Long sequences, streaming
Causal support	Masking (still $O(n^2)$ )	Cumulative sum (true $O(n)$ )

Fill in the Blanks

Linear attention replaces with a kernel function, enabling the computation to be reordered from O(n²d) to O( ). The trade-off is that attention patterns become less .

Fill in every blank, then check your answer.

Real Architectures

Several production-grade architectures use linear or sub-quadratic attention:

Linear Transformers (Katharopoulos et al., 2020): The original linear attention paper. Uses $\phi(x) = \text{elu}(x) + 1$ as the kernel, achieving $O(nd^2)$ complexity. Simple to implement but loses some quality compared to softmax.

Performer (Choromanski et al., 2020): Uses random feature maps (FAVOR+) to approximate the softmax kernel more accurately. Better quality than simple kernels but adds a feature dimension $d'$ that increases constant factors.

Mamba (Gu & Dao, 2023): Takes a different approach entirely. Instead of approximating attention with kernels, it uses a state-space model (SSM) with selective gating. Each token updates a fixed-size hidden state (like an RNN), giving $O(n)$ time and $O(1)$ memory per step. Mamba matches or beats transformer quality on many benchmarks while being much faster for long sequences.

Mamba-2 (Dao & Gu, 2024): Shows that the SSM approach and linear attention are actually connected. The structured state-space duality (SSD) framework unifies them, with the SSM update rule being equivalent to a specific form of linear attention with causal masking.

Causal Linear Attention

For autoregressive generation (where token $i$ can only attend to tokens $1 \ldots i$ ), linear attention has an elegant property. The summary matrix $S_i = \sum_{j=1}^{i} \phi(k_j) v_j^\top$ can be maintained as a running sum:

$S_i = S_{i-1} + \phi(k_i) v_i^\top$

This means causal linear attention naturally works as a recurrence, like an RNN. At each step, update the $d \times d$ state and produce an output. No attention matrix, no cache growing with sequence length. True $O(1)$ per-step memory.

This is why some researchers argue that the line between “attention” and “recurrence” is blurring. Linear attention with causal masking is structurally a linear RNN with a matrix-valued hidden state.

Why Not Just Use Bigger Kernels?

A natural question: if the kernel approximation is the weak point, why not use a higher-dimensional feature map $\phi$ to get a better approximation? The catch is that kernel dimension $d'$ appears in the $O(nd'^2)$ cost. As $d'$ grows, the cost approaches (and eventually matches) standard attention. The whole point was to keep $d'$ small, which limits how well $\phi$ can approximate softmax.

Flash Attention: A Third Path

Flash Attention (Dao et al., 2022) takes yet another approach: keep exact softmax attention but never materialize the $n \times n$ matrix in GPU memory. By tiling the computation and fusing operations, it achieves $O(n^2 d)$ time (same asymptotic cost) but $O(n)$ memory. It doesn’t change the math at all, just how the hardware executes it.

In practice, Flash Attention with exact softmax often beats linear attention variants on quality while being fast enough for most sequence lengths. Linear attention’s advantage emerges only for very long sequences (tens of thousands of tokens) where even Flash Attention’s $O(n^2)$ time becomes the bottleneck.

Quick Summary

Step	What happens	Intuition
Standard attention	$(QK^\top)V$ : build $n \times n$ matrix first	Every query scores against every key
The reorder trick	$Q(K^\top V)$ : build $d \times d$ summary first	Summarize all key-value pairs, then each query reads the summary
Softmax blocks this	Softmax couples entries within each row	Can’t summarize without seeing all scores first
Kernel replaces softmax	$\phi(Q)\phi(K)^\top V$ : independent per entry	Apply $\phi$ to each vector separately, then reorder freely
Linear attention	$\phi(Q) \cdot (\phi(K)^\top V)$ : $O(nd^2)$	Summary is $d \times d$ (small, fixed), linear in $n$

Learning Checklist

I can explain why matrix associativity lets us reorder the attention computation
I understand why softmax prevents this reordering
I can describe what a kernel function does and why it enables linear scaling
I know when linear attention is faster and when standard attention wins ( $n$ vs. $d$ )
I can explain the quality trade-off (sharp vs. diffuse attention patterns)
I’m aware of the connection between causal linear attention and RNNs

Now that you understand how attention can be made efficient, see how the KV Cache tackles the same problem from a different angle: caching rather than approximating.