State-Space Models

A class of sequence models that map continuous-time physical systems to discrete sequence representations, combining the parallel training of Convolutions with the $O(1)$ inference of Recurrent Neural Networks.

In One Sentence

"State-Space Models compress historical sequence context into a fixed-size latent state, enabling linear-time sequence processing during training via parallel convolutions or scans, and constant-time step generation during inference."

Intuition

Imagine tracking a moving vehicle. To predict its next position, you don't need to review the entire log of every coordinate it has ever visited (which is what Attention does, scaling quadratically with time). Instead, you maintain a running estimate of its current position and velocity—this is the latent state.

Classic State-Space Models (SSMs) act like physical filters. They process inputs through fixed mathematical transformations. Because these transformations are Linear Time-Invariant (LTI), the model behaves identically regardless of the content of the sequence. This stability allows us to unroll the entire sequence computation into a single massive, highly parallelizable convolution during training.

However, LTI models struggle with language. If a model must remember a specific password mentioned early in a prompt, an LTI system will slowly dilute that password over time as more words pass through.

Selective SSMs (like Mamba) solve this by making the transition parameters depend dynamically on the input. If the model encounters an important token, it opens the gates to write it to memory; if it processes filler words, it bypasses the state update, preserving critical historical context indefinitely.

Formal Definition

At their core, SSMs are inspired by continuous-time differential equations mapping a 1-D input signal $x(t) \in \mathbb{R}$ to an output $y(t) \in \mathbb{R}$ through an $N$-dimensional latent state $h(t) \in \mathbb{R}^N$:

$$h'(t) = A h(t) + B x(t)$$ $$y(t) = C h(t) + D x(t)$$

To apply this to discrete digital sequences (like text tokens), the continuous system must be discretized using a step size parameter $\Delta$. Using the standard Zero-Order Hold (ZOH) method, the continuous matrices $(A, B)$ are transformed into discrete parameters $(\bar{A}, \bar{B})$:

\bar{A} = \exp(\Delta A)$$ $$\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B

Once discretized, the model can be formulated as a recurrent sequence transition:

h_t = \bar{A} h_{t-1} + \bar{B} x_t$$ $$y_t = C h_t + D x_t

The Dual Representations

The magic of standard SSMs lies in their mathematical duality:

Recurrent View (Inference): Step-by-step updates require only $O(1)$ computation and $O(1)$ memory, as we only need to store the current state $h_t$.
Convolutional View (Training): If the parameters $(\bar{A}, \bar{B}, C)$ are constant over time (Time-Invariant), the entire recurrence can be computed as a single convolution:
$$y = x * \bar{K} \quad \text{where} \quad \bar{K} = (C\bar{B}, C\bar{A}\bar{B}, \dots, C\bar{A}^{L-1}\bar{B})$$
This allows training to scale at $O(L \log L)$ using Fast Fourier Transforms (FFTs).

Trade-offs & Decisions

When to Use

Ultra-long contexts: Sequences spanning 10k to 1M+ tokens where Transformer attention becomes computationally prohibitive.
Resource-constrained inference: Edge devices or high-throughput servers where keeping a massive KV-cache in GPU memory is impossible.
Continuous data modalities: Audio waveforms, sensor streams, and genomic data where local context and continuous signals dominate.

When to Avoid

Exact retrieval tasks: Complex multi-hop QA or code execution where the model must retrieve highly specific facts across arbitrary distances without lossy compression.
No custom hardware support: If you cannot compile specialized Triton or CUDA kernels, selective SSMs (which bypass convolutions) will run slowly on standard deep learning frameworks.

History & Evolution

Classic Structured SSMs

Models like S4 introduced structured state matrices (such as the HiPPO matrix) to solve the vanishing/exploding gradient problems of traditional RNNs, enabling long-range dependency tracking via parallel convolutions.

May 2023

RWKV: Reinventing RNNs

The RWKV architecture (2305.13048) formulated a linear attention mechanism as a recurrent network, showing that dense RNNs could scale to billions of parameters and perform on par with Transformers.

December 2023

Mamba & Selective State Spaces

Gu and Dao (2312.00752) introduced Mamba. By making SSM parameters input-dependent (selective), they solved the content-based reasoning bottleneck of LTI models, using a hardware-aware parallel scan to maintain training efficiency.

Interactive Walkthrough

Compare how a Linear Time-Invariant (LTI) SSM and a Selective SSM (Mamba) manage state compression when encountering a critical "Secret Key" token followed by filler words.

State Decay Rate ($\bar{A}$): 0.85

Sequence Processing & Memory Retention

Architectural Behavior

Loading simulation...

Selective SSM Step (PyTorch-style)

mamba_step.py Selective Scan

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelectiveSSMStep(nn.Module):
    def __init__(self, d_model, d_state):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # Continuous parameter A (fixed, structured)
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        
        # Projection layers to generate input-dependent B, C, and Delta
        self.x_proj = nn.Linear(d_model, d_state * 2 + d_model, bias=False)
        self.dt_proj = nn.Linear(d_model, d_model)

    def forward(self, x, h_prev):
        """
        x: (batch, d_model) - Current input token representation
        h_prev: (batch, d_model, d_state) - Previous hidden state
        """
        # 1. Project input to get selective parameters
        # B, C are dynamic (input-dependent)
        proj = self.x_proj(x)
        B, C, dt_raw = torch.split(proj, [self.d_state, self.d_state, self.d_model], dim=-1)
        
        # 2. Compute dynamic step size Delta
        delta = F.softplus(self.dt_proj(dt_raw)) # (batch, d_model)
        
        # 3. Discretize A and B (Zero-Order Hold approximation)
        # delta: (batch, d_model, 1), A: (d_model, d_state)
        delta = delta.unsqueeze(-1)
        A_bar = torch.exp(delta * self.A.unsqueeze(0)) # (batch, d_model, d_state)
        
        # B: (batch, d_state) -> (batch, 1, d_state)
        # B_bar = delta * B
        B_bar = delta * B.unsqueeze(1) # (batch, d_model, d_state)
        
        # 4. Update latent state
        # h_prev: (batch, d_model, d_state)
        # x: (batch, d_model) -> (batch, d_model, 1)
        h_next = A_bar * h_prev + B_bar * x.unsqueeze(-1)
        
        # 5. Compute output
        # C: (batch, d_state) -> (batch, 1, d_state)
        y = torch.einsum("bmd,bd->bm", h_next, C)
        
        return y, h_next

Common Variants

S4 (Structured State Spaces)

The foundational structured SSM. Uses a fixed, structured $A$ matrix (HiPPO) to capture long-range dependencies, computed purely via parallel convolutions during training.

Mamba (Selective SSM)

Introduces input-dependent selection parameters. Bypasses standard convolution in favor of a hardware-aware parallel scan, enabling high-quality language modeling.

Linear Attention (e.g., RWKV)

While mathematically related to SSMs, these models reformulate traditional attention mechanisms to avoid the softmax bottleneck, allowing RNN-like recurrent inference.

Representative Flash Papers

yudgctsnk7

Attention Is All You Need

The baseline sequence transduction model that SSMs aim to challenge by eliminating the quadratic scaling bottleneck of self-attention.

y74k5w8ubz

Sequence to Sequence Learning with Neural Networks

The foundational work on recurrent sequence-to-sequence mapping, establishing the classical RNN paradigms that modern SSMs seek to reinvent.

Made with Flash Papers — make your own ↗