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.