Selective State Space Models

Start from a linear time-invariant system mapping a scalar input stream to an output through an $N$-dimensional latent state.

To run it on tokens you discretise with a step size $\Delta$. Zero-order hold gives the matrices the recurrence actually uses.

With $A$, $B$, $C$ constant the recurrence is a convolution, computable for the whole sequence in one FFT. That is S4, and it is also its ceiling: a time-invariant system cannot choose what to remember, so it cannot solve selective copying or induction.

Selection

Mamba makes $B$, $C$ and $\Delta$ functions of the current token. The system stops being time-invariant, the convolution disappears, and what remains is a recurrence that must be scanned.

$\Delta_t$ is the interesting one. Large $\Delta$ makes $\bar{A} \to 0$ and the state resets to the current input; small $\Delta$ makes $\bar{A} \to I$ and the token is ignored. It is a gate in the LSTM sense, arrived at from the discretisation rather than bolted on.

The cost model inverts the transformer’s. Decoding is $O(1)$ per token in context length with a fixed-size state, but that state is a lossy summary — exact retrieval from far back is the thing attention buys and this does not. Hybrids that interleave a few attention layers among many SSM layers are the current answer.

import torch
from torch import Tensor


def selective_scan(
    x: Tensor,        # [B, L, D]   input
    A: Tensor,        # [D, N]      log-negative, learned
    B: Tensor,        # [B, L, N]   input-dependent
    C: Tensor,        # [B, L, N]   input-dependent
    delta: Tensor,    # [B, L, D]   input-dependent, softplus'd
) -> Tensor:
    batch, length, dim = x.shape
    state = torch.zeros(batch, dim, A.size(-1), device=x.device)
    out = []

    for t in range(length):
        dt = delta[:, t].unsqueeze(-1)                  # [B, D, 1]
        a_bar = torch.exp(dt * A)                       # [B, D, N]
        b_bar = dt * B[:, t].unsqueeze(1)               # [B, 1, N] -> broadcast
        state = a_bar * state + b_bar * x[:, t].unsqueeze(-1)
        out.append((state @ C[:, t].unsqueeze(-1)).squeeze(-1))

    return torch.stack(out, dim=1)

$A$ is stored as $-\exp(A_{\log})$ so the eigenvalues stay negative and the system stays stable; $\bar{A} = \exp(\Delta A)$ then lies in $(0, 1)$ and the state cannot blow up. The $\bar{B}$ above drops the $(\Delta A)^{-1}(\exp(\Delta A) - I)$ factor for the first-order approximation $\Delta B$, which is what Mamba itself uses.