State space model (Mamba)

What it is

A state space model in deep learning represents a sequence as a hidden state that evolves step by step with the input, rather than comparing every position with every other position as attention does^[1]. It is a continuous-time model with both recurrent and convolutional properties; it can handle irregularly sampled data and keeps an unbounded context^[1].

How it works

The model keeps a state, updates it once per input, and produces an output at each step. Because the state is carried forward, the model does not need to store every earlier position. Reading a sentence, for example, it updates its internal state at each token, and the final state serves as a summary of everything read so far.

Gu, Goel and Re proposed the structured state space sequence model, S4, in 2021, building on the state space form from control theory^[1]. S4 was published at NeurIPS 2021 as a linear state space layer combining recurrent, convolutional and continuous-time models^[1]. The linear structure lets the same update rule be run step by step or unrolled into a convolution and computed in parallel.

Mamba

Mamba was proposed by Albert Gu and Tri Dao in 2023^[1]. The paper notes that foundation models are almost universally based on the Transformer architecture and its core attention module^[2], and it identifies the inability to perform content-based reasoning as a key weakness of subquadratic architectures such as linear attention, gated convolution and structured state space models^[2].

Mamba's key mechanism is a selective state space, in which the model's parameters change with the input, turning a time-invariant system into a time-varying one that can filter out irrelevant information^[2]. Because the parameters are functions of the input, the model can selectively propagate or forget information along the sequence depending on the current token^[2].

The change has a cost: once parameters vary with the input, efficient convolutions are ruled out^[2], and the time-varying computation creates efficiency challenges^[1]. The authors therefore design a hardware-aware parallel algorithm in recurrent mode^[2], using kernel fusion, parallel scan and recomputation to run efficiently on GPUs^[1].

Mamba also merges the state space model and the MLP block into a single homogeneous architecture^[1], and it is designed for long sequences with linear time complexity, as the paper title "Linear-Time Sequence Modeling" indicates^[1]. It is claimed to train and infer more efficiently than transformers on language, audio and genomics tasks^[1].

Later work

A later variant, MoE-Mamba in 2024, reached comparable performance with 2.2 times fewer training steps^[1].

Limitations

The selective mechanism answers the content-based reasoning weakness identified in earlier subquadratic architectures^[2], but it makes the computation time-varying, so efficiency depends on the dedicated hardware-aware algorithm^[1]. The comparison with transformers is the authors' claim^[1]; what a given task actually gains depends on the task and the implementation.