Understand SSMs &
Mamba

Chapter 1: Linear Recurrences

The simplest alternative to attention is a linear recurrence. At each step, we maintain a hidden state h and update it:

The matrix A controls how the previous state decays and mixes (it's the "memory" operator). The matrix B controls how much of the new input is written into the state. The output y_t = C · h_t reads a projection of the state. Processing each token costs constant time: one matrix-vector product for A·h, one for B·x, one for C·h. No matter how long the sequence, each new token costs the same.

If the state has dimension N and the input has dimension D, the per-step cost is O(N² + ND) for the matrix multiplications. Compare this to attention's O(L · D) per token at generation (where L is the sequence length so far). For long sequences, the SSM wins because N is fixed (typically 16–64), while L keeps growing.

Chapter 2: The State Space Model

A state space model (SSM) comes from control theory. It starts in continuous time:

This is a differential equation: h' is the derivative of the hidden state. To use it on discrete sequences (like text), we discretize it using a step size Δ. The most common method (zero-order hold) gives:

Why O(L), Not O(L²)?

The recurrence h_k = Ā h_k-1 + B̄ x_k processes each token in constant time — one matrix-vector multiply, regardless of how long the sequence is. For a sequence of length L, the total cost is O(L). Attention, by contrast, computes pairwise scores between all L tokens: O(L²). The difference is enormous at scale.

Unrolling the recurrence reveals something beautiful. Substituting repeatedly:

This is a convolution: y = K * x where the kernel K_i = C Āⁱ B̄. When A, B, C are fixed, this kernel is precomputable. You convolve it with the input using FFT in O(L log L). Training uses this parallel path; inference uses the sequential recurrence for constant memory.

Chapter 3: The S4 Trick

The breakthrough of S4 (Structured State Spaces for Sequence Modeling) was solving the long-range memory problem. Previous SSMs forgot quickly — information from 1000 steps ago had decayed to nearly zero. S4 fixed this with two key ideas:

1. HiPPO Initialization

Instead of random initialization, the A matrix is set to the HiPPO (High-Order Polynomial Projection Operator) matrix. This special matrix is designed to optimally compress the history of a continuous signal into a fixed-size state. Each element of h_t stores a coefficient of a polynomial approximation of the entire input history.

2. Diagonal + Low-Rank Structure

Computing Ā = exp(ΔA) naively for an N×N matrix costs O(N³). S4 decomposes A into diagonal + low-rank form, reducing this to O(N). This makes training practical for state sizes of 64 or more.

What S4 Actually Computes

S4 decomposes A into Diagonal Plus Low-Rank (DPLR) form: A = Λ − PQ*, where Λ is diagonal and P, Q are low-rank matrices. This is critical because computing exp(ΔA) for a general N×N matrix costs O(N³). With DPLR structure, it reduces to O(N). The convolution kernel K can then be computed via a Cauchy kernel formula, making the entire training step O(L log L) via FFT.

In subsequent work, S4D (diagonal state spaces) simplified this further: just use a diagonal A matrix, initialized from the HiPPO eigenvalues. This drops the low-rank components entirely, making implementation much simpler while retaining most of the performance. Mamba uses diagonal A by default.

The HiPPO matrix specifically captures a polynomial projection: each element of the state h_t stores a Legendre polynomial coefficient of the input history. Element h_t[k] approximates the k-th Legendre coefficient of x(s) over [0, t]. This means the state literally stores a polynomial approximation of everything it has seen — a lossless-ish compression of the entire history into a fixed-size vector.

Chapter 4: Selective Scan (Mamba)

S4 has a critical limitation: A, B, C are the same for every input token. The model can't decide to "pay more attention" to important tokens or "forget" irrelevant ones. Mamba (Gu & Dao, 2023) fixes this by making B, C, and Δ functions of the input.

This is the selective in "selective scan." When the model sees an important token, it can increase Δ (larger step = more input mixed in) and adjust B to write strongly to the state. For irrelevant tokens, it can shrink Δ and effectively skip them.

What "Selective" Actually Means

Think of Δ_t as a gate width. When Δ is large, exp(ΔA) shrinks the old state more and B̄ grows — the new input floods in. When Δ is small, the old state is preserved and the input barely registers. The model learns to open the gate wide for important tokens (names, keywords, punctuation) and keep it narrow for filler words.

B_t controls what is written to the state, C_t controls what is read out, and Δ_t controls how much the state changes. All three are linear projections of the current input x_t — making the entire SSM dynamics content-dependent. This is what S4 couldn't do.

The parameter generation is lightweight. For each token x_t in R^d_inner: B_t = W_B x_t where W_B is [d_state, d_inner], producing B_t in R^d_state. Similarly for C_t. For Δ_t: a linear projection to R^d_inner followed by softplus ensures Δ > 0. The total parameter overhead for selectivity is minimal — just three small linear layers per block.

python
import torch
import torch.nn.functional as F

def selective_ssm_step(x_t, h_prev, A, W_B, W_C, W_delta, bias_delta):
    # 1. Generate input-dependent parameters
    B_t = W_B @ x_t                         # [d_state]
    C_t = W_C @ x_t                         # [d_state]
    delta_t = F.softplus(W_delta @ x_t + bias_delta)  # [d_inner], > 0

    # 2. Discretize (A is in log-space, shape [d_inner, d_state])
    A_real = torch.exp(A)                   # negative reals
    A_bar = torch.exp(delta_t.unsqueeze(1) * A_real)  # [d_inner, d_state]
    B_bar = delta_t.unsqueeze(1) * B_t.unsqueeze(0)   # [d_inner, d_state]

    # 3. Update state
    h_t = A_bar * h_prev + B_bar * x_t.unsqueeze(1)  # [d_inner, d_state]

    # 4. Read output
    y_t = (h_t * C_t.unsqueeze(0)).sum(dim=1)  # [d_inner]

    return y_t, h_t

Chapter 5: Hardware-Aware Scan

Making B, C, Δ input-dependent means we can't use convolution for training. The naive approach — a sequential for-loop — would be unbearably slow on GPUs. Mamba uses a parallel scan algorithm that computes the full recurrence in O(log n) parallel steps instead of O(n) sequential ones.

The Parallel Scan

The key insight: a linear recurrence h_t = a_t h_t-1 + b_t is an associative operation. Like addition, we can rearrange the order of computation. Instead of left-to-right, we use a binary tree:

The Fused CUDA Kernel

Mamba doesn't just use a parallel scan — it fuses the entire SSM computation (discretization, scan, output projection) into a single GPU kernel. The key bottleneck on modern GPUs isn't compute but memory bandwidth. Reading and writing intermediate tensors of shape [B, L, d_inner, d_state] to global memory (HBM) would dominate runtime.

Instead, Mamba's kernel keeps the state (d_state=16 floats per channel) entirely in SRAM (the GPU's fast on-chip scratchpad). The associative scan runs in registers, and only the final output [B, L, d_inner] is written back to HBM. This reduces memory I/O by a factor proportional to d_state, making the selective SSM nearly as fast as a simple matrix multiply.

