State-Space Models and Mamba

A recurrent network runs in $O(n)$ time but trains through a product of Jacobians that vanishes or explodes, so long-range signal is hard to learn. The Transformer learns any-distance dependencies directly through attention, but pays $O(n^2)$ in time and memory. A state-space model (SSM) targets the corner neither reaches: $O(n)$ cost and stable long-range memory. It does so by making the recurrence linear in the hidden state, a restriction that turns out to expose two equivalent computations from one set of weights.¹

The continuous linear state-space

The starting point is not a neural network but a linear dynamical system from control theory: a hidden state $x(t)\in {\mathbb{R}}^N$ driven by a scalar input signal $u(t)$, evolving by a linear ordinary differential equation and read out linearly.

The $Du(t)$ term is a plain residual, so we drop it from the derivations and restore it at the end. What matters is how the state $x(t)$ integrates the past of $u$ and how $C$ reads it out.

Solving the linear ODE gives an explicit integral: the state at time $t$ is a convolution of the input's past against the matrix exponential of $A$.

The output is a convolution against the kernel $h(s)=C{e}^{As}B$. Every efficient algorithm in this lesson exploits this fact.

Discretization

Sequences are sampled, not continuous, so we discretize the ODE at a step size $\Delta$. The standard choice is the zero-order hold (ZOH): assume the input $u$ is constant across each interval $[k\Delta ,(k+1)\Delta )$ and integrate the dynamics exactly over that interval. This converts $(A,B)$ into discrete matrices $\overline{A},\overline{B}$.

The derivation is the continuous solution applied over one step. Integrate from $(k-1)\Delta$ to $k\Delta$ with $u\equiv u_k$ pulled out of the integral:

which is exactly $\overline{A}=e^{\Delta A}$ and $\overline{B}=A^{-1}(e^{\Delta A}-I)B$, matching the ZOH formula after grouping the $\Delta$. The step size $\Delta$ is a learned parameter: it sets the timescale at which the continuous system is sampled.

For example, take a scalar state ($N=1$) with $A=-1$, $B=1$, and step size $\Delta =0.5$. Then

The recurrence $x_k=0.607x_{k-1}+0.393u_k$ is a leaky integrator: each step keeps $60.7\%$ of the old state and admits $39.3\%$ of the new input. Halving the step to $\Delta =0.25$ gives $\overline{A}=e^{-0.25}\approx 0.779$, retaining more per step so the state persists longer, exactly the long-memory regime the remark describes. Notice $\overline{A}+(-\overline{B}/A)=e^{\Delta A}+(1-e^{\Delta A})=1$: the ZOH split conserves mass, which is why the discrete system inherits the continuous system's stability.

In an $N$-dimensional state the same formulas apply with matrices. $\overline{A}=e^{\Delta A}$ is the $N\times N$ matrix exponential, and $\overline{B}$ is $N\times 1$; when $A=\Lambda$ is diagonal (the S4D case below) the exponential is elementwise, ${\overline{A}}_{ii}=e^{\Delta {\lambda }_i}$, so discretization costs $O(N)$ rather than the $O(N^3)$ of a dense matrix exponential.

With the discrete matrices fixed, the SSM is now an ordinary linear recurrence, identical in shape to an RNN cell but with no nonlinearity inside the state update. That linearity is what we exploit next.

Two views of one model

Unrolling the recurrence from $x_0=0$ writes each output as a weighted sum of past inputs. Expand $x_k$ step by step:

so reading out $y_k=C{x}_k$ gives a convolution of the input against a fixed kernel whose taps are $C{\overline{A}}^j\overline{B}$.

Continuing the scalar example ($\overline{A}=0.607$, $\overline{B}=0.393$) with $C=1$, the kernel taps are ${\overline{K}}_j=C{\overline{A}}^j\overline{B}=0.393\times 0.607^j$:

a decaying exponential filter. Both algorithms must give the same $y$. Feed the impulse $u=(1,0,0,\dots )$: the recurrence gives $x_1=0.393$, then $x_2=0.607\times 0.393=0.239$, then $x_3=0.145$, so $y_k=C{x}_k$ reproduces $\overline{K}$ tap for tap. The kernel is literally the SSM's impulse response, and convolving any input against it is the parallel restatement of running the recurrence, taps computed once and reused for every position.

The same weights describe two algorithms, and either can be used as the task requires. Training has the whole sequence in hand, so it materializes the kernel $\overline{K}$ once and convolves in parallel; autoregressive inference has tokens arriving one at a time, so it steps the recurrence and keeps only the $O(N)$ state.

