State-space models

Problem

A Transformer compares every token with every earlier token. Training on a sequence of length L costs time proportional to L², and generating each new token reads a KV cache that grows by one entry per token. At 8,000 tokens that is manageable; at half a million it dominates both memory and time. An RNN has the opposite profile: a fixed-size state and constant cost per token at inference, but training must step through time one position at a time, and gradients fade over long distances.

State-space models (SSMs) aim for both: constant-size state and linear cost at inference, and parallel training over the whole sequence. The price is a linear recurrence, which on its own cannot decide what to keep. The history of the family is how to make that recurrence remember long and, in Mamba, how to let it choose.

Intuition

Picture a set of leaky buckets, the state. Each time step some input pours in, every bucket leaks a fixed fraction, and the output is a weighted reading of the levels. Because the leaking and pouring are linear, an input at time 0 contributes to the output at time t by a fixed amount that depends only on the gap t. So you can compute the output either by simulating the buckets step by step (a recurrence) or by sliding one fixed filter over the input (a convolution). Training uses the convolution, which is parallel; generation uses the recurrence, which needs only the current levels.

The leak rate is controlled by a step size Δ. A large Δ means the state is mostly replaced by the new input; a tiny Δ means the input is ignored and the state carries on unchanged. Mamba's idea is to let each token pick its own Δ: an important token takes a large step and is written in, filler takes a tiny step and is skipped. The filter is then different at every position, so the convolution view is lost, and an associative scan recovers the parallelism.

Mechanics

The continuous system

An SSM maps a one-dimensional input signal u(t) to an output y(t) through an N-dimensional state h(t):

A is N × N, B is N × 1 and C is 1 × N. A model with D channels runs D independent copies, one per channel, usually with N between 4 and 64. The eigenvalues of A must have negative real parts, or the state blows up.

Discretization

Tokens arrive at discrete steps, so the equation is turned into a recurrence h_t = Ā h_t−1 + B̄ u_t with a step size Δ. Three rules are common:

• Euler: Ā = I + ΔA, B̄ = ΔB. Cheap, but unstable once Δ is large: for scalar a = −1, |1 − Δ| > 1 when Δ > 2.
• Bilinear (Tustin): Ā = (I − ΔA/2)⁻¹(I + ΔA/2), B̄ = (I − ΔA/2)⁻¹ΔB. Stable for every Δ whenever A is. S4 uses it.
• Zero-order hold (ZOH): assumes the input is constant over each step, giving Ā = exp(ΔA) and B̄ = A⁻¹(exp(ΔA) − I)B. Exact for piecewise-constant input. S4D and Mamba use it; for a diagonal A both are elementwise.

At small Δ all three agree (bilinear is closer to exact than Euler at Δ = 0.1). At Δ = 1 Euler forgets everything in one step; at Δ = 2.5 its factor of −1.5 makes the state grow and flip sign, while bilinear stays below 1 in magnitude and ZOH is exact by construction.

Recurrence and convolution

Unroll the recurrence from h₋₁ = 0: y_t = Σ_k=0..t CĀ^kB̄ u_t−k. That is a causal convolution with kernel K_k = CĀ^kB̄. In training, the whole sequence is known, so the kernel is computed once per layer and applied with an FFT in O(L log L) or as a Conv1d. In generation, the recurrence advances the N-dimensional state by one token at constant cost. The kernel is the system's impulse response: feed a single 1 followed by zeros through the recurrence and the outputs are K.

S4: a structured state matrix

A random A forgets quickly and its powers are expensive. S4 (Gu, Goel and Ré, 2021) starts from HiPPO-LegS, a matrix derived so that the state holds the coefficients of a Legendre-polynomial approximation of the whole input history so far:

This A is not diagonalizable in a numerically useful way, but it is normal plus low rank (NPLR): adding pp^T with p_n = √(n + ½) gives −½I plus a skew-symmetric matrix, which is normal and so unitarily diagonalizable. S4 uses that split and the Woodbury identity to compute the length-L kernel in nearly linear time instead of raising Ā to L powers. Later work found simpler choices nearly as good: S4D (Gu et al., 2022) keeps only a diagonal A, initialized from the diagonal part of the NPLR form or with real values −(n+1), and the kernel becomes a Vandermonde product.

Mamba: selection

An LTI system applies the same filter to every input. It cannot do something as simple as "remember the last value that was flagged and ignore everything else", because which inputs matter depends on their content. Mamba (Gu and Dao, 2023) makes Δ, B and C functions of the current token:

softplus keeps Δ positive. A large Δ_t drives Ā_t towards 0 and writes the token in; a Δ_t near 0 leaves Ā_t ≈ 1 and B̄_t ≈ 0, so the token is skipped. This is a gate, the same idea as an LSTM's forget gate, but the recurrence stays linear in h. A stays fixed and diagonal (S4D-real initialization), stored as log(−A). Each channel of the expanded inner dimension has its own A and Δ, and all channels share B_t and C_t.

A Mamba block wraps the selective SSM: a linear projection to twice the expanded width, a depthwise causal convolution of width 4 and a silu on one half, the SSM, multiplication by silu of the other half, and a projection back. Blocks are stacked with residual connections and normalization, with no attention and no MLP.

The parallel scan

