Draft:State space model

State Space Models (SSMs) are a class of neural network architectures for processing time series data that model sequences using principles from control theory. SSMs have emerged as efficient alternatives to Transformer and recurrent neural network (RNN) architectures, particularly for handling long-range dependencies in sequence modeling tasks. Unlike Transformers which have quadratic complexity with respect to sequence length, SSMs achieve linear or near-linear time complexity, making them particularly effective for processing very long sequences.

Overview

State Space Models in deep learning are based on continuous-time state space representations from classical control theory. At their core, SSMs map a one-dimensional input signal u(t) to an output signal y(t) through a hidden state x(t) using a system of differential equations. The basic SSM is defined by the equations:

x'(t) = Ax(t) + Bu(t)
y(t) = Cx(t) + Du(t)

where A is the state matrix, B is the control matrix, C is the output matrix, and D is a direct feedthrough term (often treated as a skip connection in deep learning applications).

SSMs offer several key advantages: they can naturally handle continuous data, automatically adapt to different sampling rates without retraining, and provide mathematically tractable analysis of their dynamics. Through discretization, SSMs can be viewed from three complementary perspectives: as continuous-time systems, as recurrent networks during inference, and as convolutional models during training.

History and Development

Origins in Neuroscience

The application of state space models to deep learning traces back to theoretical neuroscience research. In 2018, Aaron R. Voelker and Chris Eliasmith from the University of Waterloo proposed that the dynamic system in SSMs can effectively model "time cells" present in the hippocampus and cortex, leading to their work on applying SSMs to neural networks.

Legendre Memory Units (2019)

The Legendre Memory Unit (LMU), introduced by Voelker, Kajić, and Eliasmith in 2019, was among the first successful applications of SSMs in deep learning. LMUs are mathematically derived to orthogonalize continuous-time history by solving coupled ordinary differential equations, with their phase space mapping onto sliding windows of time via Legendre polynomials. LMUs demonstrated the ability to handle temporal dependencies spanning 100,000 time steps and achieved state-of-the-art performance on permuted sequential MNIST, exceeding 97% accuracy.

HiPPO Framework (2020)

The High-Order Polynomial Projection Operators (HiPPO) framework, introduced by Gu et al. in 2020 who were from Stanford University, provided a unified mathematical foundation for memory in sequence models. HiPPO optimally projects continuous signals onto polynomial bases, yielding linear dynamics for the projection coefficients. This framework produces several instantiations including HiPPO-LegS (scaled Legendre) and HiPPO-LegT (translated Legendre), which achieve timescale robustness and bounded gradients. The HiPPO framework achieved 98.3% accuracy on permuted MNIST, surpassing previous RNN approaches by a significant margin.

Parallelization (2021)

Chilkuri and Eliasmith proposed and demonstrated a method to efficiently train SSMs in parallel on GPUs. This overcomes concerns that the recurrence in SSMs would be difficult to train on GPUs, since other recurrent networks like LSTMs fell out of favour for this reason. Subsequently the first large language model (LLM) using SSMs was demonstrated to scale better than either LSTMs or Transformers using this method.

Structured State Space Models (S4, 2021)

The Structured State Space sequence model (S4), introduced by Gu, Goel, and Ré in 2021, marked a breakthrough in making SSMs practical for large-scale deep learning. S4 addressed the computational challenges of naive SSM implementations through a novel parameterization involving:

1. Structured initialization: Using the HiPPO matrix for the state matrix A
2. Normal plus low-rank (NPLR) decomposition: Allowing A to be diagonalized stably
3. Efficient computation: Reducing the SSM to a Cauchy kernel computation

S4 achieved remarkable results across multiple domains:

• 91% accuracy on sequential CIFAR-10 (matching 2D ResNets with no data augmentation)
• State-of-the-art on all tasks in the Long Range Arena benchmark
• First model to solve the Path-X task (16,000 sequence length) with 88% accuracy
• 60× faster generation than Transformers on language modeling

