Linear Attention Mechanisms
Where the quadratic cost in attention actually comes from
Standard self-attention computes $\text{softmax}(QK^T / \sqrt{d_k})V$. The $QK^T$ term produces an $n \times n$ matrix of pairwise scores for a sequence of length $n$ — this is the source of the quadratic time and memory cost. The softmax nonlinearity sitting between $QK^T$ and the multiplication by $V$ is what forces this full $n \times n$ matrix to be materialized: you can't factor $\text{softmax}(QK^T)V$ into a cheaper associative form the way you could without the softmax in between.
The linear attention trick: drop or approximate the softmax
Linear attention replaces the softmax kernel with a different similarity function $\phi(Q)\phi(K)^T$ that can be reassociated: because there's no nonlinearity sitting between $Q$ and $K$, matrix multiplication is associative, so $(\phi(Q)\phi(K)^T)V$ can instead be computed as $\phi(Q)(\phi(K)^T V)$. The second form never materializes an $n \times n$ matrix at all — $\phi(K)^T V$ is only $d \times d$, independent of sequence length — reducing both compute and memory from $O(n^2 d)$ to $O(n d^2)$, which is a large win whenever $n \gg d$ (long sequences).
import numpy as np
def softmax_free_linear_attention(Q, K, V, feature_map):
"""Q, K: (n, d), V: (n, d_v). feature_map: elementwise nonlinearity, e.g. ELU+1.
Reassociating avoids ever forming the (n, n) attention matrix."""
phi_Q = feature_map(Q) # (n, d)
phi_K = feature_map(K) # (n, d)
KV = phi_K.T @ V # (d, d_v) -- NOT (n, n); this is the key trick
numerator = phi_Q @ KV # (n, d_v)
normalizer = phi_Q @ phi_K.sum(axis=0, keepdims=True).T # (n, 1)
return numerator / (normalizer + 1e-6)This reassociation also enables a true recurrent form
What's lost: exact retrieval capability
Softmax attention's key practical strength is that it can sharply attend to one or a few specific tokens out of many, because softmax is a genuinely competitive normalization — a few large scores dominate, the rest are pushed toward zero. Linear attention's kernel-based similarity doesn't have this same sharp, competitive behavior, so empirically it tends to underperform full softmax attention on tasks requiring precise retrieval of a specific fact or token from far back in context, even though it matches or exceeds it on raw throughput and long-sequence memory efficiency.
Softmax attention vs linear attention
| Property | Softmax attention | Linear attention |
|---|---|---|
| Time/memory complexity | O(n^2) | O(n) in sequence length |
| Precise retrieval of specific tokens | Strong — competitive softmax normalization | Weaker — smoother, less sharp weighting |
| Recurrent/streaming form | Not natural (needs full KV cache) | Natural — fixed-size running state |
| Empirical quality (long-context tasks) | Generally higher | Competitive on throughput-sensitive tasks, often lower on retrieval-heavy ones |
Hybrids are the common production answer, not a clean replacement
What's Next?
Linear attention's reassociation trick is one point on a spectrum of attention-efficiency techniques that also includes sparse attention patterns (restricting which pairs attend at all) and windowed/sliding attention — each makes a different tradeoff between full quadratic expressivity and linear-time scalability.
I build these systems professionally.
Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.