Mamba (Selective State-Space Model)

Definition

Mamba (Gu & Dao, 2023; arXiv:2312.00752) is a selective state-space model (S6) that makes the SSM parameters $B,C,\Delta$ functions of the input, recovering the content-based routing that plain linear/LTI models lack. It keeps linear-time training (via a hardware-aware parallel scan) and constant-state, linear-time inference, matching or beating Transformers of similar size on language, audio, and genomics while scaling far better in sequence length.

Key math / How it works

S4’s efficiency came from being LTI — the same dynamics at every position — which is also its weakness: it cannot decide, from the current token’s content, what to remember or ignore. Mamba makes the parameters input-dependent:

$B_t=s_B(x_t),C_t=s_C(x_t),{\Delta }_t=\mathrm{softplus}(s_{\Delta }(x_t)),$

with $A$ learned but its effective discretization ${\bar{A}}_t=\exp ({\Delta }_{t}A)$ now input-dependent through ${\Delta }_t$. Intuition: a large ${\Delta }_t$ writes the current token strongly into state (and decays old state) — “pay attention”; a small ${\Delta }_t$ makes the model coast — “ignore.” This is a per-token content-based forget/keep gate, the routing that vanilla linear attention (fixed $\varphi$, no gating) does not have.

Cost of selectivity. Once $B_t,C_t,{\bar{A}}_t$ vary with $t$ the system is non-LTI, so the convolutional form (and FFT training) is lost. Mamba recovers parallel training with a hardware-aware parallel (associative) scan kept in SRAM. Complexity: train $O(bsNd)$ parallel scan; inference state $O(Nd)$ per layer, constant in $s$. At $d=4096$, $L=32$, $N=16$, fp16, one stream’s Mamba state is $\approx 0.17$ GB vs. a 128k-context attention KV cache of $\approx 67$ GB — and it does not grow.

Why it matters

Mamba is the flagship attention-alternative: constant-state, linear-time, with content selectivity closing much of the quality gap to softmax attention. But the gap is narrowed, not erased — a fixed-size state is still a lossy compression, so pure SSMs trail full attention on verbatim recall (copying, associative recall, needle-in-a-haystack). This drives hybrids: interleaving a few full-attention layers among mostly Mamba/linear layers (Jamba, arXiv:2403.19887; DeepMind Griffin/Hawk, arXiv:2402.19427) recovers most recall at a large efficiency win [verify exact layer ratios/benchmarks — moving fast]. [recent]

Taught in

• 07-linear-and-subquadratic — §3.4 selectivity (S6) and the parallel scan; §3.5 hybrids.

See also

• state-space-models — the S4/LTI foundation Mamba generalizes
• linear-attention — the other constant-state family it improves on
• kv-cache · attention-glossary · attention-mechanisms