Selective State Space Models (Mamba): Mathematical Foundations, Discretization Dynamics, and Linear-Time Sequence Modeling

The dominant paradigm in natural language processing and modern foundation models relies almost exclusively on the Transformer architecture. While standard multi-head self-attention delivers strong expressivity and in-context learning capabilities, its computational requirements present fundamental scaling bottlenecks: training complexity scales quadratically with sequence length $O(L^2)$ , and autoregressive token generation requires storing key-value pairs in high-bandwidth memory (HBM), creating a linear KV cache memory footprint $O(L)$ per active sequence.

To circumvent these computational limits, researchers developed Structured State Space Sequence Models (SSMs), culminating in architectures such as S4, Mamba, and Mamba-2. State space models originate from classical control theory, mapping a continuous input signal to an output signal through an implicit latent state. By introducing input-dependent parameter selection and hardware-aware associative scan kernels, modern selective SSMs achieve linear time complexity $O(L)$ during training, constant time complexity $O(1)$ per generated token during inference, and competitive performance against standard Transformers across dense language modeling benchmarks.

Continuous-Time State Space Representations

At their core, continuous-time state space models define a linear mapping from a 1D continuous sequence $x(t) \in \mathbb{R}$ to an output $y(t) \in \mathbb{R}$ via an $N$ -dimensional latent state $h(t) \in \mathbb{R}^N$ . The continuous system is governed by a pair of linear differential equations:

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

In this formulation:

$A \in \mathbb{R}^{N \times N}$ represents the state transition matrix, governing how internal latent representations evolve over continuous time.
$B \in \mathbb{R}^{N \times 1}$ is the input projection vector, modulating the continuous input signal $x(t)$ into the latent state space.
$C \in \mathbb{R}^{1 \times N}$ is the output projection vector, mapping the latent state $h(t)$ back to the scalar output space.
$D \in \mathbb{R}^{1 \times 1}$ represents a direct feedthrough connection (often omitted or parameterized as a skip connection).

In deep learning implementations, models operate over multi-dimensional feature vectors with hidden dimension $D$ . The state space transformation is applied independently across each feature channel, resulting in a latent state dimension of $D \times N$ .

Discretization Mechanics: Bridging Continuous Systems to Discrete Sequences

Because digital language models process discrete sequences of tokens rather than continuous analog signals, the continuous differential system must be discretized using a timescale step size parameter $\Delta > 0$ .

The standard discretization method in modern SSMs is the Zero-Order Hold (ZOH) transformation, which assumes that the continuous input signal remains constant over each discrete interval $[t_k, t_k + \Delta]$ . Applying ZOH yields the discrete-time state space formulation:

h_t = A_bar * h_{t-1} + B_bar * x_t
y_t = C_bar * h_t + D * x_t

The discrete transition matrices $\bar{A}$ and $\bar{B}$ are derived mathematically from continuous parameters $A$ , $B$ , and timescale $\Delta$ :

A_bar = exp(Delta * A)
B_bar = (Delta * A)^(-1) * (exp(Delta * A) - I) * (Delta * B)

In first-order Taylor approximations, $\bar{B}$ can be simplified as $\bar{B} \approx \Delta B$ . The output matrix remains invariant during discretization: $\bar{C} = C$ .

Alternative discretization strategies include the Bilinear (Tustin) transform and Euler discretization:

Bilinear Transform: $\bar{A} = (I - \frac{\Delta}{2} A)^{-1} (I + \frac{\Delta}{2} A)$ , preserving frequency-domain properties and stability along the imaginary axis.
Euler Discretization: $\bar{A} = I + \Delta A$ , which offers simpler arithmetic but suffers from numerical instability when $\Delta \|A\|$ is large.

The timescale parameter $\Delta$ acts as a dynamic gating mechanism:

When $\Delta \to \infty$ , $\bar{A} \to 0$ and $\bar{B} \to B$ , causing the system to reset its latent state and focus entirely on the current input token $x_t$ .
When $\Delta \to 0$ , $\bar{A} \to I$ and $\bar{B} \to 0$ , causing the model to preserve its existing memory state $h_{t-1}$ while ignoring current input perturbations.

Selective State Space Models Architecture

