---
title: "State-Space Models and Mamba"
module: Architectures
moduleNumber: 5
lessonNumber: 9
order: 509
summary: >
  A state-space model carries a continuous linear hidden state through a sequence,
  and that linearity buys two equivalent algorithms from one set of weights: a
  recurrence that runs in linear time with constant memory, and a global convolution
  that trains in parallel. Long-range memory comes from how the transition matrix is
  initialized (HiPPO) and parameterized (S4's diagonal-plus-low-rank form). Mamba
  breaks the convolution on purpose, making the parameters input-dependent so the model
  can select what to remember, recovered at speed by a hardware-aware parallel scan.
topics: [Architectures]
sources:
  - book: Goodfellow
    ref: "Ch. 10 — sequence modeling (deep SSMs postdate the 2016 text)"
---

A [recurrent network](/deep-learning/architectures/recurrent-networks) 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](/deep-learning/architectures/the-transformer-architecture)
learns any-distance dependencies directly through [attention](/deep-learning/architectures/attention-and-transformers),
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.[^gf-seq]

| Model | Train cost | Inference cost / step | Long-range memory |
| --- | --- | --- | --- |
| RNN / LSTM | $O(n)$ sequential | $O(1)$ state | hard (vanishing gradient) |
| Transformer | $O(n^2)$ parallel | $O(n)$ KV-cache | direct (attention) |
| State-space model | $O(n)$ or $O(n \log n)$ parallel | $O(1)$ state | by construction (HiPPO) |

## 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.

> **Definition (Continuous state-space model).** For matrices
> $A \in \mathbb{R}^{N \times N}$, $B \in \mathbb{R}^{N \times 1}$,
> $C \in \mathbb{R}^{1 \times N}$, $D \in \mathbb{R}$, the SSM maps a $1$-D input
> $u(t)$ to a $1$-D output $y(t)$ through an $N$-dimensional latent state:
> $$
> x'(t) = A\,x(t) + B\,u(t),
> \qquad
> y(t) = C\,x(t) + D\,u(t).
> $$
> $A$ is the **state matrix** governing the latent dynamics, $B$ the input map,
> $C$ the output map, and $D$ a skip connection. A deep SSM layer runs $H$ such
> systems in parallel, one per feature channel.

The $D\,u(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.

$$
% caption: The continuous SSM as a block diagram. The input drives an integrator through $B$;
% state feeds back through $A$ and is read out through $C$, with $D$ as a direct skip.
\begin{tikzpicture}[>=stealth, font=\small,
  blk/.style={draw, black, thick, minimum width=10mm, minimum height=8mm, align=center},
  sum/.style={draw, black, thick, circle, inner sep=0pt, minimum size=6mm}]
  \definecolor{acc}{HTML}{2348F2}
  \definecolor{green}{HTML}{1F9D4D}
  \node (u)    at (0,0)    {$u(t)$};
  \node[blk, draw=acc, text=acc] (B) at (1.8,0) {$B$};
  \node[sum] (s)  at (3.6,0) {$+$};
  \node[blk, draw=acc, text=acc, minimum width=14mm] (int) at (5.4,0) {integrate};
  \node[blk, draw=acc, text=acc] (C) at (7.8,0) {$C$};
  \node[sum] (s2) at (9.6,0) {$+$};
  \node (y)    at (11.2,0) {$y(t)$};
  \draw[->, thick] (u) -- (B);
  \draw[->, thick] (B) -- (s);
  \draw[->, thick] (s) -- (int) node[midway, above, font=\footnotesize] {rate};
  \draw[->, thick] (int) -- (C) node[midway, above, font=\footnotesize] {state};
  \draw[->, thick] (C) -- (s2);
  \draw[->, thick] (s2) -- (y);
  % feedback through A
  \node[blk, draw=acc, text=acc] (A) at (5.4,-2.0) {$A$};
  \draw[->, acc, thick] (int.south) -- (5.4,-1.2) -- (A.north);
  \draw[->, acc, thick] (A.west) -- (2.7,-2.0) -- (2.7,-0.5) -- (s.south);
  \node[acc, font=\footnotesize, anchor=west] at (6.0,-2.0) {state feedback};
  % skip D
  \node[blk, draw=green, text=green] (D) at (3.6,2.0) {$D$};
  \draw[->, green, thick] (0.4,0.25) -- (0.4,2.0) -- (D.west);
  \draw[->, green, thick] (D.east) -- (9.6,2.0) -- (s2.north);
  \node[green, font=\footnotesize, anchor=south] at (6.5,2.0) {direct skip};
\end{tikzpicture}
$$

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$.

> **Theorem (Continuous solution).** With $x(0) = 0$, the system
> $x' = Ax + Bu$, $y = Cx$ has the closed-form response
> $$
> x(t) = \int_{0}^{t} e^{A(t - \tau)} B\,u(\tau)\,\d\tau,
> \qquad
> y(t) = \int_{0}^{t} C\,e^{A(t-\tau)} B\,u(\tau)\,\d\tau .
> $$

> **Proof.** Write $x' - Ax = Bu$ and multiply by the integrating factor
> $e^{-At}$, so $\frac{\d}{\d t}\parens{e^{-At} x} = e^{-At} B u$. Integrate from
> $0$ to $t$ with $x(0) = 0$ to get $e^{-At} x(t) = \int_0^t e^{-A\tau} B u(\tau)\,\d\tau$,
> then left-multiply by $e^{At}$. Substituting into $y = Cx$ gives the second
> identity, a convolution with kernel $C e^{A s} B$. $\qed$

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}$.

