Visual walkthrough — State-space models (Mamba, S4)

Before any symbol appears, meet the cast in plain words.

---

Step 1 — One box that remembers: the recurrence

WHAT. We draw a single "memory box". Each tick it does exactly two things: it shrinks the old bag a little, then adds a splash of the new input.

WHY. This is the smallest possible machine with memory. "Shrink then add" is the entire mechanic of an RNN too (see RNNs and LSTMs); an SSM is this box with a particularly clean shrink rule, which is what lets us do the magic in later steps.

PICTURE.

The one governing line, term by term:

Then we read out the bag through a lens $\mathbf{C}$:

$y_k=\mathbf{C}{h}_k$

$\mathbf{C}$ is just which mixture of the bag's numbers we report. (We set the skip term $\mathbf{D}=0$, as the parent did.)

---

Step 2 — Start empty, take one step

WHAT. We commit to starting with an empty bag: $h_0=0$. Then we run the box for exactly one tick.

WHY. Starting from zero removes any "initial condition" clutter, so every number in the bag can be traced back to some specific input. That traceability is the whole trick.

PICTURE.

$h_1=\bar{\mathbf{A}}\underset{=0}{\underbrace{h_0}}+\bar{\mathbf{B}}{x}_1=\bar{\mathbf{B}}{x}_1$

Read the arrows in the figure: the old bag was empty, so all of $h_1$ is fresh input $x_1$ scaled by $\bar{\mathbf{B}}$. Nothing else could contribute — there was nothing to shrink.

---

Step 3 — Take a second step: the past starts fading

WHAT. Feed in $x_2$. Now the box shrinks the previous bag (which held $x_1$) and adds fresh $x_2$.

WHY. This is the first moment memory shows up: $x_1$ is still present, but multiplied by $\bar{\mathbf{A}}$ once — it has aged one tick.

PICTURE.

$h_2=\bar{\mathbf{A}}{h}_1+\bar{\mathbf{B}}{x}_2=\underset{x_1,\text{aged once}}{\underbrace{\bar{\mathbf{A}}\bar{\mathbf{B}}{x}_1}}+\underset{x_2,\text{fresh}}{\underbrace{\bar{\mathbf{B}}{x}_2}}$

• The term $\bar{\mathbf{A}}\bar{\mathbf{B}}{x}_1$ is the echo of $x_1$, dimmed by one factor of $\bar{\mathbf{A}}$.
• $\bar{\mathbf{B}}{x}_2$ is the newcomer, at full strength.

Look at the two coloured contributions in the picture: same shape, different age.

---

Step 4 — The pattern locks in: powers of $\bar{\mathbf{A}}$

WHAT. Take one more step and watch the ages stack.

WHY. After three steps the rule is undeniable: the older an input, the more times it got multiplied by $\bar{\mathbf{A}}$. Age $=$ exponent.

PICTURE.

$h_3=\underset{\text{aged twice}}{\underbrace{{\bar{\mathbf{A}}}^2\bar{\mathbf{B}}{x}_1}}+\underset{\text{aged once}}{\underbrace{\bar{\mathbf{A}}\bar{\mathbf{B}}{x}_2}}+\underset{\text{fresh}}{\underbrace{\bar{\mathbf{B}}{x}_3}}$

Each input carries a power of $\bar{\mathbf{A}}$ equal to how many ticks ago it entered:

input	ticks ago at step 3	weight it carries
$x_1$	2	${\bar{\mathbf{A}}}^2\bar{\mathbf{B}}$
$x_2$	1	${\bar{\mathbf{A}}}^1\bar{\mathbf{B}}$
$x_3$	0	${\bar{\mathbf{A}}}^0\bar{\mathbf{B}}$

The staircase in the figure is the point: decay is geometric — each older step is another rung down.

---

Step 5 — Read out, and a kernel is born

WHAT. Push $h_k$ through the lens $\mathbf{C}$. Do it for the general step $k$.

WHY. Applying $\mathbf{C}$ to the sum-of-aged-inputs turns the state recurrence into a formula for the output directly in terms of past inputs. The list of weights that appears is the kernel.

PICTURE.

$y_k=\mathbf{C}{h}_k={\sum }_{i=1}^k\underset{\text{weight for input }{x}_i}{\underbrace{\mathbf{C}{\bar{\mathbf{A}}}^{k-i}\bar{\mathbf{B}}}}{x}_i$

