llms.blog
AI

State Space Models in Large Language Models: How Mamba, S4, and Selective Recurrence Challenge Transformer Attention

Modern large language models rely almost universally on the Transformer architecture. However, the core mechanism powering Transformers, softmax multi-head self-attention, exhibits fundamental scaling limitations. Specifically, standard self-attention requires quadratic time and memory complexity relative to sequence length during prefilling, alongside a linear memory expansion for the key-value (KV) cache during autoregressive token generation. To circumvent these computational bottlenecks, re

axc
axc
6 min
State Space Models in Large Language Models: How Mamba, S4, and Selective Recurrence Challenge Transformer Attention

Modern large language models rely almost universally on the Transformer architecture. However, the core mechanism powering Transformers, softmax multi-head self-attention, exhibits fundamental scaling limitations. Specifically, standard self-attention requires quadratic time and memory complexity relative to sequence length during prefilling, alongside a linear memory expansion for the key-value (KV) cache during autoregressive token generation.

To circumvent these computational bottlenecks, researchers developed State Space Models (SSMs). Rooted in classical continuous control theory and modernized through deep learning, SSM architectures such as S4, Mamba, and Mamba-2 demonstrate linear time scaling with sequence length and constant memory requirements during autoregressive decoding.

State Space Models vs Attention Complexity

The Attention Bottleneck: Quadratic Cost and Memory Bound Decoding

The Transformer self-attention operation computes pairwise interaction scores across all tokens in an input sequence. For an input sequence of length N and hidden dimension D, computing the attention matrix requires O(N^2 D)* floating-point operations.

During inference, autoregressive token generation generates one token at a time. To avoid recomputing past key and value projections at every generation step, inference engines cache historical states in GPU high-bandwidth memory (HBM). This KV cache expands linearly with sequence length:

  • Memory per token = 2 * (Layers) * (Hidden Dimension) * (Bytes per Parameter)

For multi-thousand token contexts and concurrent batch processing, the KV cache quickly consumes tens of gigabytes of VRAM. This shifts the serving bottleneck from compute saturation (arithmetic intensity) to memory bandwidth limits, throttling inference throughput.

Linear attention variants historically attempted to remove the softmax nonlinearity to allow recurrent formulation, but suffered severe degradation in language modeling quality and struggled to retain information across long contexts.

Continuous-Time State Space Foundations and S4

State Space Models represent physical systems mapped through continuous differential equations. A continuous-time 1D linear state space system maps an input signal x(t) to a latent state representation h(t) and output y(t) via linear differential equations:

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

Here, A is the state transition matrix governing latent dynamics, B maps input signals into the state, C projects the state to outputs, and D represents a direct feedthrough skip connection.

Discretization

Digital hardware processes discrete sequences rather than continuous signals. Continuous SSMs are discretized using a step size parameter Δ via methods such as the Zero-Order Hold (ZOH):

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

Once discretized, the system operates across two equivalent computational views:

  1. Recurrent View (Inference): The system computes the next state via h_t = A_bar h_{t-1} + B_bar * x_t* and output y_t = C h_t. This yields O(1)* compute and strictly constant memory footprint per generation step.
  2. Convolutional View (Training): For parallel training on known sequences, unrolling the recurrence yields a 1D convolution y = x K_bar, where the convolution kernel K_bar = (C * B_bar, C * A_bar * B_bar, ..., C * A_bar^(N-1) * B_bar). Using Fast Fourier Transforms (FFT), convolution over sequence length N* computes in O(N log N) parallel time.

Structured State Spaces (S4) and HiPPO

Early discretized SSMs suffered from vanishing and exploding gradients over long contexts. In 2021, Albert Gu, Karan Goel, and Christopher Ré introduced the Structured State Space sequence model (S4).

S4 incorporated the High-order Polynomial Projection Operators (HiPPO) framework. HiPPO constructs a specialized transition matrix A whose continuous dynamics mathematically maintain optimal polynomial projections of past input history over decaying time windows. S4 parameterized A as a normal-plus-low-rank matrix, enabling stable diagonalization and reducing kernel computation to a Cauchy kernel calculation.

The Selection Bottleneck: Why Linear Time-Invariant SSMs Failed at Language

Despite strong results on audio and synthetic benchmarks, S4 and related Linear Time-Invariant (LTI) SSMs trailed Transformers on natural language tasks.

In an LTI system, the matrices (A_bar, B_bar, C, Δ) remain fixed across every timestep t, irrespective of what token x_t contains. The model applies an identical convolution filter across the entire sequence.

Natural language modeling requires dynamic, content-aware processing:

  • Selective Filtering: The model must ignore irrelevant prompt tokens while storing critical facts.
  • Associative Recall: The model must dynamically retrieve specific key-value associations established hundreds of tokens earlier.
  • Induction Heads: The model must track repetition patterns and adjust outputs based on specific contextual triggers.

Because LTI SSMs treat every token through uniform linear dynamics, they cannot compress information adaptively based on context.

Mamba: Selective State Spaces and Hardware-Aware Scans

In December 2023, Albert Gu and Tri Dao introduced Mamba (S6). Mamba solved the expressivity limitation by transforming SSMs from linear time-invariant systems to selective, input-dependent systems.

Dynamic Parameterization