> **Definition (Zero-order-hold discretization).** Holding $u$ constant on each
> step of width $\Delta$ and integrating $x' = Ax + Bu$ exactly yields the discrete
> recurrence
> $$
> x_k = \overline{A}\,x_{k-1} + \overline{B}\,u_k,
> \qquad
> y_k = C\,x_k,
> $$
> with
> $$
> \overline{A} = e^{\Delta A},
> \qquad
> \overline{B} = \parens{\Delta A}^{-1}\parens{e^{\Delta A} - I}\,\Delta 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:

$$
x_k = e^{\Delta A} x_{k-1} + \parens{\int_{0}^{\Delta} e^{A s}\,\d s} B\,u_k
    = e^{\Delta A} x_{k-1} + A^{-1}\parens{e^{\Delta A} - I} B\,u_k,
$$

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.

> **Remark (Why $\Delta$ matters).** Small $\Delta$ makes $\overline{A} \approx I + \Delta A$,
> a near-identity recurrence that holds state for many steps (long memory); large
> $\Delta$ shrinks the eigenvalues of $\overline{A}$ toward $0$, forgetting quickly.
> $\Delta$ is the model's control over how fast the past decays.

For example, take a scalar state ($N = 1$)
with $A = -1$, $B = 1$, and step size $\Delta = 0.5$. Then
$$
\overline{A} = e^{\Delta A} = e^{-0.5} \approx 0.607,
\qquad
\overline{B} = A^{-1}\parens{e^{\Delta A} - I}B = \frac{e^{-0.5} - 1}{-1} \approx 0.393 .
$$
The recurrence $x_k = 0.607\,x_{k-1} + 0.393\,u_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.

```algorithm
caption: $\textsc{SsmRecurrent}(u_{1:n}, \overline{A}, \overline{B}, C)$ — linear-time, constant-state inference
$x_0 \gets 0$ // hidden state, $O(N)$ memory
for $k \gets 1$ to $n$ do
  $x_k \gets \overline{A}\,x_{k-1} + \overline{B}\,u_k$ // one matrix-vector update
  $y_k \gets C\,x_k$ // read out the output
return $y_{1:n}$
```

## 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:

$$
x_1 = \overline{B} u_1,
\quad
x_2 = \overline{A}\,\overline{B} u_1 + \overline{B} u_2,
\quad
x_3 = \overline{A}^{2}\overline{B} u_1 + \overline{A}\,\overline{B} u_2 + \overline{B} u_3,
$$

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}$.

