Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Transformers won foundation modeling with attention: every token can look at every other token. That power is also the bill — training cost grows with sequence length squared, and autoregressive decode must carry a growing key/value (KV) cache. A long line of subquadratic models (linear attention, gated convolutions, structured state space models) tried to keep a compact state instead. On language, they usually lost.

Mamba’s punchline is simple enough to draw: prior structured SSMs were linear time-invariant (LTI) — their dynamics did not depend on the token content — so they could not selectively remember or forget. Make the SSM parameters functions of the input, compute the resulting time-varying recurrence with a hardware-aware scan, and wrap it in a simplified block. The result is a homogeneous architecture without attention that the authors report matches much larger Transformers on language while scaling linearly in sequence length.

What this paper explains

Structured state space models (S4 and relatives) map a 1-D sequence through a latent state $h_t \in \mathbb{R}^{N}$ with linear dynamics. After discretization they look like a recurrence

or, when parameters are constant in time, an equivalent global convolution. That dual form made LTI SSMs fast to train. It also capped them: a static kernel cannot decide “this token is noise; skip it” the way attention can.

Gu and Dao ask:

1. What principle separates models that compress context well from those that do not?
2. Can we add content-aware selection to SSMs without giving up hardware efficiency?
3. How simple can the surrounding neural architecture be if the sequence mixer itself is strong?

Their answers are selectivity, a selective scan implementation, and the Mamba block (SSM fused into a gated MLP-style pathway, stacked homogeneously).

Prior limits

• Attention does not compress: decode stores the entire context as KV, so memory and bandwidth grow with length.
• LTI SSMs / global convolutions compress into a fixed-size state and train efficiently, but their $(\Delta, A, B, C)$ do not depend on $x_t$. Synthetic tasks that require content-aware filtering (selective copying) or associative recall (induction heads) expose the failure.
• Earlier attempts at input-dependent SSM pieces often hit the wall that the convolution trick disappears once parameters vary over time — materializing a huge state tensor $(B, L, D, N)$ is expensive on GPUs.
• Hybrid stacks that bolt weak LTI mixers next to attention inherit the same long-context costs attention was meant to avoid.

The mechanism

Selection as compression control

The paper frames sequence modeling as compressing context into a state. Attention is effective because it refuses to compress; RNNs/SSMs are efficient because they must. Selectivity is the control knob: given $x_t$, decide how information propagates along the sequence dimension.

Mechanically, Mamba’s selective SSM (sometimes called S6 in the ablations) makes several parameters input-dependent:

• $B_t = s_B(x_t)$, $C_t = s_C(x_t)$ — linear projections to the state dimension $N$
• $\Delta_t = \tau_\Delta(\operatorname{Broadcast}_D(\operatorname{Linear}_1(x_t)))$ with $\tau_\Delta = \operatorname{softplus}$

$A$ stays structured/parameterized as in prior SSMs; discretization turns $(\Delta_t, A, B_t)$ into time-varying $(\overline{A}_t, \overline{B}_t)$. The recurrence is no longer LTI, so the convolution dual is gone.

Gating connection

When $N=1$, $A=-1$, $B=1$, and $\Delta$ is produced as above, the selective recurrence collapses to classical RNN gating (their Theorem 1):

Large $\Delta_t$ (large $g_t$) focuses on the current token and resets the state; small $\Delta_t$ persists the state and ignores the token. That is the teaching picture behind the slider below.

Hardware-aware selective scan

Naively expanding $h$ to shape $(B,L,D,N)$ blows memory bandwidth. Mamba keeps the expanded state in fast levels of the GPU memory hierarchy (SRAM), loads $(\Delta, A, B, C)$ from HBM, discretizes, runs a parallel associative scan in recurrent mode, and writes only the smaller outputs $y$ back to HBM. Same algorithmic recurrence; different materialization strategy.

Mamba block

Instead of interleaving an H3-style mixer with a separate MLP, Mamba uses one repeated block: expand the model width by factor $E$ (experiments fix $E=2$), apply a local conv + selective SSM on the main branch, gate with SiLU/Swish, project back, residual + norm. Stacking two such blocks matches the parameter pattern of a Transformer’s attention+MLP pair at width $D$.