In Mamba, the input projection matrices and discretization step size become dynamic functions of the current input token:

  • B_t = Linear_B(x_t)
  • C_t = Linear_C(x_t)
  • Δ_t = Softplus(Parameter + Linear_Δ(x_t))

The step size Δ_t functions as a continuous gating mechanism. When Δ_t is large, the state update incorporates the current input x_t heavily; when Δ_t is small, the model bypasses the input and preserves its existing recurrent state h_t.

The Hardware-Aware Parallel Associative Scan

Making parameters input-dependent breaks the convolutional training view. Because A_bar_t and B_bar_t change at each position t, training cannot use FFT convolution.

A naive recurrent loop over sequential tokens is memory-bandwidth bound on modern GPUs. Moving intermediate state tensors of dimension (Batch, Length, Dimension, State) between GPU high-bandwidth memory (HBM) and processor registers creates massive I/O overhead.

Mamba resolved this with a hardware-aware parallel scan kernel:

  • The kernel loads input sequences and parameters from HBM directly into high-speed GPU on-chip SRAM once.
  • It executes an associative prefix scan directly in SRAM across thread blocks.
  • It writes only the final output representations back to HBM, avoiding materialization of intermediate hidden states in slow memory.

This implementation delivers training throughput on par with FlashAttention while preserving linear O(N) time scaling.

Mamba-2 and Structured State Space Duality (SSD)

In 2024, Tri Dao and Albert Gu introduced Mamba-2, formalizing Structured State Space Duality (SSD).

SSD establishes a theoretical bridge between selective state space models and attention mechanisms. The authors demonstrated that selective SSMs and 1-semiseparable matrix transformations represent identical mathematical structures.

Under the SSD framework:

  • Selective SSMs can be computed via block-wise matrix multiplication.
  • The recurrence within small blocks is computed via parallel scan, while cross-block interactions are computed via matrix multiplication on GPU Tensor Cores.
  • Mamba-2 achieves 2x to 8x higher processing speed than Mamba-1 during training by replacing scalar recurrences with structured Tensor Core matrix operations.

Architectural Trade-Offs and Hybrid Deployments

While pure state space models solve the quadratic scaling wall, they introduce distinct trade-offs when compared against standard Transformers:

  1. In-Context Associative Recall: A fixed-size state vector h_t has finite information capacity. On complex synthetic retrieval benchmarks and dense multi-hop reasoning tasks, Transformers with full quadratic attention retain absolute access to every historical token without compression loss.
  2. Memory Footprint: SSMs maintain a constant memory state during generation regardless of whether the prompt is 1,000 or 1,000,000 tokens long. In contrast, Transformer KV cache footprints expand linearly.
  3. Throughput at Long Contexts: SSM inference throughput remains stable across extended sequence lengths, whereas Transformer generation slows down as KV cache lookups saturate memory bandwidth.

Hybrid Architectures in Production

To balance inference efficiency with precise in-context recall, modern implementations frequently combine SSM layers with sparse attention layers:

  • AI21 Labs Jamba: Interleaves Transformer attention layers and Mamba SSM layers in an 1:7 ratio alongside Mixture-of-Experts (MoE) routing, slashing KV cache memory footprint by up to 8x while maintaining full attention retrieval performance.
  • Samba and Nemotron-4: Combine state space blocks for local feature aggregation with periodic global attention layers for dense cross-document retrieval.

State Space Models have established linear-time sequence modeling as a practical alternative to pure Transformer architectures, providing a scalable foundation for long-context generation and resource-constrained edge serving.

Sources

AILLMArchitectureEducational
ShareX (Twitter)LinkedIn

Written by

axc

More to read

  • Document Parsing and Visual Retrieval for Production RAG: Architecture, Benchmarks, and Serving Trade-Offs for Docling, Marker, MinerU, and ColPali

    Document Parsing and Visual Retrieval for Production RAG: Architecture, Benchmarks, and Serving Trade-Offs for Docling, Marker, MinerU, and ColPali The retrieval quality of a Retrieval-Augmented Generation (RAG) system is strictly bounded by the fidelity of its document ingestion pipeline. In enterprise environments, the vast majority of domain knowledge remains locked in unstructured Portable Document Format (PDF) files, scanned reports, technical manuals, and multi-column research papers. Na

    1 min

  • Context Window Extension in Large Language Models: How Position Interpolation, YaRN, and LongRoPE Scale Sequence Lengths

    Large language models are bounded during pretraining by a fixed sequence length, typically between 2,048 and 8,192 tokens. When standard autoregressive transformers attempt to process sequences beyond this pretraining context window, performance degrades immediately. Perplexity rises sharply and the model loses coherence within a few dozen tokens past the training boundary. Extending this context window by training from scratch on long sequences is computationally prohibitive due to the quadrat

    1 min

  • Round Hill Files $1B Copyright Infringement Lawsuits Against Anthropic and Suno

    Independent music rights administrator Round Hill Music has filed twin copyright infringement lawsuits against generative AI music platform Suno and frontier foundation model developer Anthropic. The complaints, filed in the U.S. District Court for the Northern District of California, allege that both companies unlawfully scraped, ingested, and reproduced copyrighted musical compositions without licenses, authorization, or compensation to build and train their commercial AI models. Round Hill M

    1 min