> **Theorem (Convolutional form).** The discrete SSM output equals the input
> convolved with the **SSM convolution kernel**
> $$
> \overline{K} = \parens{C\overline{B},\; C\overline{A}\,\overline{B},\; C\overline{A}^{2}\overline{B},\; \dots,\; C\overline{A}^{n-1}\overline{B}},
> \qquad
> y = \overline{K} \ast u .
> $$

> **Proof.** By the unrolled state, $x_k = \sum_{j=0}^{k-1} \overline{A}^{j} \overline{B}\, u_{k-j}$.
> Apply $C$: $y_k = \sum_{j=0}^{k-1} \parens{C \overline{A}^{j} \overline{B}}\, u_{k-j}$,
> which is the discrete convolution $\parens{\overline{K} \ast u}_k$ with the $j$-th
> kernel tap $\overline{K}_j = C \overline{A}^{j} \overline{B}$. $\qed$

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}$:
$$
\overline{K} = (0.393,\; 0.239,\; 0.145,\; 0.088,\; 0.053,\; \dots),
$$
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.

$$
% caption: One SSM, two equivalent algorithms. The recurrence (left) is sequential, $O(1)$ state;
% the convolution (right) is parallel over the length, both computed from $\overline{A}, \overline{B}, C$.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  st/.style={draw, black, thick, circle, minimum size=8mm},
  io/.style={draw, black, thick, minimum width=7mm, minimum height=6mm}]
  \definecolor{acc}{HTML}{2348F2}
  % ---- recurrent view (left) ----
  \foreach \i in {1,2,3,4} {
    \node[io] (u\i) at (\i*1.5,-1.6) {$u_{\i}$};
    \node[st, draw=acc, text=acc] (x\i) at (\i*1.5,0) {$x_{\i}$};
    \node[io] (y\i) at (\i*1.5,1.6) {$y_{\i}$};
    \draw[->, black, thick] (u\i) -- (x\i);
    \draw[->, black, thick] (x\i) -- (y\i);
  }
  \foreach \i [evaluate=\i as \j using int(\i+1)] in {1,2,3}
    \draw[->, acc, thick] (x\i) -- (x\j) node[midway, above, font=\scriptsize, text=acc] {$\overline{A}$};
  \node[black, font=\footnotesize] at (3.0,-2.6) {\texttt{recurrent: sequential,} $O(1)$ \texttt{state}};
  % ---- convolutional view (right) ----
  \begin{scope}[xshift=9.2cm]
    \node[black, font=\footnotesize] (uu) at (0,1.0) {input $u$};
    \node[acc, font=\footnotesize] (kk) at (0,0.0) {kernel $K$};
    \draw[->, black, thick] (1.4,1.0) -- (2.4,1.0);
    \draw[->, acc, thick] (1.4,0.0) -- (2.4,0.0);
    \node[draw, black, thick, circle, minimum size=8mm] (conv) at (3.1,0.5) {conv};
    \draw[->, black, thick] (conv) -- (4.7,0.5) node[right, font=\footnotesize] {output $y$};
    % kernel taps, written as a stacked column to avoid comma/\dots glyph issues
    \node[acc, font=\scriptsize, anchor=west] at (-1.5,-1.2) {$K_0 = C\overline{B}$};
    \node[acc, font=\scriptsize, anchor=west] at (-1.5,-1.8) {$K_1 = C\overline{A}\,\overline{B}$};
    \node[acc, font=\scriptsize, anchor=west] at (-1.5,-2.4) {$K_2 = C\overline{A}^{2}\overline{B}$};
    \node[black, font=\footnotesize] at (3.2,-2.6) {convolution: parallel over length};
  \end{scope}
\end{tikzpicture}
$$

| Property | Recurrent view | Convolutional view |
| --- | --- | --- |
| Update | $x_k = \overline{A} x_{k-1} + \overline{B} u_k$ | $y = \overline{K} \ast u$ |
| Parallelism | sequential in $k$ | parallel over the sequence |
| Memory | $O(N)$ state | $O(n)$ kernel + FFT buffers |
| Best for | autoregressive generation | training, full-sequence encoding |
| Cost | $O(n N)$ | $O(n \log n)$ via FFT |

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:

