Linear Attention & State-Space Models (Mamba)

Attention & Transformers names "sparse/linear attention variants for very long contexts" as one line in a list of fixes for attention's quadratic cost. This page is that line made real: the actual kernel-trick math behind linear attention, and the actual mechanism — not just "it's an RNN alternative" — behind Mamba, the state-space model that made this family competitive with attention on language in the first place.

Linear Attention: Reordering the Multiplication

What is it? A reformulation of attention that replaces the softmax similarity function with a kernel feature map, which turns out to make the underlying matrix multiplication associative — and associativity is the entire trick.

How does it work? Standard attention computes $\text{softmax}(QK^T/\sqrt{d_k})V$: because softmax is a nonlinear function applied to $QK^T$, it must be computed first, in full, before $V$ can be multiplied in — that's what forces the full $n \times n$ matrix into memory. Katharopoulos et al. (2020) replace $\exp(q^Tk/\sqrt{d})$ with a similarity of the form $\varphi(q)^T\varphi(k)$, for a feature map $\varphi(x) = \text{elu}(x) + 1$ (elu chosen simply to keep every output positive, since a similarity score should never go negative). Removing the nonlinear softmax wrapper makes the product associative:

$\big(\varphi(Q)\varphi(K)^T\big)V = \varphi(Q)\big(\varphi(K)^TV\big)$

The right-hand side is the whole point: $\varphi(K)^TV$ is a $d \times d$ matrix — its size depends only on the feature dimension $d$, never on sequence length $n$. Compute that first, once, and every query just multiplies against it.

Why is it useful? This flips the complexity from $O(n^2d)$ to $O(nd^2)$ — for the long sequences this is built for, $d \ll n$, so this is a real asymptotic win, not a constant-factor tweak. The paper's own benchmark reports up to 4000× faster autoregressive generation on very long sequences, at comparable quality to standard attention.

Limitation: The associativity that makes this fast is also what makes it lossy — softmax's exponential sharply amplifies the highest-similarity keys and suppresses the rest (a soft version of "attend mostly to the one or two most relevant tokens"), while $\varphi(q)^T\varphi(k)$ with a simple positive feature map produces a much flatter, less selective similarity distribution. This isn't a minor implementation detail; it's the direct mathematical reason linear attention variants generally underperform full softmax attention on tasks that need to pick out one precise past token from many similar-looking ones — the trade-off this whole page keeps coming back to.

The Same Trick, Seen as a Recurrence

The associativity above has a second consequence worth making explicit: it turns attention into something that looks exactly like an RNN during generation.

What is it? Autoregressive linear attention, written as a running state update instead of one big matrix multiplication.

How does it work? Define a running state matrix $S_i = \sum_{j \le i} \varphi(k_j)v_j^T$ and a running normalizer $z_i = \sum_{j \le i}\varphi(k_j)$. Generating token $i$ only needs the update, not a re-scan of every prior token:

$S_i = S_{i-1} + \varphi(k_i)v_i^T \qquad z_i = z_{i-1} + \varphi(k_i) \qquad y_i = \frac{\varphi(q_i)^TS_i}{\varphi(q_i)^Tz_i}$

Why is it useful? $S_i$ is a fixed-size $d \times d$ matrix regardless of how many tokens have been generated — generating token 10,000 costs exactly as much as generating token 10, with no growing cache to scan. Standard attention's KV cache, by contrast, grows linearly with every generated token (see the KV cache problem) — this recurrent view is a genuinely different way of paying for long context: constant memory per step, instead of a cache that keeps growing.

Limitation: $S_i$ is the entire memory of everything the model has generated so far, compressed into one fixed-size matrix — nothing about a specific earlier token is separately retrievable from $S_i$ the way it's separately retrievable from an uncompressed KV cache. This is the same compression trade-off as the section above, now made unavoidable: there is no "look up token 47's key exactly," only "read whatever the running sum currently encodes."

That compression-for-speed trade is exactly the design space state-space models occupy too — from a completely different starting point.

State-Space Models: A Different Foundation

What is it? A sequence model built from classical control theory's state-space representation, rather than from attention at all — the mathematical ancestor of Mamba.

How does it work? A continuous-time linear state-space model defines a hidden state $h(t)$ evolving as $h'(t) = Ah(t) + Bx(t)$, with output $y(t) = Ch(t)$. To run on discrete tokens, Gu & Dao (2023) (building on Gu et al.'s earlier S4) discretize this with a step size $\Delta$ via zero-order hold: $\bar{A} = \exp(\Delta A)$, $\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I)\cdot\Delta B$, giving a clean recurrence:

$h_t = \bar{A}h_{t-1} + \bar{B}x_t \qquad y_t = Ch_t$

Because $\bar A, \bar B, C$ don't change over time, this recurrence can also be computed as one big global convolution — which is what let S4 train efficiently in parallel, the same way attention does.

Why is it useful? $h_t$ is a fixed-size vector — same constant-memory-per-step property linear attention's $S_i$ has, and the same linear-time-in-sequence-length training the convolutional view provides.

Limitation, the one that mattered most: because $A$, $B$, and $C$ are fixed constants, S4 has no way to let the content of $x_t$ influence what gets remembered or forgotten — every input is processed through the identical fixed filter. Gu & Dao's own paper names this directly as S4's core weakness: an "inability to perform content-based reasoning." A model that can't selectively decide "this token matters, remember it" vs. "this token is filler, forget it" struggles on exactly the tasks language modeling needs — which is precisely the gap Mamba's contribution closes.