The convolution is global: $\overline{K}$ has one tap per position, so the kernel is as long as the sequence. Computing it naively costs a matrix power ${\overline{A}}^j$ per tap, and convolving a length-$n$ kernel by FFT costs $O(n\log n)$. The open problem the early SSM work solved was making that kernel both cheap to compute and good at remembering, which is entirely a question about $A$.

Dimensions of a deep SSM layer

The definition above is a single-input-single-output (SISO) system: it maps one scalar channel $u(t)$ to one scalar output. A real layer processes a batch of sequences with many channels, so the SISO system is replicated. For a batch of $B$ sequences, each of length $L$, with $H$ feature channels and state size $N$ per channel, the tensors are:

The $H$ channels are independent SSMs: channel $h$ has its own $(A_h,B_h,C_h,{\Delta }_h)$ and never sees the others inside the SSM. Mixing across channels happens outside the recurrence, in the linear projections that sandwich the layer, exactly as attention mixes across positions and a following MLP mixes across channels. The state that inference must carry is $(B,H,N)$ numbers, independent of $L$, which is the constant-memory claim made concrete.

Long-range memory: HiPPO

A random $A$ does not remember. Its kernel taps $C{\overline{A}}^j\overline{B}$ decay like the eigenvalues of $\overline{A}$ raised to the $j$, so the same geometric forgetting that plagues the vanilla RNN reappears unless $A$ is chosen deliberately. The fix is HiPPO (high-order polynomial projection operators): pick $A$ so that the hidden state holds the optimal coefficients of a polynomial that approximates the entire input history.

For the (scaled Legendre) HiPPO measure the state matrix has an explicit lower-triangular form, with entries growing with the row index $n$:

which is what lets the high-order coefficients capture fine recent detail while the low-order ones retain the distant past. Initializing $A$ this way is the difference between a model that forgets in a dozen steps and one that carries signal across thousands.

The remaining obstacle is cost. HiPPO's $A$ is dense, so each kernel tap needs a dense matrix power, and materializing $\overline{K}$ becomes $O(n{N}^2)$, far too expensive. S4 (structured state spaces) makes the kernel cheap by writing $A$ as a diagonal-plus-low-rank (DPLR) matrix, $A=\Lambda -P{P}^{\ast }$, whose powers can be summed with a fast Cauchy-kernel computation rather than naive matrix powers. Two simplifications followed almost immediately: the diagonal SSM (DSS / S4D) drops the low-rank term and keeps only a diagonal $\Lambda$, which is far simpler and nearly as accurate, and S5 uses a single multi-input-multi-output system computed with a parallel associative scan instead of a convolution.

Selective state spaces: Mamba

Every model so far shares one limitation: $\overline{A},\overline{B},C,\Delta$ are fixed once trained, so the convolution kernel is the same for every input. That is what makes the convolution possible, and it is also a ceiling. A linear time-invariant system cannot do content-based reasoning: it cannot decide to ignore a filler word or to specifically remember a name, because its dynamics do not depend on what the input is.

This is the selection mechanism of Mamba. Input-dependent ${\Delta }_k$ acts as a learned gate: a large ${\Delta }_k$ resets the state toward the current token (remember this), a small ${\Delta }_k$ holds the existing state (skip this). Input-dependent $B_k,C_k$ decide what gets written into and read out of memory.

This sacrifices the convolution. Once the kernel depends on the input, it is different at every step, so there is no single $\overline{K}$ to FFT. Mamba is back to a sequential recurrence, which on a GPU would be slow. The fix is algorithmic: the linear recurrence is an associative scan, and associative scans parallelize.

Mamba's contribution past the scan itself is to make it hardware-aware: the state is expanded only inside fast GPU SRAM, the scan runs there, and only the outputs are written back to slow HBM, so the large intermediate state never touches main memory.

A full Mamba block wraps the selective scan in a gated structure reminiscent of a gated MLP: the input is projected up, passed through a short causal convolution and the selective SSM on one branch, multiplied by a SiLU-gated branch, and projected back down. Mamba-2 later simplified this by showing that selective SSMs and attention are two views of one operation (a structured masked matrix product), a duality called state-space duality (SSD) that lets the scan be cast as a matrix multiplication and run even faster on tensor-core hardware.