• The sum runs over every past input $x_i$ (from the very first up to now).
• The exponent $k-i$ is the age: how many ticks between when $x_i$ arrived ($i$) and now ($k$).
• Each weight $\mathbf{C}{\bar{\mathbf{A}}}^{k-i}\bar{\mathbf{B}}$ is one number — how loudly an input of age $k-i$ still speaks in today's output.

Collect those weights by age and you get one fixed list, the SSM kernel:

---

Step 6 — The two views are the same picture

WHAT. Recognise the sum in Step 5 as a convolution: slide the kernel across the input and take weighted sums.

WHY. This is the payoff. A convolution can be computed with the FFT in $O(L\log L)$ time and — crucially — all outputs at once, in parallel. The recurrence (Step 1) is sequential and slow to train; the convolution is parallel and fast. Same numbers, two speeds.

PICTURE.

$y=\bar{\mathbf{K}}\ast x⟺y_k={\sum }_{i=1}^k{\bar{\mathbf{K}}}_{k-i}{x}_i$

The convolution view is exactly the same fixed-filter idea you meet in Convolutional Neural Networks — except here the filter's shape ($\bar{\mathbf{K}}$) is derived from the decay $\bar{\mathbf{A}}$ rather than learned as raw weights.

---

Step 7 — Edge & degenerate cases (never let the reader fall through)

WHAT / WHY / PICTURE for the corners the smooth story skips.

Each strip in the figure shows the same kernel machinery with a different $\bar{\mathbf{A}}$: dying, flat, exploding, ringing. One rule, every regime.

---

Step 8 — Where Mamba departs

WHAT. In Steps 1–6 the weights $\bar{\mathbf{A}},\bar{\mathbf{B}},\mathbf{C}$ were frozen — the same for every token — which is exactly what let us collect them into one fixed kernel $\bar{\mathbf{K}}$.

WHY it matters. Mamba makes those weights depend on the current input ($\bar{\mathbf{B}}(x_k)$, $\mathbf{C}(x_k)$, $\Delta (x_k)$). The moment the weight for age $k-i$ changes with $k$, you can no longer collect one fixed list — the convolution collapses, and you must run the recurrence (Step 1) directly. This is the same "content-dependent mixing" spirit as Linear Attention.

PICTURE.

The picture contrasts a fixed comb (S4: one kernel for all) with a per-token comb (Mamba: the filter reshapes at each step). The price is losing FFT parallelism; the prize is choosing what to remember.

---

The one-picture summary

Read it left to right: an empty bag → shrink-and-add each tick → each input aged by a power of $\bar{\mathbf{A}}$ → read out with $\mathbf{C}$ → the ages line up into the kernel $\bar{\mathbf{K}}$ → the same numbers are a convolution. That single chain is S4.

Recall Feynman retelling — say it back in plain words

Imagine a bucket. Each second you pour out a fixed fraction of what's in the bucket (that's $\bar{\mathbf{A}}$, the forgetting), then pour in a splash of whatever's arriving now (that's $\bar{\mathbf{B}}x$). You read the water level through a tinted glass ($\mathbf{C}$). Because you always pour out the same fraction, an input from $n$ seconds ago has been diluted exactly $n$ times — that's the power ${\bar{\mathbf{A}}}^n$. So today's reading is just every past input, each dimmed by its age. Line those age-weights up into a list and you've got a filter you can slide across the whole input in one shot — that's the convolution, and it's why training is fast. Mamba's twist: let the pour-out fraction change per token, so the bucket can decide to hold onto a name and dump a filler word. That flexibility costs you the "one-shot slide", so Mamba runs the bucket step by step with clever GPU code instead.

Recall Quick self-test

Why does an input $x_i$ carry the factor ${\bar{\mathbf{A}}}^{k-i}$ at output step $k$? ::: Because between arriving (step $i$) and being read (step $k$) it survived $k-i$ shrink operations, each multiplying by $\bar{\mathbf{A}}$. What single object are the recurrent and convolutional views both computing? ::: The kernel / impulse response $\bar{\mathbf{K}}=(\mathbf{C}\bar{\mathbf{B}},\mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}},\dots )$. Why can't Mamba use one fixed kernel? ::: Its weights depend on the input, so the "weight for age $k-i$" changes with position and can't be collected into a single fixed list. What happens to the kernel if $\bar{\mathbf{A}}=0$? ::: It becomes $(\mathbf{C}\bar{\mathbf{B}},0,0,\dots )$ — the model is memoryless, output depends only on the current input.