Mamba: State Space Models for Sequence Processing

A ground station tracking an ISRO satellite does not keep a growing logbook of every radar ping it has ever received. It keeps a small state vector — position and velocity, six numbers — and at every new measurement it updates that state and discards the raw reading. This is a state space model: a system where a compact, fixed-size summary is carried forward through time, updated by a simple linear law, instead of re-reading the entire history at every step. Kalman filters, which do exactly this, have flown on spacecraft since the 1960s. Mamba (Gu & Dao, 2023) asks a strange question: what if a language model's hidden state worked the same way? Not a Transformer re-attending to every previous token at every step, but a state vector that gets updated once per token and never grows, no matter how long the sequence gets. The catch, and the reason this took sixty-plus years to become competitive with attention, is that a Kalman filter's update rule is the same at every timestep — it cannot decide that this particular reading matters more than that one. Mamba's actual contribution is making that update rule content-aware without losing the linear-time efficiency that made state space models attractive in the first place.

The equations underneath: control theory, not NLP

A continuous-time linear state space model is defined by two equations. A hidden state x(t) evolves according to

dx(t)/dt = A·x(t) + B·u(t)

and an output is read off the state by

y(t) = C·x(t)

Here u(t) is the input signal, x(t) is the internal state (for the satellite: position and velocity), A describes how the state evolves on its own (physics: velocity changes position), B describes how new input perturbs the state (a thruster firing), and C reads out what we actually observe. In control theory A, B, C are matrices sized by the state dimension N and input/output dimension D. For sequence modelling, u(t) becomes the token embedding at position t, and the question becomes: can a rule this simple — three fixed matrices and a linear update — compress an entire sentence, paragraph, or genome into a state that a decoder can read?

From a continuous law to a token-by-token recurrence

Text arrives in discrete steps, so the continuous equation has to be discretized. Given a step size Δ (how much "time" one token represents), the exact solution of the ODE over the interval, assuming the input is constant across that interval, is

A_d = e^ΔA      B_d = A⁻¹(e^ΔA − I)B

This is the standard zero-order hold discretization, and it turns the ODE into an exact recurrence over tokens:

h_t = A_d·h_t−1 + B_d·x_t      y_t = C·h_t

Notice what this recurrence looks like: it is structurally an RNN update. That resemblance is exactly the point of the misconception addressed later — and exactly why it is misleading.

Why this recurrence was interesting even before Mamba

The predecessor architecture, S4 (Structured State Space Sequence model), fixed A, B, C once as constants that do not change across the sequence — a linear time-invariant (LTI) system. An LTI recurrence has a remarkable property: because A_d and B_d are the same at every step, the entire output sequence can be written as one fixed convolution kernel applied to the input, y = x * K, where K is derived once from A_d, B_d, C. A convolution over a full sequence can be computed with an FFT in O(L log L) time, fully parallel across positions — no step-by-step loop needed. So S4 could be trained like a convolutional network (parallel, fast) and deployed like a recurrent network (constant memory per step, no need to store the whole sequence). A also had to be initialized carefully — S4 uses the HiPPO scheme, which sets up the state's decay rates so that a fixed-size vector can represent a long history without the vanishing-gradient collapse that plagues randomly initialized RNNs. This recurrence/convolution duality, not the selection mechanism, is what state space models inherited from control theory and what made them viable competitors to attention in the first place.

The problem the LTI assumption creates

A fixed A_d, B_d, C means the model treats every token identically: the state always decays at the same rate, every input is injected with the same weight, regardless of what the token actually is. Consider a task where a sequence contains a few "important" tokens scattered among long runs of padding, and the model must copy only the important ones while ignoring everything else, regardless of how much padding separates them. A Transformer solves this trivially: attention scores are computed from content (query·key), so the model can learn to attend hard to the important tokens and ignore the rest, independent of distance. An LTI SSM cannot: its "attention" to a given position is baked into A_d and decays purely with elapsed time, never with content. It cannot choose to remember one token and forget an adjacent one of similar magnitude. This is not a minor gap — it is precisely the class of task on which pre-Mamba SSMs measurably underperformed Transformers.

Mamba's fix: make the update itself depend on the token

Mamba's change is to stop treating Δ, B, and C as fixed parameters and instead compute them from the current input:

Δ_t, B_t, C_t = Linear(x_t)

The raw matrix A stays a fixed, learned, per-channel constant — it is not predicted from the input. But because A_d = e^ΔA, and Δ is now a function of the token, the effective discretized transition A_d,t = e^Δ_tA varies from token to token even though A itself never moves. This is the entire selection mechanism: a large Δ_t pushes A_d,t toward the input-dominated regime (the token strongly overwrites the state — "pay attention, remember this"), while a small Δ_t pushes A_d,t toward the identity (the state is nearly frozen — "this token carries nothing new, preserve what I already have"). The model learns, per token and per channel, how much of the past to keep versus how much new information to write in — a content-based gate implemented entirely through the step size of a differential equation, not through a query-key dot product. The paper's authors found that letting Δ alone carry the input-dependence (rather than also making A itself a function of the input) was already sufficient, which keeps the parameter count and the recurrence structure simple.