The model demonstrated the ability to handle sequences exceeding 10,000 steps while maintaining linear scaling in sequence length.

Key Architectural Innovations

Mamba (2023)

Mamba, introduced by Gu and Dao in December 2023, represents a major advancement in SSM architectures through the introduction of selective state space models. Traditional SSMs use time-invariant parameters, meaning the matrices A, B, and C remain constant across the sequence. Mamba's key innovation is making these parameters functions of the input, allowing the model to selectively propagate or forget information based on content.

Key features of Mamba:

• Selective SSMs (S6): Input-dependent parameters (B, C, and step size Δ) that enable content-based reasoning
• Hardware-aware algorithms: Parallel scan, kernel fusion, and selective recomputation to achieve efficient training
• Linear-time complexity: O(n) scaling with sequence length, compared to O(n²) for Transformers
• Simplified architecture: Replaces attention and MLP blocks with a unified SSM block

Mamba achieved competitive or superior performance compared to Transformers while providing 5× higher throughput and linear scaling to million-length sequences. On language modeling, Mamba-3B matched Transformers twice its size in both pretraining and downstream evaluation.

Mamba-2 (2024)

In May 2024, Dao and Gu introduced Mamba-2 through their "Transformers are SSMs" paper, which established theoretical connections between SSMs and attention mechanisms via structured semiseparable matrices. The State Space Duality (SSD) framework enabled the design of Mamba-2, which is 2-8× faster than Mamba while maintaining competitive performance with Transformers on language modeling.

Mamba-2 achieves faster computation by leveraging matrix multiplication primitives and tensor cores on modern GPUs, allowing for larger state expansion (typically N=128-256 compared to N=16 in Mamba) while remaining computationally efficient. The model also enables better system-level optimizations including tensor parallelism and sequence parallelism.

Hybrid Architectures

Jamba (2024)

AI21 Labs introduced Jamba in March 2024, the first production-grade model combining Mamba SSM layers with Transformer attention and mixture-of-experts (MoE) components. Jamba features:

• A hybrid architecture interleaving Transformer and Mamba layers at a 1:7 ratio
• 52B total parameters with only 12B active parameters during inference
• Support for 256K token context windows
• 3× throughput improvement on long contexts compared to Mixtral 8x7B

The architecture demonstrated that hybrid approaches can effectively balance the strengths of both SSMs (efficiency, long context) and Transformers (performance, in-context learning). Jamba 1.5, released in August 2024, scaled to 398B total parameters with 94B active, representing the largest hybrid SSM-Transformer architecture to date.

Other Hybrid Models

Recent work has explored various hybrid architectures:

• Vision Mamba (Vim): Bidirectional Mamba blocks for visual data processing
• MambaByte: Byte-level language modeling without tokenization
• MoE Mamba: Integration of mixture-of-experts with Mamba, requiring 2.2× fewer training steps than standard Mamba

Mathematical Framework

Continuous Representation

The continuous-time SSM is defined by linear ordinary differential equations:

x'(t) = Ax(t) + Bu(t)
y(t) = Cx(t) + Du(t)

where:

• x(t) ∈ ℝᴺ is the state vector (N-dimensional latent state)
• u(t) ∈ ℝ is the input signal
• y(t) ∈ ℝ is the output signal
• A ∈ ℝᴺˣᴺ is the state transition or dynamics matrix
• B ∈ ℝᴺˣ¹ is the control matrix
• C ∈ ℝ¹ˣᴺ is the output matrix
• D ∈ ℝ is the feedthrough term

Discretization

To implement SSMs on digital computers, the continuous system must be discretized. The most common approach uses the zero-order hold (ZOH) method with step size Δ:

x̄ₖ = Āxₖ₋₁ + B̄uₖ
yₖ = Cx̄ₖ + Duₖ

where the discrete parameters are:

Ā = exp(ΔA)
B̄ = (exp(ΔA) - I)A⁻¹B

This discretization highlights two complementary views:

1. Recurrent view: Efficient O(N) linear-time inference by maintaining state
2. Convolutional view: Parallel O(N log N) training via FFT-based convolutions

The convolution kernel K̄ can be precomputed as:

K̄ₖ = CĀᵏB̄

This duality allows SSMs to combine the inference efficiency of RNNs with the training parallelism of CNNs and Transformers.

Computational Complexity and Efficiency

Comparison with Transformers

A fundamental advantage of SSMs is their computational complexity compared to Transformers:

Transformers:

• Training complexity: O(N²D) where N is sequence length, D is dimension
• Inference complexity: O(N²D) due to attention over all previous tokens
• Memory: O(N) for KV cache, growing linearly with sequence length

State Space Models:

• Training complexity: O(N log N) via FFT for convolutional view
• Inference complexity: O(N) linear time per token (only update hidden state)
• Memory: O(1) constant for state, independent of sequence length

This means SSMs scale better during training and achieve linear-time generation, while Transformers have quadratic complexity that becomes prohibitive for very long sequences. At sequence lengths beyond 8,000-16,000 tokens, SSMs typically become significantly faster than Transformers.

Hardware Optimization

Modern SSM implementations employ several hardware-aware optimizations:

1. Kernel fusion: Combining multiple operations to minimize memory transfers between GPU memory hierarchies
2. Parallel scan algorithms: Efficient parallelization of recurrent computation during training
3. Selective recomputation: Recomputing intermediate values during backpropagation rather than storing them
4. Tensor parallelism: Distributing computation across multiple GPUs

These optimizations have been crucial for achieving competitive wall-clock time performance with highly optimized Transformer implementations.

Applications

Natural Language Processing

SSMs have demonstrated strong performance on various NLP tasks:

• Language modeling competitive with Transformers of similar or larger size
• Long-document understanding with contexts up to 256K tokens
• Strong performance on question answering, summarization, and text classification

The selective mechanism in Mamba has proven particularly effective for discrete modalities like language, addressing early limitations of S4 in this domain.

Computer Vision

Vision applications of SSMs include:

• Image classification on ImageNet
• Sequential image tasks (e.g., sequential CIFAR-10)
• Video understanding and generation
• Vision Mamba (Vim) achieving competitive results with Vision Transformers

Audio and Speech

SSMs excel at audio tasks due to their continuous-time formulation:

• Speech generation with models like SaShiMi
• Audio classification on benchmarks like Speech Commands (96.3% accuracy with S4)
• Music generation
• Robustness to different sampling rates without retraining

Time Series and Scientific Computing

The continuous nature of SSMs makes them well-suited for:

• Time series forecasting
• Genomic sequence modeling (million-length DNA sequences)
• Climate and weather prediction
• Medical time series analysis
• Financial modeling

Limitations and Open Challenges

Despite their advantages, SSMs face several challenges:

1. Associative recall: Pure SSM architectures may struggle with tasks requiring precise retrieval of specific information from long contexts, where attention mechanisms excel
2. Training speed at short sequences: Highly optimized Transformer implementations can be faster than SSMs at sequence lengths below 2,000-4,000 tokens
3. State capacity: Fixed-size hidden states may saturate on extremely long sequences, though this can be mitigated with larger state dimensions
4. Discrete modalities: While Mamba addressed this, earlier SSMs like S4 showed higher perplexity on language tasks compared to Transformers

These limitations have motivated hybrid architectures that combine SSMs with attention mechanisms to leverage the strengths of both approaches.

See Also

• Recurrent neural network
• Transformer (machine learning model)
• Attention (machine learning)
• Long short-term memory
• Control theory
• Mamba (deep learning architecture)

External Links

• Mamba official GitHub repository
• S4 official GitHub repository
• The Annotated S4 tutorial
• HiPPO official code
• AI21 Jamba models