A state space model (SSM) in the context of deep learning is a sequence modeling architecture inspired by classical control theory. SSMs map input sequences to output sequences through a continuous latent state, offering linear-time complexity with respect to sequence length. This makes them a compelling alternative to transformers, whose self-attention mechanism scales quadratically.

How It Works

An SSM defines a linear dynamical system with four matrices: A (state transition), B (input projection), C (output projection), and D (skip connection). The continuous system is discretized for processing discrete sequences like text or audio. The key challenge is parameterizing the A matrix so the model can efficiently capture long-range dependencies.

S4 (Structured State Spaces for Sequence Modeling) solved this by initializing A as a HiPPO matrix, which is mathematically designed to compress continuous signals into a fixed-size state. This initialization enables the model to remember information across thousands of time steps. S4 can be computed as either a recurrence (for autoregressive inference) or a convolution (for parallel training), giving it the best of both worlds.

Mamba extended S4 by making the B, C, and discretization parameters input-dependent (selective), allowing the model to filter irrelevant information based on content. This selectivity, combined with a hardware-aware parallel scan implementation, enabled Mamba to match transformer quality on language modeling while scaling linearly with sequence length.

Why It Matters

SSMs address the fundamental computational bottleneck of transformers for long sequences. Where a transformer’s attention mechanism requires O(n^2) computation, SSMs require O(n). This enables processing very long sequences (genomics, audio, extended documents) at substantially lower cost. Hybrid architectures combining SSM layers with attention layers are emerging as a practical middle ground.

Practical Considerations

SSMs are still maturing relative to transformers in terms of ecosystem support and pre-trained model availability. Mamba models are available through Hugging Face, and hybrid architectures like Jamba integrate SSM and attention layers. For teams with long-sequence workloads, SSMs offer meaningful cost savings. Evaluate them against transformer baselines on your specific task, as performance advantages vary by domain.

Sources

  • Gu, A., Goel, K., & Ré, C. (2022). Efficiently modeling long sequences with structured state spaces. ICLR 2022. (S4; foundational structured SSM with HiPPO initialization.)
  • Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv:2312.00752. (Mamba; selective SSM achieving transformer-level language modeling quality at linear complexity.)
  • Kalman, R.E. (1960). A new approach to linear filtering and prediction problems. ASME Journal of Basic Engineering, 82, 35–45. (Classical state space model; theoretical ancestor of deep SSMs.)