State Space Models and Mamba (Mamba-1 and Mamba-2): Mathematical Foundations, Selective State Spaces, Structured State Space Duality (SSD), and Linear-Time Sequence Modeling

State Space Models (SSMs) and their modern selective formulations, most notably Mamba-1 and Mamba-2, represent a foundational alternative to the standard Transformer architecture for sequence modeling. While multi-head self-attention scales quadratically with sequence length ($O(T^2)$) and requires an ever-expanding Key-Value (KV) cache during autoregressive generation ($O(T)$), State Space Models achieve linear time complexity ($O(T)$) during training and constant memory footprint ($O(1)$) per step during inference.

Originating from continuous-time dynamical systems in classical control theory, SSMs were historically limited by time-invariance, preventing them from performing content-dependent reasoning. The development of Selective State Spaces (Gu and Dao, 2023) and Structured State Space Duality (Dao and Gu, 2024) bridged this gap, enabling subquadratic models to match or exceed Transformer performance across language, audio, and genomic modeling while unlocking up to 5x higher inference throughput.

Theoretical Foundations: Continuous and Discrete State Spaces

A continuous-time State Space Model maps a 1-dimensional continuous input signal $x(t) \in \mathbb{R}$ through an $N$-dimensional latent state $h(t) \in \mathbb{R}^N$ to produce an output signal $y(t) \in \mathbb{R}$. The system is governed by a set of linear ordinary differential equations (ODEs):

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

where:

• $A \in \mathbb{R}^{N \times N}$ is the state transition matrix, determining how the internal state evolves over time.
• $B \in \mathbb{R}^{N \times 1}$ is the input projection matrix, mapping the scalar input into the latent state space.
• $C \in \mathbb{R}^{1 \times N}$ is the output projection matrix, mapping the latent state back to the output space.
• $D \in \mathbb{R}$ is a direct feedthrough (skip connection) parameter, typically omitted or handled as a separate residual connection.

Continuous-to-Discrete Discretization

Because digital computers process discrete token sequences $(x_0, x_1, \dots, x_{T-1})$ rather than continuous signals, the continuous system must be discretized over a continuous timescale parameter $\Delta > 0$ (the step size or sample resolution).

Using the standard Zero-Order Hold (ZOH) discretization rule, which assumes the input signal $x(t)$ remains constant over each time interval $[t\Delta, (t+1)\Delta)$, the continuous parameters $(A, B)$ are transformed into discrete parameters $(\bar{A}, \bar{B})$:

$$\bar{A} = \exp(\Delta A)$$ $$\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot (\Delta B) \approx \Delta B$$

The discrete state space formulation is then expressed as:

$$h_t = \bar{A} h_{t-1} + \bar{B} x_t$$ $$y_t = C h_t + D x_t$$

The Duality of Inference and Training in Linear Time-Invariant SSMs

When the parameters $(\Delta, A, B, C)$ are constant over time (known as a Linear Time-Invariant or LTI system), the model possesses a dual representation:

1. Linear Recurrence (Inference Mode): Tokens are processed sequentially step-by-step. Generating token $t$ requires only updating the fixed $N$-dimensional vector $h_t$ from $h_{t-1}$ in $O(1)$ memory and $O(N)$ operations per step, completely eliminating the linear memory growth of a Transformer KV cache.
2. Global Convolution (Training Mode): Expanding the recurrence from $h_0 = 0$ yields:

$$y_t = C \bar{A}^t \bar{B} x_0 + C \bar{A}^{t-1} \bar{B} x_1 + \dots + C \bar{A} \bar{B} x_{t-1} + C \bar{B} x_t$$

The entire output sequence $y \in \mathbb{R}^T$ can be computed as a discrete 1D convolution:

$$\bar{K} = \left( C\bar{B}, C\bar{A}\bar{B}, C\bar{A}^2\bar{B}, \dots, C\bar{A}^{T-1}\bar{B} \right) \in \mathbb{R}^T$$ $$y = x * \bar{K}$$

Using the Fast Fourier Transform (FFT), this global convolution computes the entire sequence in parallel in $O(T \log T)$ time, enabling efficient training on modern GPUs.

The LTI Expressivity Bottleneck and HiPPO

Early structured SSM architectures such as S4 (Gu et al., 2021) addressed numerical stability and long-range memory by structuring $A$ using the HiPPO (High-order Polynomial Projection Operator) framework. HiPPO derives closed-form continuous transition matrices that maintain an optimal continuous orthogonal polynomial projection (e.g., Legendre polynomials) of historical input signals.

