Mamba SSM Math

The state-space recurrence at the core

A state-space model maps an input sequence to an output through a hidden state h_t that is updated one step at a time. In discrete form the two equations are the whole engine:

h_t = Ā h_(t-1) + B̄ x_t     (state update)
y_t = C h_t                (readout)

Here x_t is the input at step t, h_t is the hidden state (a vector of size N, the state dimension), and y_t is the output. Ā is an N×N transition matrix that decides how much of the old state survives; B̄ writes the new input into the state; C reads the state out. The bar on Ā and B̄ marks them as discretized versions of underlying continuous parameters — more on that next. Structurally this is a linear RNN: no non-linearity sits between h_(t-1) and h_t, which is exactly the property that will later let us compute all steps in parallel instead of strictly left-to-right.

From continuous to discrete: the step size Delta

Mamba does not learn Ā and B̄ directly. It learns a continuous-time system with matrices A and B and then discretizes it with a per-step timescale Δ — the amount of ‘time’ one token advances. The standard zero-order-hold rule gives:

Ā = exp(Δ A)
B̄ = (Δ A)^(-1) (exp(Δ A) − I) · Δ B   ≈  Δ B

The intuition: Δ is a knob on the recurrence’s memory. A small Δ makes Ā = exp(Δ A) ≈ I, so the state barely changes and the current input is largely ignored — the model ‘holds’ its state. A large Δ pushes Ā toward its decay and lets B̄ write the input in strongly — the model ‘refreshes’ on this token. So Δ behaves like a continuous, learnable gate between remembering and updating. In a plain SSM Δ is fixed; the leap Mamba makes is to compute it, and the write/read maps, from the data itself.

The selectivity innovation: input-dependent B, C, and Delta

This is the idea that names the model. In a classical SSM the parameters are time-invariant: the same A, B, C, Δ act on every token, so the system cannot treat one token differently from another based on content. Mamba makes three of them functions of the input:

B_t = Linear_B(x_t)        C_t = Linear_C(x_t)
Δ_t = softplus(Linear_Δ(x_t))    (kept positive)

Now the recurrence reads h_t = Ā_t h_(t-1) + B̄_t x_t with Ā_t = exp(Δ_t A). Because B̄_t, C_t, and Δ_t change token by token, the dynamics are data-dependent: the model can, on the fly, decide to ignore a filler token (drive Δ_t small so the state passes through untouched) or to latch onto a salient one (large Δ_t, strong B̄_t). Notably A itself stays input-independent — but multiplying by the input-dependent Δ_t inside exp(Δ_t A) makes the effective transition selective anyway.

Why selectivity is the whole point

The linear-time-invariant SSM has a fatal weakness for language: it processes every token with identical dynamics, so it cannot perform content-based reasoning — it cannot look at a token and decide whether this one matters. The canonical stress test is selective copying: reproduce a few marked tokens scattered in a stream of noise. A fixed SSM blurs signal and noise together because its B̄ writes everything with the same strength; attention solves it trivially because it can point directly at the marked tokens. Selectivity gives the SSM the same power without attention: an input-dependent B̄_t can gate noise out of the state and salient tokens in, and an input-dependent C_t can choose what to surface at readout. The cost is that the parameters now vary with t, so the elegant convolutional shortcut that a time-invariant SSM enjoys no longer applies — which is why Mamba needs a different fast algorithm, covered below.

Linear scaling and a constant-size state

Run the recurrence and count the work. Each step multiplies an N×N-structured transition, writes the input, and reads out — a fixed amount of arithmetic that does not grow with how many tokens came before. Do that for L tokens and total compute is O(L · N · D) for state size N and model width D — linear in sequence length. Just as important, at inference the model only ever holds h_t, a vector of size N per channel. Generating token 10,000 costs the same as generating token 10, and needs the same memory.

Contrast the transformer. Attention forms an L×L score matrix — O(L^2) compute — and during generation caches every past key and value, a KV cache that grows linearly with context. Mamba replaces that unbounded, ever-growing cache with a single fixed-size state. The history is not stored token by token; it is compressed into h_t. That compression is the source of both Mamba’s efficiency and its main limitation.

The complexity comparison, side by side

The asymptotics make the trade concrete. For sequence length L, state size N (small — often 16), and model width D:

Worked example: at L = 64K tokens, attention’s L^2 term is roughly 4×10^9 pairwise interactions per head, and the KV cache holds all 64K keys and values. Mamba does ~64K constant-cost recurrence steps and carries one N=16 state vector. The gap widens as context grows — which is exactly why SSMs are attractive for very long sequences where the quadratic term dominates.

The hardware-aware parallel scan

Selectivity broke the convolution trick, so how does Mamba train fast instead of grinding through L sequential steps? The answer is that the recurrence h_t = Ā_t h_(t-1) + B̄_t x_t is an associative scan (a prefix-scan / ‘cumulative’ operation). Any linear recurrence can be computed with a parallel scan in O(log L) sequential depth rather than O(L), letting the GPU chew through the whole sequence in parallel at training time while still running as a simple loop at inference.

Mamba goes further with a hardware-aware implementation. The per-token materialized states are large, so naively writing them to GPU HBM would make the model memory-bandwidth-bound. Mamba fuses the discretization, the scan, and the readout into one kernel that keeps the expanded state in fast on-chip SRAM and recomputes it during the backward pass instead of storing it — the same kernel-fusion and recomputation philosophy behind FlashAttention. This is why the paper stresses ‘hardware-aware’: the algorithm and the memory hierarchy are co-designed, and without it the selective scan would be slow in practice.

The accuracy and throughput trade vs transformers

What do you actually get? On throughput, the win is unambiguous: with a constant-size state and no growing KV cache, Mamba’s generation is several times faster than a similarly sized transformer at long context, and its memory footprint stays flat as the sequence grows. On quality, selective SSMs match or beat transformers of comparable size on language modeling and many long-sequence tasks, closing most of the gap that earlier (non-selective) SSMs left open.

But the compression is a genuine trade. Because history lives in a fixed N-dimensional state rather than an exact per-token cache, tasks that need precise recall of arbitrary earlier tokens — exact copying, in-context lookup of a specific string — are where attention’s explicit memory still has an edge. This is why much production work uses hybrid architectures: mostly Mamba layers for cheap linear-time mixing, with a few attention layers sprinkled in to restore exact recall. You buy linear scaling with a bounded memory budget, and you pay in fidelity of long-range exact retrieval.

Practical notes and common pitfalls

A few things trip people up. First, state size N is small (commonly 16) — Mamba’s capacity comes from having many channels each with their own tiny SSM, not from one huge state; do not conflate N with the model width D. Second, Δ is passed through softplus so it stays positive; a well-initialized Δ range matters, because it sets the default memory horizon before training adapts it. Third, remember that A is not selective — only B, C, and Δ depend on the input — yet Δ_t still makes the effective transition exp(Δ_t A) data-dependent, which is often misremembered. Finally, on CPU-bound and small-model deployment the constant-size recurrent state is a real gift: there is no KV cache to grow, per-token cost is flat regardless of context, and memory is predictable — attractive when you cannot lean on a large GPU and want long context without the quadratic wall.