Concretely, a residual input of shape $(B,L,D)$ enters. An input projection widens it to $(B,L,E)$ with an expansion factor (typically $E=2D$), splitting into two branches of width $E$. The main branch runs a short depthwise causal convolution (kernel width $\approx 4$, mixing a few neighboring positions) and a SiLU nonlinearity, then the selective SSM with $H=E$ channels and state size $N$ (often $N=16$). The gate branch is a SiLU applied to its half, multiplied into the main branch elementwise. A final projection contracts $(B,L,E)$ back to $(B,L,D)$ for the residual add.

Cost and quality versus the Transformer

The cost comparison is the main argument for SSMs. Self-attention spends $O(n^{2}d)$ to form an $n\times n$ score matrix and caches a growing key-value buffer at inference. An SSM never builds an $n\times n$ object: training is $O(n)$ work (linear, with an $O(\log n)$ scan depth), and inference carries a fixed-size state independent of how long the context already is.

Quality is measured on the Long Range Arena (LRA), a benchmark of six tasks built to demand dependencies over thousands of tokens, where Transformers struggle both in accuracy and in the memory cost of long sequences. S4 was the first model to clear all six tasks by a wide margin, including the Path-X task that every prior Transformer variant had failed outright, and selective SSMs carried the gains to language modeling.

The trade-off: attention keeps a perfect, addressable record of every past token and pays for it with quadratic compute and a context-length-bounded cache; an SSM compresses the past into a fixed state and pays for that with a lossy summary. The selection mechanism makes the compression content-aware, so the loss falls on the tokens weighted as unimportant. For very long sequences, linear time and constant state are decisive.

The deep SSM lineage

Deep state-space models are entirely post-2016, so every result here rests on the recent literature rather than the standard references.

• The primary line. HiPPO (Gu et al., NeurIPS 2020) derived the transition matrix that lets a fixed-size state hold an optimal polynomial memory of the whole past. S4 (Gu, Goel & Ré, Efficiently Modeling Long Sequences with Structured State Spaces, ICLR 2022) made that kernel cheap with the diagonal-plus-low-rank form and cleared the Long Range Arena. Mamba (Gu & Dao, Linear-Time Sequence Modeling with Selective State Spaces, 2023) added input-dependent parameters and the hardware-aware scan, matching Transformers on language at a fraction of the generation cost. The simplifications — diagonal-only S4D (Gu et al., 2022) and S5 (Smith et al., ICLR 2023) — are what made the models easy to implement.
• One family, many names. The linear recurrence at the core of an SSM is the same object as linear attention (attention with the softmax removed, which makes it associative and therefore expressible as a recurrence) and as the RWKV and RetNet architectures. Reading all of these as attention without softmax, run as a scan unifies a scattered literature: the price of dropping the softmax is losing the perfect content-addressable lookup, and the selection mechanism in Mamba is one way to buy back the content-awareness the softmax provided.
• Hybrids in practice. The clean trade — attention's exact recall versus the SSM's linear cost — pushed production models toward interleaving the two: a few full-attention layers for precise copying and lookup, many SSM layers for cheap long-context mixing (Jamba, Lieber et al., 2024). The lesson mirrors ViT and Graph Transformers: the strongest architectures rarely pick one primitive outright.

Takeaways

• A state-space model carries a continuous linear state, $x^{\prime }=Ax+Bu$, $y=Cx+Du$, and its solution is a convolution of the input against the kernel $C{e}^{As}B$ — every algorithm below follows from this linearity.
• Zero-order-hold discretization turns the ODE into a linear recurrence $x_k=\overline{A}{x}_{k-1}+\overline{B}{u}_k$ with $\overline{A}=e^{\Delta A}$; the learned step size $\Delta$ sets the memory timescale.
• One set of weights yields two equivalent algorithms: a sequential recurrence ($O(1)$ state, ideal for generation) and a global convolution $y=\overline{K}\ast u$ with ${\overline{K}}_j=C{\overline{A}}^j\overline{B}$ ($O(n\log n)$, ideal for training).
• Long-range memory comes from HiPPO, which initializes $A$ so the state holds the optimal polynomial approximation of all the past; a random $A$ forgets geometrically. S4 makes that kernel cheap with a diagonal-plus-low-rank $A$ (simplified to diagonal in S4D/DSS and S5).
• Mamba makes $\overline{B},C,\Delta$ input-dependent (selectivity), gaining content-based gating at the cost of the static convolution, and recovers speed with a hardware-aware parallel scan ($O(\log n)$ depth, on-chip).
• Against the Transformer, an SSM trains in linear time and generates with a constant-size state rather than a growing KV-cache; on the Long Range Arena S4 and its selective successors win the long-context tasks attention cannot reach.