However, LTI SSMs suffered from a fundamental theoretical limitation: time-invariance. Because the convolutional filter $\bar{K}$ is static and independent of input tokens, an LTI model cannot dynamically change its attention or memory based on sequence context. It cannot solve basic associative recall tasks, such as copying dynamic key-value pairs or filtering out irrelevant filler tokens based on content.

Mamba-1: Selective State Spaces and Hardware-Aware Scans

To overcome the expressivity ceiling of LTI models while retaining linear computational scaling, Mamba-1 (Gu and Dao, 2023) introduced Selective State Spaces (S6).

The Selection Mechanism

Mamba makes the discretization parameters dynamic functions of the current input token $x_t \in \mathbb{R}^D$:

$$B_t = W_B x_t \in \mathbb{R}^{N}$$ $$C_t = W_C x_t \in \mathbb{R}^{N}$$ $$\Delta_t = \text{softplus}\left(\text{Parameter} + W_\Delta x_t\right) \in \mathbb{R}^D$$

where $W_B \in \mathbb{R}^{N \times D}$, $W_C \in \mathbb{R}^{N \times D}$, and $W_\Delta \in \mathbb{R}^{D \times D}$. The transition matrix $A \in \mathbb{R}^{D \times N}$ is parameterized as a diagonal matrix. The discretized parameters become time-dependent:

$$\bar{A}_t = \exp(\Delta_t A), \quad \bar{B}_t = \Delta_t B_t$$

The dynamic step size $\Delta_t$ acts as a continuous, input-dependent gating mechanism:

• When $\Delta_t \to 0$, $\bar{A}_t \to I$ and $\bar{B}_t \to 0$, causing the model to ignore the current token $x_t$ and retain its accumulated historical state $h_{t-1}$.
• When $\Delta_t \to \infty$, $\bar{A}_t \to 0$, resetting the historical state and focusing computational capacity entirely on the new token $x_t$.

+-------------------------------------------------------------------------+
|                           Mamba-1 Block Layer                           |
|                                                                         |
|                Input Vector x_t in R^D                                  |
|                 /                    \                                  |
|        Linear Expansion (E=2)     Linear Projection                     |
|                |                           |                            |
|       1D Conv (k=4) + SiLU                 |                            |
|                |                           |                            |
|     Selective SSM (S6 Engine)              |                            |
|   Delta_t, B_t, C_t = f(x_t)               |                            |
|     h_t = A_t h_{t-1} + B_t x_t            |                            |
|     y_t = C_t h_t                          |                            |
|                \                          /                             |
|                 Multiplicative Gating (SiLU)                            |
|                              |                                          |
|                    Linear Projection Out                                |
+-------------------------------------------------------------------------+

Breaking the Convolutional Constraint

Because $\bar{A}_t, \bar{B}_t, C_t$ vary with each token $t$, the system loses time-invariance. The global convolution operator $\bar{K}$ can no longer be parameterized as a fixed filter, making FFT-based parallel convolution mathematically impossible.

A naive sequential implementation of the recurrence $h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t$ on a GPU would require materializing the intermediate state tensor $h \in \mathbb{R}^{B \times L \times D \times N}$ in High-Bandwidth Memory (HBM) at each token position. For batch size $B=64$, sequence length $L=2048$, model dimension $D=1024$, and state expansion $N=16$, materializing $h$ requires transferring gigabytes of intermediate data across the memory bus at every layer, resulting in severe GPU memory bandwidth starvation.

Hardware-Aware Parallel Associative Scan

Mamba-1 resolves this memory bottleneck through a dedicated GPU kernel utilizing a parallel associative scan:

1. Kernel Fusion in SRAM: Parameters $(\Delta, A, B, C)$ are loaded from GPU HBM into fast local on-chip Shared Memory (SRAM) once per block.
2. On-Chip Discretization: The discretization step $(\bar{A}_t, \bar{B}_t)$ is computed directly in SRAM without writing intermediate discretized matrices back to HBM.
3. Associative Scan Tree: The recurrence is formulated as an associative binary operation:

$$(a_j, b_j) \bullet (a_i, b_i) = (a_j \cdot a_i, a_j \cdot b_i + b_j)$$

Using a parallel prefix sum (Blelloch scan tree) executed across GPU thread warps, the sequential recurrence over length $T$ is computed in parallel across the sequence dimension in $O(\log T)$ parallel step depth.

