Mamba and Selective State Spaces

An S4 layer's transition matrices are fixed after training — every input sequence, whatever it contains, gets pushed through the exact same linear dynamics. That is precisely what makes S4 fast (the whole sequence can be pre-computed as one convolution), but it is also a real limitation: the model cannot look at an input and decide, on the fly, "this particular token matters a lot, keep it in memory" versus "this one is noise, let it decay quickly." An LSTM's gates can do exactly that — decide, input by input, what to keep — but at the cost of a sequential, step-by-step recurrence that can't be pre-computed as a parallel convolution. Mamba's contribution is making the state-space transition matrices themselves depend on the current input, recovering LSTM-style content-aware selectivity, while finding a different computational trick (not the convolution trick, since input-dependent dynamics break that) to keep training fast anyway.

The analogy: a highlighter that knows what matters

An S4 layer is like reading a document with a fixed rule: highlight every fifth word, no matter what the document says. It's fast and mechanical, but it can't tell a critical number from a filler word — both get the same treatment. An LSTM is like a careful human reader who highlights based on judgment, deciding word by word what's important, but who has to read strictly left to right, one word at a time, which is slow over a long document. Mamba is a reader who highlights based on judgment — content-aware, like the human — but who has found a way to do that highlighting so efficiently that reading a whole book still runs at nearly the speed of the mechanical fixed-rule approach. The selectivity of an LSTM's gates, at close to the speed of S4's parallel scan.

The maths

Recall S4's discretized recurrence, where $\bar A$, $\bar B$ are fixed matrices shared across every timestep:

$$x_k = \bar{A}\, x_{k-1} + \bar{B}\, u_k, \qquad y_k = C x_k$$

Mamba's core change is to make the parameters that control this recurrence functions of the current input $u_k$, rather than constants:

$$B_k = f_B(u_k), \qquad C_k = f_C(u_k), \qquad \Delta_k = f_\Delta(u_k)$$

In words: instead of one fixed $B$ and $C$ used at every step, small learned linear layers $f_B$, $f_C$, $f_\Delta$ look at the current input $u_k$ and produce input-specific versions of these parameters at every timestep. $\Delta_k$ (the discretization step size) is the most important of the three — it directly controls how much the current input is allowed to affect the state versus how much of the old state is preserved.

The recurrence becomes:

$$x_k = \bar{A}(\Delta_k)\, x_{k-1} + \bar{B}_k\, u_k, \qquad y_k = C_k\, x_k$$

In words: at each step, the effective transition and input matrices are recomputed based on what the input actually is, right before being applied. When $\Delta_k$ is small for a given input, $\bar A(\Delta_k)$ stays close to the identity — the state barely updates, meaning this input is effectively skipped, its effect on memory minimal. When $\Delta_k$ is large, the state updates substantially — this input is treated as important and strongly written into memory. This is functionally the same job an LSTM's forget and input gates do, but arrived at through a different mechanism: a continuous discretization step size rather than a discrete sigmoid gate.

The catch: because $\bar A$ now changes at every timestep (it depends on $\Delta_k$, which depends on $u_k$), the neat trick that made S4 fast — collapsing the whole sequence into one fixed convolution kernel computed once — no longer works, since there is no single fixed kernel anymore. Mamba instead uses a hardware-aware parallel scan: a way of computing the sequential recurrence that still parallelizes across the sequence length (it's mathematically an associative operation, so partial results can be combined out of order and merged), implemented to minimize slow memory transfers on a GPU. The output is the same sequential-recurrence answer, computed with much better hardware utilization than a naive step-by-step loop.

Worked example 1: input-dependent Δ changes the outcome

Take a 1-D toy state (as in the S4 page), with $\bar A(\Delta) = e^{-\Delta}$ (a standard discretization form) and $\bar B = 1$, $C = 1$, starting from $x_0 = 0$. Feed two inputs, $u_1 = 5$ (an unusual, informative spike) and $u_2 = 5$ (an identical value, but suppose the model has learned it's routine noise in this context, so it produces a much smaller $\Delta$ for it).

Suppose $\Delta_1 = 1.0$ (large, "this matters") and $\Delta_2 = 0.05$ (small, "skip this"):

$$\bar A(\Delta_1) = e^{-1.0} \approx 0.368, \qquad \bar A(\Delta_2) = e^{-0.05} \approx 0.951$$ $$x_1 = 0.368(0) + 1(5) = 5.0$$ $$x_2 = 0.951(5.0) + 1(5) \approx 4.755 + 5 = 9.755 \quad \text{(if treated as also important)}$$

But with $\Delta_2$ small, the input contribution is also scaled down proportionally in a real Mamba implementation (via $\bar B(\Delta_2)$, not just $\bar A$) — the point being that the same raw value $u_2=5$ produces a very different state update depending on what $\Delta_2$ the model assigns it. An S4 layer, by contrast, would apply the exact same $\bar A$, $\bar B$ to both inputs regardless of their context, because those matrices are fixed constants, not functions of $u_k$.

Worked example 2: a selective copy task

Consider a synthetic test that's standard in this literature: a sequence with a handful of meaningful tokens scattered among many "filler" tokens, where the task is to recall the meaningful tokens later in the sequence, in order, ignoring the filler entirely — say, positions $3, 47, 112$ hold real values and everything else is a repeated filler symbol.

A fixed-dynamics model (plain S4, or a plain RNN) processes filler and meaningful tokens identically — every position updates the state by the same rule, so the filler dilutes the state's capacity right alongside the meaningful signal, since nothing in the recurrence itself knows to treat position 3's content differently from position 4's. A selective model can learn to assign near-zero $\Delta$ to filler tokens (barely updating state on them) and a large $\Delta$ to the three meaningful ones (writing them cleanly into state), so the state ends up holding close to just the three values that matter, undiluted by the hundred-plus filler positions in between. This filtering behavior — telling signal from filler based on content — is exactly the class of task where selective state spaces measurably outperform plain S4, and it's a direct, testable consequence of making $\Delta$ input-dependent rather than fixed.

Selectivity is a routing decision made per input, similar in spirit to how a classifier's decision boundary treats different regions of input space differently rather than applying one uniform rule everywhere — Mamba's $\Delta_k$ is effectively learning where in input space to "turn the state update on" versus "leave it alone."

What this means in practice

Selective state-space models are attractive wherever S4-scale sequence lengths matter (very long tick histories, full order-book sessions) and the sequence genuinely contains a mix of important and irrelevant content that a fixed-dynamics model can't tell apart — for instance, a session dominated by quiet, low-information ticks punctuated by a handful of information-bearing prints. Where a sequence is fairly uniform in importance throughout, plain S4's simplicity and pure-convolution training speed may still be the more practical choice, since Mamba's added selectivity has more moving parts to tune and debug.