With per-token Ā_t there is no single kernel. But the update h_t = a_th_t−1 + b_t composes: applying (a₁, b₁) then (a₂, b₂) is the single step (a₁a₂, a₂b₁ + b₂). That combine is associative, so a prefix scan computes every h_t in log₂L parallel rounds, the same trick as a parallel prefix sum (cumsum, or cumprod for the a half). Mamba's CUDA kernel goes further: it keeps the expanded state (L × D × N) in on-chip SRAM, never writes it to GPU memory, and recomputes it in the backward pass, the same IO-aware idea as FlashAttention (see GPU execution). Mamba-2 (Dao and Gu, 2024) restricts A to a scalar times the identity per head, which turns the scan into blocks of matrix multiplication that use tensor cores, and shows the result is a form of masked linear attention.

Worked example

Implementation

From primitives: an S4-style LTI layer computed both as a kernel and as a convolution, and a Mamba-style selective SSM as a plain loop. The kernel uses torch.einsum, and the convolution uses F.conv1d with left padding so no output sees the future.

import torch
import torch.nn.functional as F

def discretize_bilinear(A, B, dt):
    I = torch.eye(A.shape[0], dtype=A.dtype)
    left = torch.linalg.inv(I - dt / 2 * A)
    return left @ (I + dt / 2 * A), left @ (dt * B)

def ssm_kernel(A_bar, B_bar, C, L):
    powers = [torch.eye(A_bar.shape[0], dtype=A_bar.dtype)]
    for _ in range(L - 1):
        powers.append(A_bar @ powers[-1])
    return torch.einsum("n,lnm,m->l", C, torch.stack(powers), B_bar)   # K[l] = C A^l B

def ssm_conv(u, K):                                   # u: (batch, L)
    L = u.shape[-1]
    x = F.pad(u.unsqueeze(1), (L - 1, 0))             # causal: pad the left only
    return F.conv1d(x, K.flip(-1).view(1, 1, L)).squeeze(1)  # conv1d correlates, so flip

def selective_scan(x, delta, A, B, C):
    # x, delta: (batch, L, D); A: (D, N); B, C: (batch, L, N)
    h = x.new_zeros(x.shape[0], x.shape[2], A.shape[1])
    ys = []
    for t in range(x.shape[1]):
        dA = torch.exp(delta[:, t, :, None] * A)          # ZOH for diagonal A
        dB = (dA - 1) / A * B[:, t, None, :]
        h = dA * h + dB * x[:, t, :, None]
        ys.append(torch.einsum("bdn,bn->bd", h, C[:, t]))
    return torch.stack(ys, dim=1)

class SelectiveSSM(torch.nn.Module):
    def __init__(self, d, n=16):
        super().__init__()
        self.A_log = torch.nn.Parameter(torch.log(torch.arange(1, n + 1).float()).repeat(d, 1))
        self.to_delta = torch.nn.Linear(d, d)
        self.to_B = torch.nn.Linear(d, n, bias=False)
        self.to_C = torch.nn.Linear(d, n, bias=False)
        self.D = torch.nn.Parameter(torch.ones(d))

    def forward(self, x):                             # x: (batch, L, d)
        A = -torch.exp(self.A_log)                    # S4D-real: -1, -2, ..., -n
        delta = F.softplus(self.to_delta(x))
        return selective_scan(x, delta, A, self.to_B(x), self.to_C(x)) + self.D * x

The snippet has not been run against PyTorch here; it uses only the calls linked above, and the helpers behind this page compute the same recurrence, kernel and selective scan in plain Python. The loop is O(L) sequential steps: fine for checking, far too slow for training. The official Mamba package replaces it with a fused CUDA scan, and its reference implementation uses the simpler B̄ = ΔB rather than the exact ZOH form for B.

Tradeoffs

Check yourself

• Write the convolution kernel of a discretized LTI SSM in terms of Ā, B̄ and C.
• Why does Euler discretization fail for large Δ while bilinear and ZOH do not?
• What does the associative combine (a₁, b₁) ∘ (a₂, b₂) compute, and why does associativity matter?
• How much memory does a Mamba layer need at inference for a 1M-token context, compared with 1K?

Connections

SSMs are linear RNNs with gates like an LSTM's, and they compete with the Transformer and its attention. The memory comparison is with the KV cache, and the per-token cost matters for latency and throughput. The convolution view is a one-dimensional CNN with a kernel as long as the sequence. Why the fused scan kernel is fast is on GPU execution; matrix powers and eigenvalues are on the linear algebra toolkit. As language models, SSMs are trained like any decoder-only LLM; long time series are covered on time series.

References

• PyTorch v2.15.0: Conv1d, _conv_forward, einsum, conv1d, softplus, silu, cumsum, cumprod.
• Gu, Dao, Ermon, Rudra and Ré. HiPPO: Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020. arXiv:2008.07669
• Gu, Goel and Ré. Efficiently Modeling Long Sequences with Structured State Spaces. ICLR 2022. arXiv:2111.00396
• Gu, Gupta, Goel and Ré. On the Parameterization and Initialization of Diagonal State Space Models. NeurIPS 2022. arXiv:2206.11893
• Gu and Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2023. arXiv:2312.00752
• Dao and Gu. Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML 2024. arXiv:2405.21060
• Blelloch. Prefix Sums and Their Applications. CMU-CS-90-190, 1990.