Chapter 6: The Mamba Block

A Mamba block looks quite different from a Transformer block. It has no attention mechanism at all. Instead, it uses a combination of linear projections, 1D convolution, and the selective SSM:

Concrete Shapes (Mamba-2.8B)

For Mamba-2.8B: d_model=2560, d_state=16, d_conv=4, expand_factor=2, so inner_dim=5120. Here's the exact data flow through one block:

The SSM state has shape [B, 5120, 16] — each of the 5120 channels maintains its own 16-dimensional state vector. During inference, this state is all you carry forward. No KV cache needed.

Chapter 7: Hybrid Models

Pure Mamba models are fast but sometimes struggle with tasks requiring precise retrieval of specific tokens from the past (like "what was the 7th word?"). Pure Transformers are great at retrieval but expensive. The natural solution: combine both.

Key Hybrid Architectures

How Jamba Interleaves

Jamba (AI21, 2024) uses a concrete pattern: for every 7 Mamba layers, insert 1 attention layer. This means only ~12% of layers are attention, but that's enough for strong retrieval performance. Combined with Mixture-of-Experts (MoE) at each layer, Jamba-1.5 fits 256K context in a single 80GB GPU — something no pure Transformer of equivalent quality can do.

Zamba (Zyphra) goes further: it uses a single shared attention layer whose weights are reused every N Mamba blocks. This dramatically reduces the parameter count from attention (which is the most parameter-heavy component) while still providing the retrieval capability where needed.

Chapter 8: Training & Inference

Mamba models are trained with the exact same objective as Transformers: next-token prediction with cross-entropy loss. The training data, tokenizer, and optimizer are all the same. The difference is entirely in the architecture.

Training: Parallel Mode

During training, the entire sequence is available. The selective SSM uses the parallel scan to process all tokens simultaneously. This is analogous to how Transformers compute all attention weights at once.

Inference: Recurrence Mode

During generation, we switch to sequential recurrence. At each step:

Memory: A Concrete Example

Consider a 2.8B model generating a 100K-token sequence:

The Mamba state is tiny: d_inner × d_state × n_layers × 2 bytes = 5120 × 16 × 48 × 2 ≈ 7.5 MB per sequence. This is why Mamba shines for long-context applications like book-length analysis, genomics, and continuous audio processing.

Chapter 9: SSM vs Transformer Tradeoffs

Neither architecture is strictly better. Each has fundamental strengths and weaknesses rooted in their core mechanisms.

The Copying Test

One of the simplest ways to expose SSM vs Transformer differences: the copying task. Present the model with "Copy this: ABCDE. Output: " and check if it reproduces ABCDE exactly. Transformers trivially solve this with induction heads — they can attend directly to the tokens to copy. SSMs struggle because the information must pass through the compressed state bottleneck. A 16-dimensional state vector can't perfectly store 5 arbitrary tokens.

This is precisely why hybrid models exist: the attention layers handle exact recall, and the SSM layers handle everything else (fluency, long-range context, style). In Jamba's design, the few attention layers disproportionately handle retrieval-like subtasks while the many Mamba layers handle the bulk of language modeling.