Mamba's Selective Mechanism (S6)

What is it? Mamba's actual innovation — not "an SSM," but a selective SSM, where the parameters that were fixed constants in S4 become functions of the input at every timestep.

How does it work? Mamba (S6) makes $B$, $C$, and the step size $\Delta$ input-dependent: $B \leftarrow \text{Linear}_N(x_t)$, $C \leftarrow \text{Linear}_N(x_t)$, $\Delta \leftarrow \text{softplus}(\text{parameter} + \text{Linear}_1(x_t))$. Each now varies per timestep instead of staying fixed for the whole sequence — the model can dilate $\Delta$ to "linger" on an important token (letting it dominate the state update) or contract it to let a filler token pass through with almost no effect on $h_t$, and adjust $B$/$C$ to control what gets written into and read out of the state, all conditioned on the actual content of $x_t$.

Why is it useful? This is the specific fix for S4's named weakness above: selectivity gives the model a genuine, content-driven way to choose what enters its fixed-size memory, which is what finally made SSMs competitive with attention on real language modeling — not a bigger model or more training, a different recurrence.

Limitation (a real engineering cost, not a modeling one): making $B$, $C$, $\Delta$ input-dependent breaks the trick that made S4 fast — a per-timestep-varying recurrence can no longer be rewritten as one fixed global convolution, so the parallel-training shortcut S4 relied on is gone.

Gu & Dao report linear-time scaling in sequence length, up to 5× higher inference throughput than same-size Transformers, and a 3B-parameter Mamba matching Transformers roughly twice its size — real numbers, not just an asymptotic argument.

The Real Trade-off: Compression vs. Exact Recall

Every mechanism on this page pays for speed and constant memory the same way: by compressing all of history into one fixed-size state instead of keeping it around uncompressed. That trade-off has a real, measured cost.

What is it? Jelassi, Brandfonbrener, Kakade, Malach & Malach ("Repeat After Me," ICML 2024) directly test this: can SSMs copy or retrieve information as reliably as attention?

How does it work? They prove a 2-layer Transformer can copy input strings of exponential length, while SSMs are provably bounded by the size of their fixed latent state — a state of size $N$ simply cannot losslessly encode more than $O(N)$ information, no matter how it's trained. They confirm this empirically on real pretrained models too: Mamba-family models "dramatically underperform" Transformer-family models (Pythia) specifically on copying and retrieval tasks.

Why is it useful (as a design signal, not a failure): this pinpoints exactly why the trade-off exists, mechanistically — attention keeps an explicit, uncompressed record of every past token (the KV cache) and can retrieve any one of them exactly; an SSM's hidden state is a lossy compression of everything seen, by construction. Neither is a "better" architecture in the abstract — they're solving for different constraints. This is precisely the finding behind the "hybrid architectures" trend replacing pure-Transformer or pure-SSM designs in current models.

Limitation: this doesn't mean SSMs are simply "worse" — for tasks that don't need needle-in-a-haystack exact retrieval (broad long-range context mixing, streaming/continuous inference, extremely long sequences where a growing KV cache is itself the bottleneck), the fixed-size state's constant memory is a genuine win attention cannot match.

Real Hybrid Architectures: Using Both on Purpose

What is it? Rather than choosing one family, several current models explicitly interleave attention layers and SSM (or SSM-like) layers, using each for what it's actually good at.

How does it work? Jamba (AI21 Labs) interleaves Transformer and Mamba layers with mixture-of-experts capacity added on top, reaching strong results at up to 256K tokens of context while fitting a single 80GB GPU — explicitly trading for high throughput and small memory footprint alongside competitive quality. Griffin (Google DeepMind) takes a related but distinct approach: it interleaves a gated linear recurrence (RG-LRU, an SSM-family mechanism) with local sliding-window attention rather than full global attention — shipped as the open-weight RecurrentGemma models. Griffin matches Llama-2's quality using 6× fewer training tokens, with the fixed-size recurrent state giving it lower latency and higher throughput at inference.

Why is it useful? Both designs make the same bet the copying-benchmark result above predicts: put cheap, constant-memory SSM/recurrent layers where long-range mixing is enough, and reserve the small number of full/local-attention layers for the positions that need precise, exact retrieval — capturing most of attention's retrieval quality at a fraction of its memory cost.

Limitation: hybrid designs add a real architecture-search dimension pure Transformers don't have — which layers should be attention vs. SSM, and in what ratio, is now a genuine design choice with no universally correct answer yet, tuned empirically per model rather than derived from first principles.

Where This Fits Today

Pure attention remains the safest default when exact long-range retrieval matters and context lengths are moderate — it's still what GPT, LLaMA, and Claude-style decoder-only models are built on. Pure linear attention and pure SSMs are the right reach when context is extremely long and cheap, constant-memory inference matters more than exact recall of any single past token — streaming applications, very-long-document processing. Hybrid designs (Jamba, Griffin) are the pragmatic middle: most of a Transformer's retrieval quality, most of an SSM's memory and throughput advantage, at the cost of one more architecture-search decision.

Next: Vision Architectures — the ViT variants and detection/segmentation models built on the same attention foundations; or back to Attention & Transformers for the full model-family comparison this page grew out of.