← All papers

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Dao, Gu · 2024 · ICML 2024

ArchitectureRead on arXiv

Establishes a formal duality between structured state space models (SSMs) and a restricted form of attention, unifying the two dominant sequence modeling paradigms under a single framework called Structured State Space Duality (SSD). Mamba 2 leverages this connection to build an architecture that is 2-8x faster than Mamba 1 while achieving competitive language modeling quality with Transformers.

Key Idea

The Transformer's quadratic attention and the state space model's linear recurrence seem like fundamentally different approaches to sequence modeling. Mamba 2 reveals that they are two views of the same computation: a restricted form of linear attention (with a specific semi-separable matrix structure) is mathematically equivalent to a structured SSM. This Structured State Space Duality (SSD) framework lets you choose the most efficient algorithm for each regime, matrix multiplication (attention-like) for training on GPUs, recurrence (SSM-like) for efficient autoregressive inference.

How It Works

  • SSD framework: The key insight is that the SSM's state transition matrices, when structured as scalar times identity (or diagonal), produce a sequence transformation that can be written as a semiseparable matrix, the same structure that appears in causal linear attention. This means you can compute the same result using either a recurrence (O(n) sequential) or a matrix multiplication (O(n²) but highly parallelizable on GPUs).
  • Mamba 2 architecture: Uses the SSD layer as its core, with larger state dimensions (64-256 vs Mamba 1's 16) enabled by the more efficient SSD algorithm. The larger state gives more expressiveness, closer to full attention's capacity to store and retrieve information.
  • Hardware-aware algorithm: A chunked computation strategy splits the sequence into blocks. Within each block, use the matrix (attention-like) form for GPU parallelism; across blocks, use the recurrence (SSM) form to avoid materializing the full n×n matrix. This gives the best of both worlds: parallel training and linear-time inference.
  • Multi-head SSM: Analogous to multi-head attention, Mamba 2 uses multiple SSM "heads" with independent state dynamics, increasing model capacity without proportional compute increase.
  • Training: Standard autoregressive language modeling objective. The SSD algorithm enables 2-8x faster training than Mamba 1 at the same model size, by better utilizing GPU tensor cores through the matrix multiplication view.

Why It Matters

  • Linear-time inference: During autoregressive generation, Mamba 2 processes each new token in O(1) time with constant memory (just the state vector), vs O(n) time and growing KV cache for Transformers. This means no KV cache explosion, inference cost doesn't grow with context length.
  • Competitive quality: Mamba 2 matches Transformer quality on language modeling perplexity up to the model sizes tested (~2.7B parameters), closing the gap that existed with Mamba 1.
  • Hybrid architectures: The SSD framework enables principled mixing of SSM and attention layers. Jamba (AI21, 2024) and Zamba (Zyphra, 2024) use ~10% attention layers with ~90% Mamba layers, getting the best of both: attention's strong in-context learning and retrieval, plus Mamba's linear-time efficiency for the bulk of computation.
  • Theoretical unification: By showing SSMs and attention are dual, the paper collapses two separate research communities into one framework. New techniques developed for either paradigm can potentially transfer to the other.
  • Practical impact: Mamba-based layers are appearing in production models (Jamba, Zamba, NVIDIA Hymba) for long-context applications where the KV cache is the primary bottleneck.

Key Takeaways for Interviews

  • SSD duality: linear attention with semiseparable structure ≡ structured SSM, same computation, two algorithms (matrix multiply for training, recurrence for inference)
  • SSM vs attention tradeoffs: SSMs give O(n) inference and constant memory but are weaker at in-context learning and precise retrieval; attention gives strong recall but O(n²) compute and growing KV cache
  • Mamba 2 vs Mamba 1: larger state dimensions (64-256 vs 16), 2-8x faster training via SSD algorithm, multi-head design
  • Hybrid architectures (Jamba, Zamba) are the practical sweet spot: ~90% Mamba + ~10% attention layers. Attention handles retrieval-heavy computation; Mamba handles the rest efficiently
  • Limitations: pure SSMs still underperform full attention on tasks requiring precise information retrieval from long context (e.g., "find the needle in a haystack"), this is why hybrids work better
  • In system design: cite Mamba 2 / SSD when discussing alternatives to Transformer attention for long-context serving, and hybrid architectures when asked how to serve 100K+ token contexts cost-effectively