| Tensor | Shape | Role |
| --- | --- | --- |
| input $u$ | $(B, L, H)$ | one scalar signal per channel per position |
| $A$ | $(H, N, N)$ | one state matrix per channel (diagonal: $(H, N)$) |
| $B$, $C$ | $(H, N)$ | input / output maps per channel |
| $\Delta$ | $(H,)$ | one step size per channel (S4); $(B, L, H)$ in Mamba |
| state $x_k$ | $(B, H, N)$ | carried across steps in the recurrence |
| kernel $\overline{K}$ | $(H, L)$ | one length-$L$ convolution filter per channel |
| output $y$ | $(B, L, H)$ | same shape as the input |

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.

$$
% caption: Tensor-shape flow through one deep SSM layer, dimensions labeled. Input $(B,L,H)$ splits
% into $H$ independent SISO systems, each an $N$-state recurrence, then recombines to $(B,L,H)$.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  box/.style={draw, black, thick, minimum width=16mm, minimum height=9mm, align=center},
  ssm/.style={draw, draw=acc, text=acc, thick, minimum width=20mm, minimum height=8mm, align=center}]
  \definecolor{acc}{HTML}{2348F2}
  \node[box] (in) at (0,0) {input\\($B$,$L$,$H$)};
  \node[ssm] (s1) at (4.2,1.5) {channel $1$\\$N$ states};
  \node[ssm] (s2) at (4.2,0.0) {channel $2$\\$N$ states};
  \node[acc, font=\scriptsize] (dots) at (4.2,-1.15) {$\vdots$};
  \node[ssm] (sH) at (4.2,-2.3) {channel $H$\\$N$ states};
  \node[box] (out) at (8.6,0) {output\\($B$,$L$,$H$)};
  \draw[->, black, thick] (in.east) -- (s1.west);
  \draw[->, black, thick] (in.east) -- (s2.west);
  \draw[->, black, thick] (in.east) -- (sH.west);
  \draw[->, acc, thick] (s1.east) -- (out.west);
  \draw[->, acc, thick] (s2.east) -- (out.west);
  \draw[->, acc, thick] (sH.east) -- (out.west);
  \node[black, font=\footnotesize, anchor=north] at (4.2,-3.0) {$H$ \texttt{independent SISO systems, state} ($B$,$H$,$N$)};
\end{tikzpicture}
$$

## 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**.

> **Definition (HiPPO state).** Fix a basis of orthogonal polynomials on the past
> of $u$. The HiPPO matrix $A$ is the one for which the linear state $x(t)$ tracks,
> at every $t$, the coefficients of the best degree-$N$ polynomial approximation of
> the history $u_{\le t}$ under a chosen measure. The state is thus a compressed,
> reconstructable summary of all the past, not a leaky average.

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

$$
A_{nk} =
\begin{cases}
\,-\sqrt{(2n+1)(2k+1)}, & k < n,\\[2pt]
\,-(n+1), & k = n,\\[2pt]
\,0, & k > n,
\end{cases}
$$

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.

$$
% caption: Memory decay of the kernel taps with sequence distance. Random $A$ (red) forgets
% geometrically; HiPPO-initialized $A$ (blue) holds a near-flat, reconstructable summary of the past.
\begin{tikzpicture}[>=stealth, font=\footnotesize, x=0.62cm, y=1.0cm]
  \definecolor{acc}{HTML}{2348F2}
  \definecolor{red}{HTML}{C0392B}
  \draw[->, thick] (0,0) -- (13,0) node[right, font=\footnotesize] {\texttt{distance into the past}};
  \draw[->, thick] (0,0) -- (0,3.8) node[above, font=\scriptsize] {retained signal};
  % HiPPO: slow, near-flat decay (blue)
  \draw[acc, very thick]
    (0.3,3.3) .. controls (4,3.05) and (8,2.75) .. (12.4,2.35);
  \node[acc, font=\footnotesize, anchor=west] at (7.2,2.95) {HiPPO $A$};
  % random A: fast geometric decay (red)
  \draw[red, very thick]
    (0.3,3.3) .. controls (1.6,1.1) and (3.0,0.35) .. (12.4,0.18);
  \node[red, font=\footnotesize, anchor=west] at (3.4,1.0) {random $A$};
  \node[black, font=\footnotesize] at (1.2,-0.45) {\texttt{recent}};
  \node[black, font=\footnotesize] at (11.0,-0.45) {\texttt{distant}};
