State Space Model (SSM)

The defining advantage of SSMs is their O(N) time and memory complexity, where N is sequence length. Unlike Transformers, which scale quadratically (O(N²)) due to pairwise token comparisons in self-attention, SSMs process sequences recurrently.

• Mechanism: The model maintains a fixed-size hidden state that summarizes all past information, updating it token by token.
• Impact: Enables processing of sequences with 1M+ tokens on a single GPU without the memory explosion seen in standard attention.
• Contrast: A 32K context in a Transformer requires ~8GB of memory for the attention matrix alone; an SSM processes the same sequence with a constant memory footprint.

SSMs are governed by a continuous-time linear differential equation: ẋ(t) = Ax(t) + Bu(t), with output y(t) = Cx(t). The HiPPO (High-order Polynomial Projection Operators) framework defines the optimal A matrix to compress the entire history of the input signal into the hidden state.

• Discretization: The continuous system is converted to a discrete recurrence using a Zero-Order Hold (ZOH) rule with a learnable step size Δ, producing: x_k = Āx_{k-1} + B̄u_k.
• Duality: This recurrence can be equivalently computed as a global convolution during training, enabling parallel processing across the entire sequence, then switched to the recurrent mode for efficient autoregressive inference.

The Mamba architecture introduces input-dependent parameters, making the SSM selective. Unlike prior SSMs (S4, H3) where A, B, and C are time-invariant, Mamba makes B, C, and Δ functions of the current input.

• Selectivity: The model learns to ignore or focus on specific inputs, filtering out irrelevant information dynamically—a capability previously exclusive to attention mechanisms.
• Hardware-Aware Algorithm: Because selectivity prevents the use of efficient static convolutions, Mamba implements a parallel scan in custom CUDA kernels (fused with memory I/O) to compute the linear recurrence without materializing the full state in slow GPU HBM.
• Result: Matches or exceeds Transformers of the same size on language modeling while maintaining linear scaling.

For a time-invariant SSM, the sequential recurrence can be unrolled into a single depthwise convolution with a kernel derived from the discretized system matrices (Ā, B̄, C).

• Parallelism: This allows the entire sequence to be processed in one pass during training, leveraging highly optimized FFT-based convolution algorithms on GPUs.
• SSM Kernel: The kernel K̄ = (CB̄, CĀB̄, ..., CĀ^{L-1}B̄) captures the system's impulse response, effectively computing all hidden states simultaneously.
• Constraint: This mode requires the parameters to be static (time-invariant), which is why Mamba's selective mechanism necessitates a return to recurrent computation with hardware-aware scan operations.

Modern SSM architectures, including Mamba, integrate the state space layer into a gated block inspired by the Gated Linear Unit (GLU) and the SwiGLU activation used in Transformers.

• Architecture: The input is projected through two linear layers. One branch passes through the SSM core and a SiLU (Swish) activation. The other branch is gated by a second SiLU activation, and the two are combined via element-wise multiplication.
• Purpose: The gating mechanism provides a non-linear information bottleneck, allowing the block to control the flow of information from the linear SSM recurrence.
• Design Philosophy: This hybrid approach treats the SSM as a powerful linear sequence mixer, while relying on proven MLP gating components for channel mixing and non-linearity.

SSMs were originally validated on the Long-Range Arena (LRA) benchmark, a suite of tasks designed to test a model's ability to handle sequences from 1K to 16K tokens across modalities including text, images, and mathematical reasoning.

• Path-X: A 16K×16K pixel image classification task requiring the model to identify whether two markers are connected by a path. SSMs (S4) were the first non-attention models to solve this task.
• Performance: SSMs achieve state-of-the-art accuracy on LRA while using a fraction of the compute and memory required by Transformer variants.
• Significance: Demonstrated that the quadratic attention mechanism is not a strict requirement for capturing long-range dependencies, validating the SSM approach for ultra-long context applications.

A memory buffer in Transformer decoders that stores pre-computed Key and Value tensors from previous time steps to avoid redundant computation during autoregressive generation. The cache size grows linearly with sequence length, creating a significant memory bottleneck. A primary advantage of recurrent SSMs like Mamba is the elimination of the explicit KV-Cache, as the model's fixed-size latent state inherently captures historical context without requiring a growing memory allocation.

A class of attention approximations that reformulate the softmax attention kernel to achieve linear O(n) complexity. By expressing attention as a kernel function and changing the order of matrix multiplication, linear attention avoids computing the full N×N attention matrix. However, many linear attention variants suffer from degraded recall quality. SSMs like Mamba are often viewed as a more principled evolution of this idea, using a data-dependent gating mechanism to dynamically filter information without the kernel approximation trade-offs.

A simple architecture that uses a gating mechanism (element-wise multiplication of two linear projections) to control information flow, popularized by architectures like Gated Linear Units (GLU). Mamba's Selective SSM block integrates this concept directly into the state space framework by making its continuous-time parameters (B, C, Δ) input-dependent functions. This data-dependence is the key innovation that allows the model to selectively propagate or forget information, overcoming the rigidity of purely linear time-invariant SSMs.