Algorithm walkthrough

1. Start from a structured SSM $(\Delta, A, B, C)$ with latent width $N$ per channel.
2. Replace static $B, C, \Delta$ with input-dependent maps $s_B, s_C, s_\Delta$ (softplus on $\Delta$).
3. For each timestep, discretize to $(\overline{A}_t, \overline{B}_t)$ (e.g. zero-order hold) and update $h_t = \overline{A}_t h_{t-1} + \overline{B}_t x_t$.
4. Train with a fused selective scan that avoids writing the full state to HBM; infer autoregressively with constant-time-per-step recurrence (no KV cache growth with history length).
5. Embed the selective SSM inside the Mamba block and stack blocks homogeneously — no attention layer required.

The opening figure shows tokens entering a compact state: relevant tokens update the latent cubes; noise tokens are gated out. The mid-article figure contrasts an LTI convolution path with the selective recurrent path.

What to notice when reading

• The synthetic tasks in §3.1 / §4.1 are not toys — they isolate content-aware failure modes that LTI models cannot patch with width alone.
• Selectivity is broader than one architecture: the paper treats it as a principle (what can depend on $x_t$), with S6 + Mamba as the concrete instance.
• The efficiency claim is not “RNNs are back”; it is “expand state in SRAM, scan in parallel, keep outputs small.”
• Ablations (Table 6 in the paper) attribute most of the language-model jump to replacing LTI SSM layers with selective ones, not merely to the outer block cosmetics.

Results and evidence

Numbers below are as reported by Gu & Dao — not re-benchmarked here.

Synthetic stress tests

Language modeling (Pile pretraining + zero-shot)

From their Table 3 (same tokenizer/dataset/training length cohort where noted):

The authors also report that a Mamba-3B matches Transformers about twice its size on their common-sense average, and that Mamba generation throughput is about 5× a same-size Transformer in their A100 decode setup (prompt 2048, generate 128; Hugging Face Transformer baseline).

Other modalities. On audio waveforms and DNA, they report beating prior SSM/Hyena/Transformer baselines and improving as context grows toward million-length sequences (including a speech FID reduction “by more than half” on a challenging generation set in their abstract summary).

Limitations

Compressed from the paper’s scope and what a careful reader should not over-claim:

• Results depend on their training recipes, tokenizers, and evaluation harness; treat Table 3 as evidence inside that protocol.
• Selective SSMs lose the pure convolution training path — performance hinges on the fused scan implementation quality.
• The homogeneous Mamba block is a design choice; later hybrids (attention + SSM) are outside this paper’s main claim.
• Million-length wins are shown most cleanly on synthetic induction and modality-specific long-range tasks; general chat assistants still need systems work beyond the architecture paper.
• Real-valued vs complex SSM parameterizations and many hyperparameters are ablated at ~350M scale; not every knob is swept at 3B.

How to read the paper

1. Abstract + Figure 1 — selectivity + hardware-aware state expansion in one picture.
2. §2 — SSM recurrence, discretization, LTI/convolution dual.
3. §3.1–3.2 — why selection, Algorithms 1–2 (input-dependent $\Delta, B, C$).
4. §3.3 — why the scan must be hardware-aware once LTI is dropped.
5. §3.4–3.5 — Mamba block; gating theorem.
6. §4.1 synthetics, §4.2 language scaling + Table 3, then modality sections you care about.
7. §4.6 ablations — S6 vs LTI is the lever.

Knowledge check

Keep reading

1. Original paper (arXiv:2312.00752) — Figures 1–3, Algorithms 1–2, Table 3.
2. Gu et al., Efficiently Modeling Long Sequences with Structured State Spaces (S4) — arXiv:2111.00396.
3. Related Fanout Daily: Attention Is All You Need, FlashAttention, Switch Transformers, GQA.

Sources

• Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, 2023 — arXiv:2312.00752
• Gu et al., Efficiently Modeling Long Sequences with Structured State Spaces (S4), ICLR 2022 — arXiv:2111.00396
• Official code/checkpoints referenced by the authors: https://github.com/state-spaces/mamba