In deep learning, state space models (SSMs) are a type of neural network used to process long streams of sequential data, such as text, speech, or sensor readings. Rather than repeatedly re-examining everything that came before, an SSM maintains a compact internal summary (the "state") that is updated as each new piece of the input arrives, an idea adapted from control theory, the branch of engineering concerned with tracking and steering changing systems.
Compact representations make SSMs an efficient alternative to the Transformer architecture that underlies most modern large language models. A Transformer compares every element of its input with every other element, so doubling the length of the input roughly quadruples the work. In contrast, an SSM's effort grows in direct proportion to the input's length. As a result, SSMs can handle very long inputs (e.g., entire books, hours of audio, or DNA sequences) and can run with less computing effort and energy. Therefore, SSMs are well suited to running directly on phones and other small devices rather than in data centers.
SSM-based methods have been adopted in commercial AI systems, including recent large language models such as Kimi K3, and prominent figures in the AI industry have described them as significant: Jensen Huang, chief executive of NVIDIA, called state space models "probably the next generation of Transformers."
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.
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.
The Legendre Memory Unit (LMU), introduced by Voelker, Kajić, and Eliasmith in 2019, was the first successful application 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.
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.
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.
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 Structured initialization which uses the HiPPO matrix for the state matrix A. As well, the model introduced a Normal plus low-rank (NPLR) decomposition, which allows A to be diagonalized stably. Finally, the model reduces the SSM to a Cauchy kernel computation to improve computational efficiency.
S4 achieved interesting results across multiple domains. These include scoring 91% accuracy on sequential CIFAR-10 (matching 2D ResNets with no data augmentation), reporting state-of-the-art performance on all tasks in the Long Range Arena benchmark, being the first model to solve the Path-X task (16,000 sequence length) with 88% accuracy, and demonstrating 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.
Mamba, introduced by Gu and Dao in December 2023, represents a major advancement in SSM architectures through the introduction of selective state space models. Earlier language-focussed SSMs use time-invariant parameters, meaning the matrices A, B, and C remain constant across the sequence. Mamba's main innovation is making these parameters functions of the input, allowing the model to selectively propagate or forget information based on content. In addition, compared to previous LLM work, Mamba provided a simplified architecture that replaces attention and MLP blocks with a unified SSM block. Finally, Mamba includes specific hardware-aware algorithms including parallel scan, kernel fusion, and selective recomputation to achieve efficient training.
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.
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.
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 include:
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.
Recent work has explored various hybrid architectures including: Vision Mamba (Vim), which has bidirectional Mamba blocks for visual data processing; MambaByte, which provides byte-level language modeling without tokenization; and MoE Mamba, which integrates mixture-of-experts with Mamba, requiring 2.35× fewer training steps than standard Mamba.
The continuous-time SSM is defined by linear ordinary differential equations:
x'(t) = Ax(t) + Bu(t) y(t) = Cx(t) + Du(t)
where:
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:
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.
A fundamental advantage of SSMs is their computational complexity compared to Transformers:
Specifically, Transformers have the following complexity profile:
In contrast, State Space Models have the following complexity profile:
This means that 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.
SSMs have demonstrated strong performance on various natural language processing (NLP) tasks, many of which are discussed above in more detail, and include: Specifically SSMs perform language modeling that is competitive with Transformers of similar or larger size, demonstrate long-document understanding with contexts up to 256K tokens, and provide 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.
Vision applications of SSMs have addressed many common areas of visual processing also tackled by Transformers, but with more of a focus on time series data. Briefly, these include image classification on ImageNet, sequential image tasks (e.g., sequential CIFAR-10), and video understanding and generation. Notably, Vision Mamba (Vim) achieved competitive results with Vision Transformers.
SSMs excel at audio tasks due to their continuous-time formulation. Audio applications include speech generation with models like SaShiMi, extremely efficient audio classification on automatic speech recognition benchmarks.
The continuous nature of SSMs makes them well-suited for a wide variety of time series problems. Generally, these include genomic sequence modeling (million-length DNA sequences), climate and weather prediction, and medical time series analysis.
Despite their advantages, pure SSMs face several drawbacks compared to pure Transformers. These include associative recall, training speed, state capacity and discrete modalities.
Associative recall refers to the fact that pure SSM architectures may struggle with tasks requiring precise retrieval of specific information from long contexts, where attention mechanisms excel. As a result, hybrid architectures are often preferred for general deployment.
Training speed captures demonstrated challenges at training short sequences. That is, highly optimized Transformer implementations can be faster than SSMs at sequence lengths below 2,000-4,000 tokens.
State capacity captures the idea that the efficiency gained from fixed-size hidden states may also come with drawbacks if those states saturate, particularly on extremely long sequences, though this can be mitigated with larger state dimensions.
Discrete modalities captures the notion that early, pure SSMs like S4 showed higher perplexity on language tasks compared to Transformers. However, recent architectures like Mamba have largely addressed this.
Overall, these limitations have motivated hybrid architectures that combine SSMs with attention mechanisms to leverage the strengths of both approaches. Often, as in many cases above, these are referred to simply as SSM architectures because they are not pure Transformer architectures.