Overview
Self-attention compares every token with every other token, so its cost grows quadratically with sequence length L. A state space model (SSM) instead carries a fixed-size hidden state h_t \in \mathbb{R}^N through the sequence and updates it with a linear recurrence, which costs time linear in L and constant memory per decoding step.
Earlier SSMs (S4 and relatives) used the same parameters at every position, so the model could not choose which tokens to keep or ignore. Mamba makes B, C and the step size \Delta functions of the input, which gives content-based selection. That change removes the global-convolution form S4 relied on for parallel training; Mamba replaces it with a hardware-aware parallel scan. Mamba-2 shows that a restricted selective SSM is equivalent to a form of linear attention.
This page works through three pieces: discretizing a continuous linear system, making it input-dependent, and parallelizing the resulting recurrence. The paper was accepted to the NeurIPS 2026 Education Track and ships with three notebooks and five exercises with solutions.
Linear recurrence
A discretized SSM maps an input sequence x_t to an output y_t through a hidden state:
\bar A and \bar B are discretized versions of continuous-time matrices A and B with step size \Delta. With fixed parameters the model is linear-time and can be unrolled into a convolution, but it treats every position identically.
Discretization
Take the scalar system h'(t) = a\,h(t) + b\,x(t) and hold the input constant, x(t) = x_k, on [k\Delta, (k+1)\Delta) (zero-order hold). Solving the linear ODE from h(k\Delta) = h_k gives
Evaluating at t = (k+1)\Delta:
For a < 0, \bar a = e^{a\Delta} \in (0,1) for every \Delta > 0. A larger \Delta forgets the previous state faster; it cannot destabilize the system. Notebook 1 compares the discrete recurrence with RK4 numerical integration of the continuous system a=-0.8,\; b=1,\; \Delta=0.05,\; L=50 and reports the maximum absolute error.
Zero-order hold calculator
| \bar a=e^{a\Delta} | |
|---|---|
| \bar b=(\bar a-1)b/a | |
| |\bar a|<1 (stable) |
For a=-4,\;b=2: \Delta=0.25 gives \bar a=0.3679,\;\bar b=0.3161. \Delta=2.0 gives \bar a=0.0003,\;\bar b=0.4998.
Selectivity
In a selective SSM the input determines how the state is updated:
Here A is diagonal and negative, and the input enters the state through B_t. Mamba additionally multiplies by the channel input and uses a per-channel \Delta_t; the notebook uses one scalar \Delta_t per timestep, shared across state channels, and the first-order form \Delta B in place of the ZOH expression for \bar b.
When \Delta_t is large, \bar a_t is near 0: the state is overwritten by the current token. When \Delta_t is near 0, \bar a_t is near 1 and the current token barely changes the state. The model learns from content which tokens to retain and which to skip.
Because the coefficients now differ at every position, the recurrence is no longer a convolution with a fixed kernel, and S4's FFT-based training does not apply.
Notebook 2 trains this block on a synthetic selective-copying task. A few marker positions carry a signal value in [−1, 1] and all other positions carry small background noise; the target at every step is the most recent marker value. The notebook compares the block with a fixed-parameter SSM and plots the learned \Delta_t against marker positions.
Tensor shapes in the reference block
For batch 16, length 30, input dimension 4 and state dimension 12:
| Tensor | Shape | Source |
|---|---|---|
A | (12,) | one decay rate per state channel |
B_t, C_t | (16,30,12) | nn.Linear(4, 12) on the last axis |
delta | (16,30,1) | nn.Linear(4, 1), \operatorname{softplus} |
a_bar, b_bar at step t | (16,12) | broadcast of (16,1) with (12,) or (16,12) |
h | (16,12) | shape unchanged across steps |
output y | (16,30,1) | sum over state axis, stack over time |
Parallel scan
The recurrence h_t = \bar a_t h_{t-1} + \bar b_t looks sequential, but each step is an affine map, and composition of affine maps is associative. Represent step t as a pair (\bar a_t, \bar b_t) with
Any associative operator can be evaluated as a scan in logarithmic depth. Notebook 3 implements a doubling scan: in each round every element combines with the element 2^k positions earlier, for k=0,1,\dots until 2^k reaches the padded length. The initial state is folded in as one extra element, so n=L+1, and the sequence is padded to the next power of two. The number of rounds is
| L | n=L+1 | Padded n | Rounds | Sequential steps |
|---|
The doubling scan performs O(n \log n) combine operations against n for the sequential loop, but its depth is \lceil \log_2 n \rceil, which a GPU can exploit. Mamba uses a fused, work-efficient, hardware-aware implementation. On CPU, the speedup in Notebook 3 comes from fewer, larger vectorized operations rather than hardware parallelism.
Duality with attention
Mamba-2 introduces structured state space duality: restricting the transition to a scalar-times-identity form makes the SSM equal to a masked, linear-attention-like operator. The same layer can therefore be computed either as a recurrence or as a matrix multiplication, and the matrix form allows a faster algorithm on hardware optimized for matmuls.
Comparison with attention and RNNs
L is sequence length, d the model dimension and N the SSM state size.
| Self-attention | Selective SSM | Vanilla RNN | |
|---|---|---|---|
| Compute per forward pass | O(L^2 d) | O(L N d) work with a work-efficient scan, O(\log L) depth | O(L d^2) |
| Training parallelism over L | Full | Via parallel scan | Sequential |
| Decoding memory per step | \mathrm{KV} cache, about 2Ld | State, O(Nd), independent of L | State, O(d) |
| Content-based selection | Yes | Yes (input-dependent \Delta, B, C) | Gated variants only |
Limits
A fixed-size state compresses the history. On tasks that need exact retrieval of arbitrary earlier tokens, such as copying or looking up a specific token far back in the context, attention's per-token cache has an advantage that a memory-only comparison does not show (Jelassi et al., 2024).
A running average of marker values is a case the Notebook 2 block cannot compute exactly. It would need \bar a_t = (n-1)/n after n markers, but \bar a_t is a function of x_t alone and the input carries no marker count. Exercise 3 works through this and the minimal architectural change that fixes it.
Teaching materials
| Item | Content |
|---|---|
| Notebook 1: From RNN to SSM | Plain linear recurrence; ZOH discretization checked numerically against the closed-form solution. |
| Notebook 2: Adding selectivity | Input-dependent B, C, \Delta; selective-copying task; fixed-parameter baseline; plot of learned \Delta_t against marker positions. |
| Notebook 3: Scan vs. sequential loop | Sequential and associative scan of the same recurrence; wall-clock benchmark with an optional GPU path. |
| Exercises and solutions | Discretization by hand, tensor shapes, breaking selective copying, scan depth and padding, SSM vs. attention compute and memory. |
Learning objectives
- State the recurrence of a discretized SSM and explain why a fixed-parameter SSM is linear-time but non-selective.
- Explain what selectivity adds and why it prevents the convolutional form used in S4.
- Implement a minimal selective-scan layer in PyTorch and run it on a toy task.
- Compare a selective SSM with self-attention and a vanilla RNN on compute, memory and long-range recall.
- Extend the notebook with a variant scan, such as a bidirectional scan, and evaluate it.
Toy tasks in Notebooks 1 and 2 run on CPU. The Notebook 3 benchmark has an optional GPU path.
References
- A. Gu and T. Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. COLM, 2024. arXiv:2312.00752.
- T. Dao and A. Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML, 2024.
- Y. Liu, Y. Tian, Y. Zhao, et al. VMamba: Visual State Space Model. NeurIPS, 2024.
- F. Rezaei Jafari, G. Montavon, K.-R. Müller, O. Eberle. MambaLRP: Explaining Selective State Space Sequence Models. NeurIPS, 2024.
- S. Jelassi, D. Brandfonbrener, S. M. Kakade, E. Malach. Repeat After Me: Transformers are Better than State Space Models at Copying. ICML, 2024.