What selectivity costs, and how the scan buys it back

The FFT convolution trick that made S4 fast to train relied on A_d, B_d being identical at every position, so one kernel could describe the whole sequence. Once A_d,t and B_d,t change every step, there is no single kernel, and the convolution view collapses — the naive fallback is a strictly sequential Python-style for-loop over L tokens, which is far too slow to train at scale on a GPU. Mamba's second contribution, arguably as important as selectivity itself, is recognizing that the recurrence is still associative: combining two consecutive steps (A₁, B₁) then (A₂, B₂) into one equivalent step (A₂A₁, A₂B₁+B₂) is a well-defined operation regardless of grouping order. An associative sequence of operations can be reduced with a parallel scan (the same prefix-sum trick used to parallelize cumulative sums): pair up neighbours, combine, repeat, halving the number of remaining elements each round, so the whole sequence resolves in O(log L) parallel steps and O(L) total work instead of a serial O(L) loop with L sequential dependencies. On top of this, Mamba is implemented as a hardware-aware CUDA kernel: the expanded per-step states are computed and combined inside the GPU's fast on-chip SRAM and only the small final outputs are written back to slower HBM memory, instead of materializing every intermediate state in HBM the way a naive implementation would. This is the same insight behind FlashAttention — the algorithm's asymptotic complexity and its real wall-clock speed are two separate engineering problems, and Mamba's actual throughput advantage depends on solving both.

The full Mamba block

One layer combines all of this into a block with two parallel branches (see the diagram below). The input is expanded by a linear projection into two paths. The content path runs through a short causal 1-D convolution (to mix in a few neighbouring tokens before the SSM), a SiLU activation, and then the selective SSM itself: Δ_t, B_t, C_t are computed from the convolved input, and the selective scan produces y_t. The gate path is a second linear projection followed by SiLU, producing z_t. The two paths are combined by elementwise multiplication, y_t × z_t, then projected back down to the model dimension D and added to the block's input as a residual connection, exactly as in a Transformer block. Each of the D channels runs its own independent SSM with a small per-channel state of size N (Mamba typically uses N = 16), so the total recurrent state for the whole layer is only D×N numbers — a few thousand floats for a realistic model, regardless of whether the sequence is 100 tokens or 100,000.

Worked example: tracing the selection mechanism by hand

To make the mechanism concrete, strip it down to a single channel with a 1-dimensional state (D = 1, N = 1). Take A = −1 (a state that decays toward zero if left alone — a "leaky" memory), B = 1, C = 1, and a sequence of three scalar inputs x = [1, 2, 0], with h₀ = 0. First, the fixed-step case, matching how S4 would run it: Δ_t = 1 for every t.

Discretize once: A_d = e^1×(−1) = e⁻¹ ≈ 0.3679, and B_d = (A_d−1)/A×B = (0.3679−1)/(−1)×1 = 0.6321. Now step through the recurrence:

t=1: h₁ = 0.3679×0 + 0.6321×1 = 0.6321, so y₁ = 0.6321
t=2: h₂ = 0.3679×0.6321 + 0.6321×2 = 0.2325 + 1.2642 = 1.4968, so y₂ = 1.4968
t=3: h₃ = 0.3679×1.4968 + 0.6321×0 = 0.5506 + 0 = 0.5506, so y₃ = 0.5506

Even though x₃ = 0 carries no information, the fixed decay still erases 63.2% of what h₂ was holding (A_d = 0.3679 keeps only 36.8%), purely because one more timestep has elapsed. Now the selective case: Δ₁ = Δ₂ = 1 (unchanged, since x₁ and x₂ are informative), but Δ₃ = 0.01, reflecting a selection rule that has learned to shrink Δ when the input carries no new information. Steps 1 and 2 are identical to before (h₂ = 1.4968). At t=3: A_d = e^0.01×(−1) = e^−0.01 ≈ 0.9900, so h₃ = 0.9900×1.4968 + (tiny B_d)×0 ≈ 1.4819, giving y₃ ≈ 1.4819 — 99.0% of h₂ is preserved instead of 36.8%. The same architecture, the same A, B, C, produces a completely different outcome at t=3 purely because Δ was allowed to look at the input. This is the entire selective-copying capability in miniature: the model chooses, per token, whether this step should overwrite memory or leave it alone.

The same computation, traced in code (verified to produce exactly these values, rounded to four decimal places):

import math

def discretize(A, B, delta):
    A_d = math.exp(delta * A)
    B_d = (A_d - 1) / A * B
    return A_d, B_d

def selective_ssm(A, B, C, xs, deltas):
    h = 0.0
    ys = []
    for x, delta in zip(xs, deltas):
        A_d, B_d = discretize(A, B, delta)
        h = A_d * h + B_d * x
        ys.append(C * h)
    return ys

