State Space Models: Mamba and Beyond

On 24 September 2014, Mangalyaan's onboard computer fired its liquid engine for exactly 24 minutes to slow the spacecraft into Martian orbit. The engine could not be controlled by a human on Earth — a radio command takes over 12 minutes each way to reach Mars, so by the time ISRO's ground station saw a problem, the burn would already be over. The spacecraft's attitude control system had to compute, thousands of times per second, how the vehicle's orientation would evolve under thruster firings, and correct it on the fly. The physics is continuous — torque acts on the spacecraft every instant, not in discrete jumps — but the onboard computer is a digital chip that can only take one step at a time. Somewhere between the continuous differential equation that describes the spacecraft's motion and the discrete loop that the flight computer actually executes, an exact translation has to happen. Get that translation wrong and the numerical spacecraft drifts from the real one.

That translation — continuous dynamics into an exact discrete recurrence — is the mathematical core of the state space model (SSM), the same object that now powers Mamba, a sequence architecture that processes language in linear time instead of the quadratic time attention requires. This chapter builds the SSM from the differential equation up: the continuous-time state equation, the discretization step that makes it computable, the HiPPO theory that tells you how to choose the system's internal dynamics so it remembers efficiently, and finally the selection mechanism that turns a fixed linear control system into something that can do content-based reasoning over a sentence. A companion chapter in this course compares Mamba's raw throughput against Transformers; this chapter is about why the machine works at all.

The continuous-time state equation

A linear time-invariant continuous system is described by two equations, borrowed directly from classical control theory (Kalman, 1960s) and used unchanged in modern deep sequence models:

dx(t)/dt = A x(t) + B u(t)
y(t)     = C x(t) + D u(t)

Here u(t) is the input signal at time t (a scalar or vector — in ISRO's case, thruster torque; in a language model, a token embedding treated as a continuous-valued channel), x(t) is a hidden internal state of dimension N (spacecraft angular velocity and orientation; in a language model, a learned latent memory vector), and y(t) is the observable output (attitude error; the model's prediction). A is an N×N matrix that governs how the state evolves on its own — its eigenvalues determine whether the system's memory decays, oscillates, or blows up. B maps the input into the state, and C reads the state back out into an output. D is usually a direct skip connection and is dropped in what follows.

This is exactly one linear ordinary differential equation per state dimension, coupled through A. It has an exact closed-form solution for any input, but a flight computer — or a GPU — cannot integrate a continuous ODE. It has to take a fixed step and update a discrete state, once per clock tick.

Discretization: zero-order hold

Suppose the input is held constant over each interval of length Δ (this is the "zero-order hold," ZOH, assumption — it is exact for a digitally sampled input, which is exactly what a token stream is). Then the ODE can be integrated exactly across one interval, giving a discrete recurrence:

Ā = exp(Δ A)
B̄ = (Δ A)⁻¹ (exp(Δ A) − I) B

x_k = Ā x_(k−1) + B̄ u_k
y_k = C x_k

Two things are worth pausing on. First, exp(ΔA) is a matrix exponential, not elementwise exponentiation — for a diagonal A it reduces to elementwise exp on the diagonal entries, which is why practical SSMs constrain A to be diagonal (or diagonalizable). Second, this is not an approximation of the continuous system — under the ZOH assumption it is the exact discrete-time solution. Let's verify that directly with a worked example, because the exactness is the whole point of choosing ZOH over a cruder scheme like Euler's method.

Worked example: discretizing a one-dimensional decay system

Take the scalar system A = −1, B = 1, C = 1, step size Δ = 0.5. This models exponential decay toward the input — a reasonable toy stand-in for, say, a satellite's angular velocity settling toward a commanded value. Discretize:

Ā = exp(Δ·A) = exp(−0.5) = 0.60653
B̄ = (Ā − 1)/A · B = (0.60653 − 1)/(−1) · 1 = 0.39347

Now drive the discrete recurrence with the input sequence u = [1, 1, 0, 0] (a unit step held for two ticks, then released), starting from x₀ = 0:

x₁ = 0.60653·0      + 0.39347·1 = 0.39347   → y₁ = 0.39347
x₂ = 0.60653·0.39347 + 0.39347·1 = 0.63212   → y₂ = 0.63212
x₃ = 0.60653·0.63212 + 0.39347·0 = 0.38340   → y₃ = 0.38340
x₄ = 0.60653·0.38340 + 0.39347·0 = 0.23254   → y₄ = 0.23254