\end{tikzpicture}
$$

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.

| Parameterization | Form of $A$ | Kernel cost | Memory quality |
| --- | --- | --- | --- |
| HiPPO (dense) | full $N \times N$ | $O(n N^2)$ | optimal, but slow |
| S4 (DPLR) | $\Lambda - P P^{\ast}$ | $O\parens{(N + n)\log n}$ | matches HiPPO |
| S4D / DSS (diagonal) | diagonal $\Lambda$ | $O(n N)$ | nearly matches |
| S5 (MIMO + scan) | diagonal, scanned | $O(n N)$ parallel | matches, simpler |

## 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.

> **Definition (Selective SSM).** Make the parameters **functions of the input**.
> At each step, project the current token $u_k$ to produce step-specific
> $B_k = \Linear_B(u_k)$, $C_k = \Linear_C(u_k)$, and
> $\Delta_k = \softplus\parens{\Linear_\Delta(u_k)}$,
> then discretize per step:
> $$
> x_k = \overline{A}_k\,x_{k-1} + \overline{B}_k\,u_k,
> \qquad
> y_k = C_k\,x_k,
> $$
> with $\overline{A}_k = e^{\Delta_k A}$. The state matrix $A$ stays fixed and
> structured (diagonal); selectivity rides on $B_k, C_k, \Delta_k$.

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.

| Quantity | Time-invariant SSM (S4) | Selective SSM (Mamba) |
| --- | --- | --- |
| $\overline{A}$ | fixed | $e^{\Delta_k A}$, varies via $\Delta_k$ |
| $\overline{B}, C$ | fixed | $B_k, C_k$ projected from $u_k$ |
| $\Delta$ | fixed scalar | $\Delta_k$ per token (a gate) |
| Convolution | available (static kernel) | none (kernel changes per step) |
| Memory behavior | uniform over the sequence | content-selective |

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.

> **Theorem (Linear recurrence is a parallel scan).** The recurrence
> $x_k = \overline{A}_k x_{k-1} + \overline{B}_k u_k$ is the prefix-scan of the
> associative operator
> $$
> (A_2, b_2) \bullet (A_1, b_1) = (A_2 A_1,\; A_2 b_1 + b_2),
> $$
> applied to the elements $(\overline{A}_k,\, \overline{B}_k u_k)$. A parallel scan
> therefore computes all states $x_{1:n}$ in $O(\log n)$ sequential depth.

> **Proof.** Take the operator's identity element $(I, 0)$ and define the partial
> products $\pi_k = (\overline{A}_k, \overline{B}_k u_k) \bullet \cdots \bullet
> (\overline{A}_1, \overline{B}_1 u_1)$. Associativity follows from the block-matrix
> form $\left(\begin{smallmatrix} A & b \\ 0 & 1 \end{smallmatrix}\right)$, whose
> products compose exactly as $\bullet$. The second component of $\pi_k$ unrolls to
> $\sum_{j} \parens{\prod_{i>j} \overline{A}_i} \overline{B}_j u_j = x_k$. Any
> balanced-tree scan evaluates an associative fold in $O(\log n)$ depth and $O(n)$
> work. $\qed$

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.

