Exercises — State-space models (Mamba, S4)

• (a) → convolutional view. One kernel $\bar{\mathbf{K}}=(\mathbf{C}\bar{\mathbf{B}},\mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}},\dots )$ is the same for every position, so it is a plain convolution. This is what training uses (parallel, FFT-friendly).
• (b) → recurrent view: $h_k=\bar{\mathbf{A}}{h}_{k-1}+\bar{\mathbf{B}}{x}_k$. This is what generation/inference uses (one step at a time).
• (c) → selective (Mamba) mechanism. Making $\bar{\mathbf{A}},\bar{\mathbf{B}},\mathbf{C}$ depend on $x$ is exactly what Mamba adds on top of S4.

$h^{\prime }(t)$ is the rate of change of the hidden state — how much $h$ moves per tiny slice of time (the derivative). The matrix $\mathbf{A}$ multiplies the current state and therefore decides how the state pushes on itself: it governs the speed and direction of that change. $\mathbf{B}$ only controls how much the input nudges the state.

Use $h_k=\bar{\mathbf{A}}{h}_{k-1}+\bar{\mathbf{B}}{x}_k$, then $y_k=\mathbf{C}{h}_k$.

Step 1 ($x_1=1$): $h_1=\bar{\mathbf{A}}\cdot 0+\bar{\mathbf{B}}\cdot 1=\left[\begin{array}{c}0.5\\ 1.0\end{array}\right]$. $y_1=2(0.5)+(-1)(1.0)=1.0-1.0=0.$

Step 2 ($x_2=0$): input term vanishes, only decay: $h_2=\bar{\mathbf{A}}{h}_1=\left[\begin{array}{c}0.8\cdot 0.5\\ 0.4\cdot 1.0\end{array}\right]=\left[\begin{array}{c}0.4\\ 0.4\end{array}\right]$. $y_2=2(0.4)+(-1)(0.4)=0.8-0.4=\mathrm{0.4.}$

Step 3 ($x_3=0$): $h_3=\bar{\mathbf{A}}{h}_2=\left[\begin{array}{c}0.32\\ 0.16\end{array}\right]$. $y_3=2(0.32)+(-1)(0.16)=0.64-0.16=\mathrm{0.48.}$

This is exactly an impulse response: we hit the system with a single spike at $k=1$ and watched it decay. The slow dimension ($0.8$) keeps ringing longer than the fast one ($0.4$).

• $\mathbf{C}\bar{\mathbf{B}}=2(0.5)+(-1)(1.0)=0.$
• $\bar{\mathbf{A}}\bar{\mathbf{B}}=\left[\begin{array}{c}0.8\cdot 0.5\\ 0.4\cdot 1.0\end{array}\right]=\left[\begin{array}{c}0.4\\ 0.4\end{array}\right]$, so $\mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}=2(0.4)-0.4=\mathrm{0.4.}$
• ${\bar{\mathbf{A}}}^2\bar{\mathbf{B}}=\left[\begin{array}{c}0.64\cdot 0.5\\ 0.16\cdot 1.0\end{array}\right]=\left[\begin{array}{c}0.32\\ 0.16\end{array}\right]$, so $\mathbf{C}{\bar{\mathbf{A}}}^2\bar{\mathbf{B}}=0.64-0.16=\mathrm{0.48.}$

Kernel $=(0,0.4,0.48)$, which matches $(y_1,y_2,y_3)$. That is the whole point: the impulse response is the kernel. Convolving any input with this kernel gives the same output as running the recurrence.

(a) Stability means the state must not blow up when input stops. Over one step the state is multiplied by $\bar{A}=e^{A\Delta }$. We need $|\bar{A}|<1$ so repeated multiplication shrinks the state. Since $\Delta >0$, $e^{A\Delta }<1$ exactly when $A<0$. If $A>0$ the state grows without bound (unstable); if $A=0$ it never forgets (marginal).

(b) $\bar{A}=e^{-0.5\cdot 0.5}=e^{-0.25}\approx \mathrm{0.7788.}$ $\bar{A}=e^{-0.5\cdot 4}=e^{-2}\approx \mathrm{0.1353.}$

(c) Smaller $\Delta$ (here $0.5$) gives $\bar{A}\approx 0.78$, closer to $1$ ⇒ slow decay ⇒ longer memory. Larger $\Delta$ ($4$) gives $\bar{A}\approx 0.14$ ⇒ fast forgetting.

⚠️ Note the subtlety: for a fixed negative $A$, increasing $\Delta$ shrinks $\bar{A}$. The "large $\Delta$ = long memory" slogan from the parent note assumes the model co-adapts $A$; do not read it as a raw mathematical identity. See the trap below.

Why softplus: $\Delta$ is a timestep and must be strictly positive (a negative or zero timestep has no physical meaning in the discretisation $e^{A\Delta }$). Softplus maps all of $\mathbb{R}$ to $(0,\infty )$ smoothly, so gradients flow everywhere and $\Delta$ can never go negative. A plain linear output could produce negative $\Delta$ and break stability.