Now check this against the continuous system directly. The ODE dx/dt = −x + 1 with x(0)=0 has the closed-form solution x(t) = 1 − e^(−t) while the input is held at 1. At t = 0.5 (one tick): 1 − e^(−0.5) = 0.39347. At t = 1.0 (two ticks): 1 − e^(−1.0) = 0.63212. Both match x₁ and x₂ exactly. Once the input is released at t=1.0, the system relaxes freely as x(t) = x(1)·e^(−(t−1)): at t=1.5, 0.63212·e^(−0.5) = 0.38340; at t=2.0, 0.63212·e^(−1.0) = 0.23254. Both again match x₃ and x₄ to five decimal places. The discrete recurrence is not tracking the continuous spacecraft approximately — for a piecewise-constant input it reproduces it exactly, tick for tick. This is why every modern SSM (S4, S5, Mamba) discretizes with ZOH or the closely related bilinear transform rather than a first-order Euler step, which would only approximate this and accumulate drift over long sequences.

HiPPO: choosing A so memory compression is optimal

The worked example fixed A = −1 by hand. In a state of dimension 1, that is the only real design choice — how fast the memory decays. But a real SSM state has hundreds of dimensions, and the question becomes: what should the N×N matrix A be, such that the state vector x(t) ∈ ℝᴺ retains as much information as possible about the entire history u(τ) for τ ≤ t, using only N numbers?

This is answered by HiPPO — "High-order Polynomial Projection Operators" (Gu, Dao, Ermon, Rudra & Ré, HiPPO: Recurrent Memory with Optimal Polynomial Projections, NeurIPS 2020). The idea: treat the state x(t) not as an arbitrary summary but as the coefficients of the best possible degree-(N−1) polynomial approximation of the input history u, measured against a weighting that says how much each past moment should matter right now. Choosing the weighting to give every point in [0, t] equal importance (the "LegS," scaled-Legendre, measure) and asking that this best-fit polynomial update optimally as t advances gives a unique, closed-form A — not a hyperparameter to search over, but a matrix derived from the calculus of the projection itself:

A[n,k] = −√(2n+1)·√(2k+1)   if n > k
       = −(n+1)              if n = k
       = 0                   if n < k

B[n]   = √(2n+1)

For a 3-dimensional state (N=3, indices 0,1,2), this evaluates to:

A = | −1.000   0       0     |     B = | 1.000 |
    | −1.732  −2.000   0     |         | 1.732 |
    | −2.236  −3.873  −3.000 |         | 2.236 |

Notice the structure: it is lower-triangular, with strictly negative diagonal entries (so every mode decays — the state is always forgetting, never exploding) and growing off-diagonal coupling. Each new state coordinate absorbs information from all the coordinates before it, weighted by how much a higher-order polynomial term needs to know about lower-order ones to stay consistent as the fitting window grows. The practical payoff, proven in the HiPPO paper, is a bound on reconstruction error that does not depend on how long the sequence has been running — a HiPPO-initialized state of fixed size N can approximate arbitrarily long histories with a guaranteed, uniform error bound, in contrast to a randomly initialized recurrent matrix, whose memory decays or blows up exponentially with sequence length depending on its eigenvalues. S4 (Gu, Goel & Ré, Efficiently Modeling Long Sequences with Structured State Spaces, ICLR 2022) is essentially "take this HiPPO matrix as a fixed initialization for A, keep the SSM linear time-invariant (LTI), and turn the whole discrete recurrence into a global convolution so it can be computed with an FFT in O(L log L) for sequence length L." That convolutional trick is what gives S4 its long-range-memory performance and its efficient training — and, as the next section shows, its principal limitation.

Why linear time-invariance is not enough: the selection problem

S4's speed trick depends on A, B, C, and Δ all being fixed, the same at every timestep, for every input. That is exactly what "linear time-invariant" means, and it is exactly what makes the discrete recurrence collapse into one fixed convolution kernel that can be precomputed and applied via FFT. But it also means the model cannot decide, based on what a token actually says, to remember it strongly or forget it immediately — the dynamics are baked in before the input is even seen. Consider a "selective copy" task: given a sequence with a mix of content tokens and filler tokens, output only the content tokens in order, ignoring the filler. This requires the model to look at each token and decide, content-dependently, whether it matters. An LTI system, applying the same Ā and B̄ regardless of content, structurally cannot do this — every token is integrated into the state with the same fixed weighting. Mamba (Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, arXiv:2312.00752, 2023) fixes this by making the discretization itself a function of the input.

Mamba's S6: input-dependent Δ, B, C

In Mamba's "S6" layer (Selective SSM), the state matrix A remains a fixed, structured, per-channel parameter — initialized in the diagonal HiPPO style so each channel still starts with the "remember uniformly across the past" property derived above. But B, C, and critically the step size Δ stop being constants and become linear projections of the current input token:

Δ_t = softplus(W_Δ · x_t)
B_t = W_B · x_t
C_t = W_C · x_t
Ā_t = exp(Δ_t · A)

Because Δ_t now depends on the token, so does Ā_t = exp(Δ_t A) — even though A itself never changes, the effective discretized decay changes on every step. This is the mechanism's real trick: you don't need to make the big matrix A input-dependent (which would be expensive and could destabilize the fixed HiPPO structure) — modulating the scalar-per-channel Δ that A gets multiplied by before exponentiating is enough to make the discretized system content-aware.

Worked example: selection in action

Reuse the scalar decay system A = −1 from the earlier ZOH example, and give it a learned selection weight W_Δ = 0.5 (bias 0), so Δ_t = softplus(0.5·x_t). Feed it two different input tokens, encoded as scalars x_t = 2 (a salient, information-bearing token) and x_t = −1 (a filler token):

x_t = 2:   Δ_t = softplus(0.5·2)  = softplus(1.0)  = ln(1+e¹)   = 1.3133
           Ā_t = exp(−1·1.3133)  = 0.2689

x_t = −1:  Δ_t = softplus(0.5·−1) = softplus(−0.5) = ln(1+e^−0.5) = 0.4741
           Ā_t = exp(−1·0.4741)  = 0.6225

The salient token produces a large Δ_t, which drives Ā_t down to 0.27 — the previous state is heavily discounted, and the new input dominates the update (a large step size means the continuous system is allowed to evolve a long "virtual time" in one tick, moving it close to whatever the input is currently pushing it toward). The filler token produces a small Δ_t, keeping Ā_t near 0.62 — the state barely moves, mostly preserving whatever it was already carrying. This is precisely a content-based gate, but derived from a control-theoretic discretization step rather than hand-designed the way an LSTM's forget gate is. The same equation that ISRO's flight computer uses to advance a spacecraft's state by one clock tick, here advances a language model's memory by one token — except now the "clock" itself speeds up or slows down depending on what it reads.

One further consequence: once Ā and B̄ vary per timestep, the discrete recurrence is no longer a single fixed convolution kernel, so S4's FFT trick no longer applies directly. Mamba instead computes the recurrence with a hardware-aware parallel scan and fused GPU kernels — a separate engineering story from the mathematics here, and the one the companion chapter's throughput comparison is built on.

Beyond Mamba: state-space duality and hybrid architectures

Mamba's S6 layer is not the end of the line. Mamba-2 (Dao & Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, ICML 2024) proves that a restricted form of selective SSM — one where A is a scalar times the identity per attention-style "head," rather than the fully general diagonal matrix S6 allows — is mathematically identical to a particular kind of structured, masked attention: the recurrence can be unrolled into a "semiseparable" attention matrix, and the identical computation can be expressed either as the sequential scan this chapter builds, or as a sequence of block matrix multiplications, the SSD (structured state space duality) algorithm. That duality lets Mamba-2 borrow the matmul-heavy, tensor-core-friendly training path that makes Transformers fast to train, while keeping the linear-time, constant-memory recurrent path at inference — and it permits a much larger state dimension N than the original Mamba, which directly sharpens the HiPPO approximation this chapter derived.

A second, more pragmatic direction accepts that pure SSMs and pure attention each have a structural weakness — an SSM compresses the entire past into a fixed-size state and can lose exact recall of one early token buried in a long context, while attention keeps every token exactly but pays the quadratic cost this chapter's companion piece measures — and interleaves both layer types in a single network instead of choosing one. Jamba (AI21 Labs, 2024) is the production example: most of its layers are Mamba, but a small fraction are ordinary Transformer attention layers dropped into the stack, trading away a little of the linear-time efficiency derived above in exchange for attention's exact long-range lookup on the tokens that need it. The general pattern — a fixed matrix A with HiPPO-style structure doing the cheap bulk compression, attention doing the expensive precise retrieval where it is worth the cost — is where most state-space research has moved since Mamba's original release.

Common misconception

Students who have just learned about Transformer context windows often assume the state dimension N in an SSM plays the same role — "a bigger N means the model can look back further, and once you run out of state slots, old tokens fall off, just like a context window." This is wrong in an important way. A Transformer's context window is a hard architectural cutoff: token L+1 in a window of size L is structurally invisible, no matter what it contains. An SSM's state dimension N is not a buffer of past tokens at all — it is the dimensionality of a compressed function approximation of the entire history, and HiPPO's guarantee is precisely that this approximation's error bound does not grow with how long the sequence has run. A HiPPO-initialized SSM with N=64 is not "a 64-token window" — it is a fixed-size summary that keeps approximating an arbitrarily long history, with old information smoothly and gracefully compressed rather than sharply discarded. In practice, larger N gives a finer-grained (lower-error) approximation of the whole past, not a longer hard cutoff; and Mamba's selectivity adds a second, orthogonal lever — the model can choose, per token, how much new information to write into that fixed-size summary at all.