$$
% caption: The selective scan. Per-token projections set $\Delta_k, B_k, C_k$ (gating); a parallel
% prefix scan over the recurrence fuses them in $O(\log n)$ depth inside fast on-chip memory.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  tok/.style={draw, black, thick, minimum width=8mm, minimum height=6mm},
  gate/.style={draw, draw=acc, text=acc, thick, minimum width=12mm, minimum height=6mm, align=center},
  node/.style={draw, black, thick, circle, minimum size=6mm}]
  \definecolor{acc}{HTML}{2348F2}
  \definecolor{green}{HTML}{1F9D4D}
  % input tokens
  \foreach \i in {1,2,3,4} \node[tok] (u\i) at (\i*2.2,0) {$u_{\i}$};
  % per-token selection
  \foreach \i in {1,2,3,4} {
    \node[gate] (g\i) at (\i*2.2,1.5) {select $\i$};
    \draw[->, black, thick] (u\i) -- (g\i);
  }
  \node[acc, font=\footnotesize, anchor=west] at (9.6,1.5) {selection};
  % scan tree (two levels) feeding the outputs
  \node[node] (p1) at (3.3,3.0) {};
  \node[node] (p2) at (7.7,3.0) {};
  \node[node, draw=green, text=green] (root) at (5.5,4.3) {};
  \draw[->, black, thick] (g1) -- (p1);  \draw[->, black, thick] (g2) -- (p1);
  \draw[->, black, thick] (g3) -- (p2);  \draw[->, black, thick] (g4) -- (p2);
  \draw[->, green, thick] (p1) -- (root); \draw[->, green, thick] (p2) -- (root);
  \node[green, font=\footnotesize, anchor=west] at (6.2,4.3) {parallel scan, $O(\log n)$ depth};
  % outputs
  \foreach \i in {1,2,3,4} {
    \node[tok] (y\i) at (\i*2.2,5.6) {$y_{\i}$};
  }
  \draw[->, green, thick] (root) -- (5.5,5.2) -- (2.2,5.2) -- (y1);
  \draw[->, green, thick] (5.5,5.2) -- (y2);
  \draw[->, green, thick] (5.5,5.2) -- (6.6,5.2) -- (y3);
  \draw[->, green, thick] (6.6,5.2) -- (8.8,5.2) -- (y4);
\end{tikzpicture}
$$

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.

$$
% caption: One Mamba block, dimensions labeled. An input projection expands $(B,L,D)$ to $(B,L,E)$
% and splits it; the main branch runs a causal conv, SiLU, and the selective SSM (state size $N$),
% gated by a SiLU branch, then a projection contracts back to $(B,L,D)$.
\begin{tikzpicture}[>=stealth, font=\footnotesize,
  proj/.style={draw, black, thick, minimum width=15mm, minimum height=8mm, align=center},
  op/.style={draw, draw=acc, text=acc, thick, minimum width=15mm, minimum height=8mm, align=center},
  gate/.style={draw, draw=green, text=green, thick, minimum width=15mm, minimum height=8mm, align=center},
  mul/.style={draw, black, thick, circle, inner sep=1pt, minimum size=6mm}]
  \definecolor{acc}{HTML}{2348F2}
  \definecolor{green}{HTML}{1F9D4D}
  \node[proj] (in)  at (0,0)     {\texttt{input proj}\\($B$,$L$,$D$)};
  \node[acc, font=\footnotesize, anchor=north] at (-0.2,-1.35) {\texttt{split to} ($B$,$L$,$E$)};
  % main branch
  \node[op]  (conv) at (3.2,1.1) {\texttt{causal conv}};
  \node[op]  (silu) at (5.9,1.1) {SiLU};
  \node[op]  (ssm)  at (8.6,1.1) {\texttt{sel. SSM}\\\texttt{state} $N$};
  % gate branch
  \node[gate] (g)   at (5.9,-1.1) {\texttt{SiLU gate}};
  \node[mul] (m)    at (10.9,0.0) {};
  \draw[black, thick] (10.75,-0.15) -- (11.05,0.15);
  \draw[black, thick] (10.75,0.15) -- (11.05,-0.15);
  \node[proj] (out) at (13.6,0.0) {\texttt{out proj}\\($B$,$L$,$D$)};
  \draw[->, black, thick] (in.east) |- (conv.west);
  \draw[->, acc, thick] (conv) -- (silu);
  \draw[->, acc, thick] (silu) -- (ssm);
  \draw[->, acc, thick] (ssm.east) -| (m.north);
  \draw[->, black, thick] (in.east) |- (g.west);
  \draw[->, green, thick] (g.east) -| (m.south);
  \draw[->, black, thick] (m) -- (out);
  \node[black, font=\footnotesize, anchor=north] at (6.8,-2.0) {\texttt{main branch width} $E = 2D$\texttt{; gate multiplies elementwise}};