From Linear Time-Invariant (LTI) Systems to Selective SSMs

Early structured SSMs, including S4 (Structured State Space for Sequences) and H3 (Hungry Hungry Hippos), enforced Linear Time-Invariance (LTI). In an LTI system, parameters $(A, B, C, \Delta)$ remain constant across all time steps $t \in [1, L]$ .

LTI systems possess a dual representation:

Linear Recurrence for autoregressive inference: $h_t = \bar{A} h_{t-1} + \bar{B} x_t$ , requiring $O(1)$ compute per token.
Global Convolution for parallel training: Expanding the unrolled recurrence reveals that $y$ is computed as a 1D convolution of the entire sequence $x$ with an explicit SSM kernel $\bar{K}$ :

K_bar = (C * B_bar, C * A_bar * B_bar, C * A_bar^2 * B_bar, ..., C * A_bar^{L-1} * B_bar)
y = x * K_bar

Using the Fast Fourier Transform (FFT), this convolution can be computed across sequence length $L$ in $O(L \log L)$ time.

The Fundamental Limitation of LTI Models

Despite their training efficiency, LTI models cannot solve foundational sequence processing primitives required for natural language reasoning:

Selective Copying: Remembering relevant tokens while filtering out irrelevant filler tokens across variable sequence intervals. Because the convolution kernel $\bar{K}$ is static and input-independent, it allocates identical weights to tokens regardless of semantic content.
Induction Heads: Detecting repeated patterns $(A \to B \dots A \to ?)$ and retrieving the associated follower token. LTI systems cannot dynamically reallocate attention or modify state transitions conditioned on observed context.

The Mamba Selection Mechanism

Mamba resolves the LTI limitation by making the parameters $\Delta$ , $B$ , and $C$ explicit functions of the input token $x_t$ at each time step:

B_t = Linear_N(x_t)
C_t = Linear_N(x_t)
Delta_t = Softplus(Linear_D(x_t) + Parameter_Delta)
A_bar_t = exp(Delta_t * A)
B_bar_t = (Delta_t * A)^(-1) * (exp(Delta_t * A) - I) * (Delta_t * B_t)

By parameterizing $B_t$ , $C_t$ , and $\Delta_t$ dynamically based on input embeddings, Mamba acts as a content-aware filter. The model dynamically flushes obsolete context from the state or amplifies salient tokens for indefinite retention.

Hardware-Aware Parallel Scan and Memory Hierarchy Optimization

Making parameters input-dependent breaks the convolutional representation because the convolution kernel $\bar{K}$ can no longer be precomputed. Evaluating the recurrence sequentially across sequence length $L$ on modern GPU accelerators would severely degrade training throughput due to memory bandwidth constraints.

Standard deep learning frameworks materialize intermediate tensors of shape $(B, L, D, N)$ in High-Bandwidth Memory (HBM). For batch size $B=16$ , sequence length $L=4096$ , model dimension $D=2048$ , and state dimension $N=16$ , intermediate tensors require tens of gigabytes of HBM read/write traffic per layer.

To eliminate memory bottlenecks, Mamba implements a hardware-aware parallel scan kernel using GPU SRAM:

Kernel Fusion in SRAM: The model loads input projections, $A$ , $B_t$ , $C_t$ , and $\Delta_t$ directly into high-speed on-chip SRAM cache (operating at over 15 TB/s on modern architectures).
Parallel Prefix Scan: The sequential recurrence $h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t$ is computed using a parallel associative scan (Blelloch algorithm) with an associative operator $\circ$ :

(A_i, B_i * x_i) o (A_{i-1}, B_{i-1} * x_{i-1}) = (A_i * A_{i-1}, A_i * B_{i-1} * x_{i-1} + B_i * x_i)

Work-Efficient Tree Reduction: The parallel scan executes in $O(L)$ total work and $O(\log L)$ parallel depth across thread blocks.
Intermediate Gradient Recomputation: During the backward pass, intermediate latent states $h_t$ are not saved to HBM. Instead, they are recomputed on-the-fly in SRAM when evaluating gradients, reducing peak memory usage to the level of standard Transformer layers.

Structured State Space Duality (SSD) and Mamba-2