xs = [1, 2, 0]
A, B, C = -1, 1, 1

fixed_y     = [round(v, 4) for v in selective_ssm(A, B, C, xs, [1, 1, 1])]
selective_y = [round(v, 4) for v in selective_ssm(A, B, C, xs, [1, 1, 0.01])]

print("fixed step:    ", fixed_y)
print("selective step:", selective_y)
# fixed step:     [0.6321, 1.4968, 0.5506]
# selective step: [0.6321, 1.4968, 1.4819]

Why linear beats quadratic at scale

Let D be the model's channel width and N the per-channel SSM state size (typically N = 16). For a sequence of length L:

At L = 100,000 tokens, L² is 10¹⁰, while L itself is 10⁵ — a five-order-of-magnitude gap in the sequence-length factor alone. Multiplying through by the D = N-vs-D difference (roughly D×16 for Mamba against D for one Transformer layer's L² term) still leaves Mamba doing on the order of a few thousand times less work for the sequence-mixing step at that length, and — more consequentially for deployment — generating each new token at a cost that never grows, instead of one that grows with every token already generated.

A common misconception

Because the discretized recurrence h_t = A_d·h_t−1 + B_d·x_t looks exactly like the update rule of a vanilla RNN, and because Mamba is often introduced as "linear time like an RNN," students commonly assume Mamba's hidden state grows the way a Transformer's KV cache does — that somewhere it must be storing a proportionally larger memory as the sequence lengthens, since it clearly "remembers" more of a long document than a short one. This is false, and it is the entire point of the architecture: the state h_t is a fixed-size vector of D×N numbers, exactly the same size whether the sequence so far is 10 tokens or 10 million. A Transformer's KV cache, by contrast, is an exact, lossless record that genuinely grows by one key-value pair per token — which is precisely why it eventually exhausts GPU memory on very long contexts. Mamba's state is a compressed summary, not a growing log; what changes with sequence length is how much information gets squeezed into that same fixed-size vector, not the vector's size. This is also the source of a fair trade-off, not just an advantage: because the state is a lossy compression, Mamba can genuinely struggle at tasks needing exact random-access recall of one specific far-back token (something an uncompressed KV cache can always retrieve exactly), even though it wins decisively on raw compute and memory scaling.

Active recall

1. Write the zero-order-hold discretization formulas for A_d and B_d given continuous A, B and step size Δ.
2. For the worked example's system (A=−1, B=1), compute A_d and B_d for Δ=2.
3. Why can't a fixed, non-selective SSM (like S4) solve a task that requires copying only the "important" tokens in a sequence while ignoring padding, regardless of how far apart the important tokens are?
4. What does it mean for the SSM recurrence to be "associative," and why does that property matter for training speed?
5. A Transformer processes a context of L=100,000 tokens. Roughly how many pairwise attention scores does one layer compute (ignoring the D factor), and how does that compare to the sequential work Mamba's scan does for the same L?
6. True or False, with justification: "Mamba's speed advantage comes purely from a better mathematical formula; how the GPU accesses memory is irrelevant."

Answers

1. A_d = e^ΔA, and B_d = A⁻¹(e^ΔA−I)B, obtained by solving the linear ODE exactly over one step assuming the input is held constant across that step.
2. A_d = e^2×(−1) = e⁻² ≈ 0.1353. B_d = (0.1353−1)/(−1)×1 = 0.8647.
3. Because Δ, A_d, B_d are the same constants at every position, the state always decays and updates at a rate that depends only on elapsed time, never on what the token actually is. The model has no mechanism to weight one token more than a neighbouring one of similar magnitude based on content, so it cannot learn "ignore padding, keep signal" regardless of gap length — only Transformers (content-based attention) or selective SSMs (input-dependent Δ, B, C) can.
4. Associative means that combining two consecutive steps, (A₁,B₁) then (A₂,B₂), into one equivalent combined step (A₂A₁, A₂B₁+B₂) gives the same result regardless of how the sequence of steps is grouped and combined. This lets the scan combine pairs in a balanced binary-tree pattern — halving the remaining elements each round — instead of a strict left-to-right loop, so the whole sequence resolves in O(log L) parallel depth on a GPU instead of L sequential dependent steps.
5. L² = (10⁵)² = 10¹⁰ pairwise scores per layer for the Transformer. Mamba's scan does O(L) = 10⁵ combine operations along the sequence per channel — five orders of magnitude fewer sequential terms, which is why Mamba's cost stays linear instead of quadratic as L grows.
6. False. The parallel scan gives good asymptotic complexity, but Mamba's practical throughput also depends on a hardware-aware kernel that keeps the expanded per-step states inside the GPU's fast on-chip SRAM rather than writing every intermediate value to slower HBM memory. Without that memory-access design, the same asymptotic algorithm would be bottlenecked by memory bandwidth and run much slower in practice — the same lesson FlashAttention taught for standard attention, whose O(L²) complexity didn't change but whose wall-clock speed did.