How the pieces fit together

Active recall

Attempt these before reading the answers.

1. A scalar SSM has A = −2, B = 1, step size Δ = 0.25. Compute Ā and B̄ using the ZOH formulas.
2. In the worked ZOH example (A=−1, Δ=0.5), suppose instead the step size were Δ = 1.0 but the input sequence is resampled to match (one input value per full unit of time: u = [1, 0] instead of [1,1,0,0]). Compute x₁ and x₂ and check them against the continuous analytic solution at t=1 and t=2.
3. Why is the HiPPO-LegS matrix A lower-triangular with strictly negative diagonal entries — what would go wrong if a diagonal entry were positive?
4. In Mamba's S6 layer, if a token produces a very small Δ_t (close to 0), what happens to Ā_t = exp(Δ_t A) and to the resulting state update? Describe the behavior in words.
5. Explain in one or two sentences why S4's FFT-based convolution trick breaks once B̄ and C become input-dependent, using the definition of the discrete recurrence.
6. A classmate says: "Since Mamba makes Δ, B, and C functions of the input, doesn't that mean A is effectively input-dependent too, since it's the thing being discretized?" Is this reasoning correct? Explain precisely what does and does not vary with the input.

Answers

1. Ā = exp(Δ·A) = exp(0.25 · −2) = exp(−0.5) = 0.60653. B̄ = (Ā−1)/A · B = (0.60653−1)/(−2) · 1 = (−0.39347)/(−2) = 0.19674.
2. Discretize with Δ=1.0: Ā = exp(−1·1.0) = 0.36788, B̄ = (Ā−1)/A = (0.36788−1)/(−1) = 0.63212. Recurrence: x₁ = 0.36788·0 + 0.63212·1 = 0.63212; x₂ = 0.36788·0.63212 + 0.63212·0 = 0.23254. Analytic check: x(t)=1−e^(−t) at t=1 gives 1−e^(−1)=0.63212 ✓, matching x₁; then with input released, x(t)=x(1)·e^(−(t−1)), at t=2: 0.63212·e^(−1)=0.23254 ✓, matching x₂. Both match exactly — note these are the same numerical values as x₂ and x₄ in the original Δ=0.5 example, because both discretizations sample the identical underlying continuous trajectory at t=1 and t=2; ZOH discretization is exact regardless of the step size chosen, as long as the input is genuinely constant over each interval.
3. Strictly negative diagonal entries mean every state coordinate decays on its own in the absence of input — this bounds the state's growth and is what makes the HiPPO error bound uniform over arbitrarily long sequences. A positive diagonal entry would make that coordinate an unstable, exponentially growing mode: with no forcing, dx_n/dt = A[n,n]·x_n would blow up as t→∞, so the "memory" would diverge instead of settling into a bounded polynomial-approximation coefficient — the opposite of the compression HiPPO is designed to guarantee.
4. A small Δ_t pushes Ā_t = exp(Δ_t A) close to exp(0) = 1 (since A is finite and negative, a tiny Δ_t makes Δ_t A close to 0). With Ā_t ≈ 1, the recurrence h_t = Ā_t h_(t-1) + B̄_t x_t keeps almost all of the previous state and barely lets the current token's contribution register — the model effectively skips or ignores that token, which is exactly the behavior needed to filter out filler content in a selective-copy-style task.
5. S4's convolution trick works because with LTI dynamics the same Ā, B̄, C are reused at every step, so the mapping from the whole input sequence to the whole output sequence is one fixed linear operator — a convolution with a kernel built once from the powers of Ā, computable via FFT. Once B̄ and C differ at every timestep (as they do under selection), the operator connecting input to output is a different, input-dependent matrix at every position, so there is no single fixed kernel left to convolve with — the computation has to fall back to an explicit (though parallelizable) scan through the recurrence.
6. Not quite correct, and the imprecision matters. The parameter A itself is not a function of the input — it is a fixed, learned, HiPPO-structured matrix, the same for every token. What is input-dependent is the scalar step size Δ_t that A gets multiplied by before exponentiating: Ā_t = exp(Δ_t A). Because Δ_t varies with the token, the effective discretized decay Ā_t does vary with the token — but this is achieved without ever changing A itself, which is what lets Mamba keep A's stable, HiPPO-derived structure intact while still getting content-dependent dynamics through Δ.

Think About It

Think about this: How would you explain state space models: mamba and beyond to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.