1. Direct Output Projection: The state $h_t$ is multiplied by $C_t$ inside SRAM, and only the final output tensor $y \in \mathbb{R}^{B \times L \times D}$ is written back to HBM. Intermediate states are recomputed during the backward pass rather than saved to memory.

Mamba-2 and Structured State Space Duality (SSD)

While Mamba-1 established competitive empirical scaling against Transformers, its execution model exposed a hardware utilization limitation on modern GPU architectures (such as NVIDIA A100 and H100). Modern AI accelerators achieve high theoretical FLOPS primarily through specialized matrix-multiplication units (Tensor Cores).

Mamba-1’s selective scan operates predominantly on matrix-vector products and point-wise state updates. These operations are memory-bandwidth bound and cannot fully exploit Tensor Core Matrix Multiply-Accumulate (MMA) pipelines.

Mamba-2 (Dao and Gu, 2024) resolved this hardware mismatch by introducing the framework of Structured State Space Duality (SSD).

The SSD Theoretical Equivalence

Dao and Gu proved that when the continuous transition matrix $A_t$ is structured as a scalar-times-identity matrix ($A_t = a_t I$), linear time-varying state space models are mathematically equivalent to a structured form of masked linear attention.

Consider an SSM with input $x \in \mathbb{R}^T$, scalar transitions $a \in \mathbb{R}^T$, and projection vectors $B, C \in \mathbb{R}^{T \times P}$. The input-to-output mapping $Y \in \mathbb{R}^{T \times P}$ can be expressed in closed matrix form:

$$Y = M X$$

where the transformation matrix $M \in \mathbb{R}^{T \times T}$ is a 1-semiseparable matrix whose elements are defined by:

$$M_{i,j} = \begin{cases} C_i^T \left( \prod_{k=j+1}^i a_k \right) B_j & \text{if } i \ge j \\ 0 & \text{if } i < j \end{cases}$$

By factoring $M$ through the lens of attention operators, this transformation can be written as:

$$Y = \left( L \circ \left( Q K^T \right) \right) V$$

where:

• Queries: $Q = C \in \mathbb{R}^{T \times P}$
• Keys: $K = B \in \mathbb{R}^{T \times P}$
• Values: $V = X \in \mathbb{R}^{T \times P}$
• Mask matrix: $L_{i,j} = \prod_{k=j+1}^i a_k$ for $i \ge j$, representing a 1-semiseparable causal decay mask.
• $\circ$ denotes the Hadamard (element-wise) product.

This formulation demonstrates that SSMs and Attention are two computational expressions of the same underlying algebraic structure:

• Attention view: Compute $Q K^T$ (matrix multiplication), multiply by mask $L$, and multiply by $V$.
• State space view: Maintain hidden state $h_t = a_t h_{t-1} + K_t^T V_t$ and project $y_t = Q_t h_t$.

+-------------------------------------------------------------------------+
|                  Chunk-Wise SSD Computation Pipeline                    |
|                                                                         |
| Sequence divided into chunks of length Q (e.g., Q = 64 tokens)          |
|                                                                         |
| 1. Intra-Chunk Computation (Tensor Cores / Dense GEMM):                 |
|    Y_intra = (L_chunk o (Q_chunk * K_chunk^T)) * V_chunk                |
|                                                                         |
| 2. Inter-Chunk State Recurrence (Linear SSM Scan):                      |
|    H_k = a_k * H_{k-1} + K_k^T * V_k  (State size N x P)                |
|                                                                         |
| 3. Output Aggregation:                                                  |
|    Y_chunk = Y_intra + Q_chunk * H_{k-1}                                |
+-------------------------------------------------------------------------+

The Chunk-Wise SSD Algorithm

To maximize hardware throughput on Tensor Cores, Mamba-2 implements a chunk-wise decomposition algorithm:

1. The sequence of length $T$ is divided into non-overlapping chunks of size $Q$ (typically $Q = 64$ or $Q = 128$).
2. Diagonal Blocks (Intra-Chunk): Within each chunk, tokens interact via dense matrix multiplications using Tensor Cores: $Y_{intra} = \left( L_{chunk} \circ (Q_{chunk} K_{chunk}^T) \right) V_{chunk}$.
3. Off-Diagonal Blocks (Inter-Chunk): Interaction between different chunks occurs strictly through boundary state vectors. Each chunk summarizes its token representations into an end-of-chunk state $H_k \in \mathbb{R}^{N \times P}$.
4. A recurrent SSM scan propagates the states across chunks: $H_k = a_k H_{k-1} + K_k^T V_k$.
5. The accumulated state is projected into the current chunk via matrix multiplication ($Q_{chunk} H_{k-1}$) and added to $Y_{intra}$.