In 2024, researchers introduced Structured State Space Duality (SSD), establishing a direct mathematical equivalence between selective state space models and structured linear attention mechanisms.

SSD demonstrates that when state transition matrix $A_t$ is structured as a scalar-times-identity matrix ( $A_t = a_t I$ ), the SSM transformation is equivalent to a 1-semiseparable matrix multiplication:

Y = (L o (Q * K^T)) * V

Where:

$Q = C$ , $K = B$ , and $V = X$ .
$L_{ij} = \prod_{k=j+1}^i a_k$ represents a strictly causal decaying mask matrix.
$\circ$ denotes Hadamard (element-wise) multiplication.

This duality enables Mamba-2 to formulate state space computation through chunked matrix multiplications. A sequence of length $L$ is partitioned into non-overlapping chunks of size $Q$ (typically 64 or 128 tokens):

Within each chunk, computations are executed as dense matrix multiplies on Tensor Cores ( $Q K^T V$ ).
Across chunk boundaries, state representations are propagated using a recurrence over chunk-level summary states.

By reformulating state space updates into matrix multiplies, Mamba-2 achieves 2x to 8x higher training throughput compared to Mamba-1 while supporting larger state expansion dimensions ( $N=64$ or $N=128$ ).

Architectural Comparison: Mamba vs Multi-Head Attention

The differences in computational complexity, memory scaling, and serving properties between standard Transformer attention, Linear Attention, and Selective SSMs can be summarized across core dimensions:

Computational Complexity

Transformer Multi-Head Attention: $O(L^2 \cdot D)$ training compute; $O(L \cdot D)$ per step inference compute.
Linear Attention: $O(L \cdot D^2)$ training compute; $O(D^2)$ per step inference compute.
Selective SSM (Mamba-1 / Mamba-2): $O(L \cdot D \cdot N)$ training compute; $O(D \cdot N)$ per step inference compute.

State Memory Footprint During Generation

Transformer Multi-Head Attention: $O(B \cdot L \cdot D)$ dynamic KV cache in HBM; scales linearly with generated sequence length.
Linear Attention: $O(B \cdot D^2)$ fixed-size recurrent state; constant memory.
Selective SSM: $O(B \cdot D \cdot N)$ fixed-size latent state $h_t$ ; constant memory regardless of sequence length.

Induction and Associative Recall

Transformer Multi-Head Attention: Exact retrieval via explicit pair-wise query-key dot products.
Linear Attention: Degraded retrieval due to unnormalized feature map decay and kernel approximations.
Selective SSM: High retrieval accuracy via input-dependent gating ( $\Delta_t, B_t, C_t$ ) and expanded state representations.

Serving Economics and Latency Dynamics

In production deployment environments, the memory and computational profiles of Mamba provide several practical advantages:

Elimination of Memory-Bound Decoding: Autoregressive decoding in large Transformers is typically memory-bandwidth bound because reading the KV cache requires transferring gigabytes of memory per token generated. Mamba maintains a small, fixed-size hidden state per layer, allowing decoding to run compute-bound at substantially higher batch sizes.
Linear Long-Context Scaling: For document processing, codebase ingestion, and multi-turn agent sessions exceeding 100,000 tokens, Mamba processes prompt context with linear time and memory scaling, avoiding the out-of-memory errors and quadratic slowdowns inherent to standard attention.
Hybrid Architectures: Modern production models such as Jamba and Bamba integrate interleaved layers of Mamba SSM blocks and Multi-Head Attention blocks (e.g. 1 attention layer for every 7 Mamba layers). This hybrid topology preserves exact associative recall across complex multi-document reasoning tasks while reducing total KV cache size by over 80 percent.

Sources

Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752. https://arxiv.org/abs/2312.00752
Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. ICML 2024. arXiv:2405.21060. https://arxiv.org/abs/2405.21060
Gu, A., Goel, K., & Re, C. (2021). Efficiently Modeling Long Sequences with Structured State Spaces. ICLR 2022. arXiv:2111.00396. https://arxiv.org/abs/2111.00396
Lieber, O., et al. (2024). Jamba: A Hybrid Transformer-Mamba Language Model. arXiv:2403.19887. https://arxiv.org/abs/2403.19887