Values: $\text{softplus}(0)=\ln (1+e^0)=\ln 2\approx \mathrm{0.6931.}$ $\text{softplus}(2)=\ln (1+e^2)\approx \ln (8.389)\approx \mathrm{2.1269.}$

${\Delta }_{\text{John}}=\text{softplus}(3)=\ln (1+e^3)\approx \ln (21.09)\approx \mathrm{3.0486.}$ ${\Delta }_{\text{the}}=\text{softplus}(-3)=\ln (1+e^{-3})\approx \mathrm{0.04859.}$

${\bar{A}}_{\text{John}}=e^{-1\cdot 3.0486}=e^{-3.0486}\approx \mathrm{0.04740.}$ ${\bar{A}}_{\text{the}}=e^{-1\cdot 0.04859}=e^{-0.04859}\approx \mathrm{0.9526.}$

Here the "the" token has $\bar{A}\approx 0.95$ (slow decay), while "John" has $\bar{A}\approx 0.047$ (fast decay). Time to decay to $1/e$: number of steps $n$ with ${\bar{A}}^n=e^{-1}$, i.e. $n=\frac{-1}{\ln \bar{A}}$.

• "the": $n=-1/\ln (0.9526)\approx -1/(-0.04859)\approx 20.6$ steps.
• "John": $n=-1/\ln (0.0474)\approx -1/(-3.049)\approx 0.328$ steps.

Interpretation & the twist: with $A$ fixed at $-1$, a large $\Delta$ makes $\bar{A}$ small (fast forget). So the raw math says "John" (large $\Delta$) is forgotten fast — the opposite of the intuitive story. This is exactly why Mamba does not fix $A$: it must learn small $|A|$ jointly so that a large $\Delta$ can push $\bar{A}$ toward $1$. This exercise deliberately shows the failure mode of holding $A$ constant.

${\bar{A}}_{\text{John}}=e^{-0.02\cdot 3.049}=e^{-0.06098}\approx \mathrm{0.9408.}$ ${\bar{A}}_{\text{the}}=e^{-0.02\cdot 0.0486}=e^{-0.000972}\approx \mathrm{0.99903.}$

Hmm — with a single tiny $A$, "the" still has larger $\bar{A}$. The real lever is that Mamba's $\bar{\mathbf{A}}$ depends on $\Delta$ and the input-dependent $\mathbf{B},\mathbf{C}$ gating how much input actually enters and exits each mode. The cleaner statement: selectivity gives the model independent control per token over how much it writes ($\bar{\mathbf{B}}(x)$), how much it decays ($\Delta (x)$ through $\bar{\mathbf{A}}$), and how much it reads ($\mathbf{C}(x)$). Memory is the joint result of all three, not $\Delta$ alone.

Numerically confirmed: ${\bar{A}}_{\text{John}}\approx 0.9408$, ${\bar{A}}_{\text{the}}\approx 0.99903$.

• Attention: $L^{2}d=(10^5{)}^2\cdot 1=10^{10}.$
• Recurrent SSM: $LdN=10^5\cdot 1\cdot 16=1.6\times 10^6.$
• FFT SSM: $L{\log }_{2}Ld=10^5\cdot {\log }_2(10^5)\approx 10^5\cdot 16.61\approx 1.661\times 10^6.$
• Ratio attention/recurrent $=10^{10}/(1.6\times 10^6)=6250.$

So attention is roughly 6250× more expensive here. This is the entire motivation for SSMs on 100K-token sequences — the quadratic term $L^2$ explodes while the SSM stays effectively linear in $L$.

1. Training = whole sequence known. All inputs $x_1,\dots ,x_L$ are available at once. The output is $y=\bar{\mathbf{K}}\ast x$ with a single fixed kernel $\bar{\mathbf{K}}$. Convolution parallelises across all positions and runs in $O(L\log L)$ via FFT — ideal for GPUs.
2. Generation = one token at a time. During autoregressive decoding future tokens do not exist yet, so we cannot form a full convolution. The recurrence $h_k=\bar{\mathbf{A}}{h}_{k-1}+\bar{\mathbf{B}}{x}_k$ carries a constant-size state $h$, giving $O(1)$ work and $O(N)$ memory per new token — far cheaper than re-attending to all past tokens.
3. Why Mamba breaks the shortcut. Mamba makes $\bar{\mathbf{A}},\bar{\mathbf{B}},\mathbf{C}$ depend on each input token. The kernel is no longer one fixed filter — it changes per position — so the operation is not a convolution and FFT no longer applies.
4. Mamba's fix. It uses a hardware-aware parallel scan (an associative-scan of the linear recurrence) with kernel fusion and careful GPU-memory management, recovering near-FFT speed while keeping the input-dependent (selective) update.

One-line synthesis: S4 buys parallelism with a fixed kernel; Mamba trades the fixed kernel for selectivity and buys back parallelism with a fused scan.

A single impulse at $k=1$ produces outputs equal to the kernel entries ::: True — the impulse response is the SSM convolution kernel (Exercises 2.1 & 2.2).