This chunk-wise algorithm achieves up to 8x faster training speed over Mamba-1 by mapping the vast majority of FLOPs to Tensor Core GEMM operations, while allowing the state dimension $N$ to be expanded from $N=16$ (in Mamba-1) to $N=64, 128$, or $256$ without incurring memory bandwidth penalties.

Multi-Head State Space Layers

Mamba-2 restructures the block architecture to match standard Multi-Head Attention (MHA) and Grouped-Query Attention (GQA) conventions. In Mamba-2:

• Hidden dimension $D$ is split into $H$ heads, each of dimension $P$ ($D = H \cdot P$).
• Projections for $Q, K, V, X$ and $\Delta$ are computed in a single unified matrix multiplication at the block entrance.
• Multi-head configurations allow sharing $A$ and $B, C$ representations across heads, mirroring Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) parameter efficiencies.

Architectural Comparison: Mamba vs. Transformer

The mathematical differences between Multi-Head Self-Attention, Linear Attention, Mamba-1, and Mamba-2 translate into distinct computational and memory profiles:

| Dimension | Standard Transformer (MHA) | Linear Attention | Mamba-1 (Selective SSM) | Mamba-2 (SSD) | | :--- | :--- | :--- | :--- | :--- | | Training Time Complexity | $O(T^2 \cdot d)$ | $O(T \cdot d \cdot N)$ | $O(T \cdot d \cdot N)$ | $O(T \cdot d \cdot N)$ | | Inference Step Compute | $O(T \cdot d)$ | $O(d \cdot N)$ | $O(d \cdot N)$ | $O(d \cdot N)$ | | Inference Memory (KV Cache) | $O(T \cdot d)$ (Grows linearly) | $O(d \cdot N)$ (Constant) | $O(d \cdot N)$ (Constant) | $O(d \cdot N)$ (Constant) | | Hardware Core Affinity | Tensor Cores (GEMM) | Vector ALUs / Mixed | Memory Bandwidth (Scan) | Tensor Cores (Chunk GEMM) | | Associative Recall | Exact (Full Softmax Matrix) | Degraded (Kernel Decay) | High (Selective Gating) | High (Multi-Head SSD) | | State Expansion ($N$) | Implicitly $N = T$ | Typically $N = d$ | Constrained ($N = 16$) | Scalable ($N = 64 \text{--} 256$) |

Inference Throughput and Memory Scaling

In standard autoregressive Transformer decoding, each step must load the historical KV cache from GPU HBM to compute attention weights. As context length $T$ reaches tens of thousands of tokens, the memory bandwidth required to transfer the KV cache exceeds the compute time of the matrix multiplications, creating an I/O bottleneck.

In contrast, Mamba models maintain a constant state representation $h_t \in \mathbb{R}^{D \times N}$. For a model with $D = 4096$ and $N = 64$, the state footprint per sequence is strictly fixed at approximately 512 KB per layer, regardless of whether the prompt is 1,000 tokens or 1,000,000 tokens long. This allows inference engines to support batch sizes up to an order of magnitude larger on identical GPU hardware.

The Trade-Off: Associative Recall and Hybrid Architectures

While selective state spaces solve the content-blindness of LTI models, pure SSMs compress historical information into a fixed-capacity vector state $h_t$. For tasks requiring exact arbitrary lookups across long contexts, such as needles in a haystack or complex multi-hop variable tracing across 100K+ tokens, softmax attention preserves an uncompressed, lossless memory representation.

To achieve an optimal balance between throughput and arbitrary associative recall, modern production systems increasingly deploy Hybrid Architectures:

• AI21 Jamba: Interleaves Mamba-1/2 layers with standard Transformer Attention layers and Mixture-of-Experts (MoE) routing, achieving up to 3x higher throughput on long prompts while retaining Transformer-grade retrieval fidelity.
• NVIDIA Nemotron-4 / Samba: Utilizes alternating blocks of Mamba and sliding-window attention to optimize training efficiency and long-context reasoning.

Sources

• Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752
• Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv:2405.21060
• Gu, A., Goel, K., & Ré, C. (2021). Efficiently Modeling Long Sequences with Structured State Spaces. arXiv:2111.00396