\end{tikzpicture}
$$

```algorithm
caption: $\textsc{SelectiveScan}(u_{1:n}, A)$ — Mamba's input-dependent linear recurrence
for $k \gets 1$ to $n$ do // per-token selection (parallel)
  $\Delta_k \gets \softplus(W_\Delta\, u_k)$ // input-dependent step size
  $B_k \gets W_B\, u_k$; $\;C_k \gets W_C\, u_k$ // input-dependent in/out maps
  $\overline{A}_k \gets \exp(\Delta_k A)$; $\;\overline{B}_k \gets \Delta_k B_k$ // discretize per step
$x_{1:n} \gets \textsc{ParallelScan}\parens{(\overline{A}_k,\, \overline{B}_k u_k)}$ // $O(\log n)$ depth, on-chip
for $k \gets 1$ to $n$ do
  $y_k \gets C_k\, x_k$ // selective read-out
return $y_{1:n}$
```

## 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.

> **Theorem (Linear-time, constant-state inference).** A selective SSM with state
> dimension $N$ and $H$ channels generates each new token in $O(N H)$ time and
> $O(N H)$ memory, independent of the sequence position $n$. A Transformer needs
> $O(n d)$ time and an $O(n d)$ key-value cache at position $n$.

> **Proof.** Each SSM step is one matrix-vector update $x_k = \overline{A}_k x_{k-1} +
> \overline{B}_k u_k$ per channel, costing $O(N)$ per channel and $O(NH)$ total,
> with the state $x$ the only thing retained. Attention at position $n$ scores the
> new query against all $n$ cached keys, $O(nd)$ time, and must store every past
> key and value, $O(nd)$ memory. The SSM bound has no $n$; the attention bound is
> linear in $n$. $\qed$

$$
% caption: Inference cost per generated token versus context length. Attention grows linearly
% with the KV-cache (blue); the SSM stays flat at a constant-size state (black).
\begin{tikzpicture}[>=stealth, font=\footnotesize, x=1.0cm, y=1.0cm]
  \definecolor{acc}{HTML}{2348F2}
  \draw[->, thick] (0,0) -- (7.4,0) node[right, font=\footnotesize] {\texttt{context length} $n$};
  \draw[->, thick] (0,0) -- (0,4.4) node[above, font=\footnotesize] {\texttt{cost per token}};
  % attention: linear growth
  \draw[acc, very thick] (0.3,0.5) -- (6.8,3.9) node[right, text=acc, font=\footnotesize] {\texttt{attention} $O(n)$};
  % SSM: flat constant
  \draw[black, very thick] (0.3,1.0) -- (6.8,1.0) node[right, text=black, font=\footnotesize] {\texttt{SSM~} $O(1)$};
  \fill[acc] (3.55,2.1) circle (1.6pt);
  \fill[black] (3.55,1.0) circle (1.6pt);
\end{tikzpicture}
$$

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.

| Model | Mechanism | Train | Inference state | LRA / long-range |
| --- | --- | --- | --- | --- |
| LSTM | gated nonlinear recurrence | $O(n)$ sequential | $O(1)$ | weak past ~$10^2$ |
| Transformer | full attention | $O(n^2)$ parallel | $O(n)$ KV-cache | quadratic memory wall |
| S4 | structured LTI SSM | $O(n \log n)$ | $O(N)$ | first to solve all LRA tasks |
| Mamba | selective SSM + scan | $O(n)$ parallel | $O(N)$ | matches Transformers, $5\times$ faster gen |

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' = 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.

[^gf-seq]: **Goodfellow**, _Deep Learning_, Ch. 10 — Sequence Modeling: the recurrent baseline ($O(n)$ but vanishing-gradient-bound) against which linear state-space models are the modern alternative; deep SSMs postdate the 2016 text.