arXiv 2023 / 2024

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu* (CMU) & Tri Dao* (Princeton)

TL;DR

Mamba introduces a new class of selective state space models that achieve the modeling power of Transformers while scaling linearly $O(L)$ in sequence length. By allowing parameters to be input-dependent, Mamba filters out irrelevant context and performs content-based reasoning, delivering $5\times$ higher inference throughput and state-of-the-art performance across language, audio, and genomics.

Linear Scaling $O(L)$ Selective SSM Hardware-Aware Scan Attention-Free

Why This Matters

Modern deep learning is dominated by the Transformer architecture. While attention mechanism excels at routing information dynamically across context windows, it suffers from a fundamental bottleneck: quadratic complexity $O(L^2)$. As sequence lengths stretch into the tens of thousands or millions (crucial for long documents, high-resolution audio, and genomic analysis), Transformers become prohibitively expensive to train and run.

Prior attempts to build subquadratic-time models (like linear attention, gated convolutions, and structured state space models) often struggled on discrete, information-dense modalities like natural language. The authors identify a core limitation of these models: their inability to perform content-based reasoning due to their time-invariant nature.

The Big Idea

Mamba's core contribution is the introduction of a Selection Mechanism into Structured State Space Models (SSMs). By making the transition parameters functions of the input, the model can dynamically decide whether to remember or ignore information at each timestep.

To overcome the computational challenge of this time-varying recurrence (which prevents using fast FFT-based convolutions), the authors design a hardware-aware parallel scan algorithm. This implementation utilizes the GPU memory hierarchy to avoid materializing the expanded state in slow high-bandwidth memory (HBM), making Mamba faster than previous architectures in practice.

How It Works

1. Continuous to Discrete State Spaces

Structured SSMs map a continuous input $x(t)$ to an output $y(t)$ through an $N$-dimensional latent state $h(t)$ using the following differential equations:

$$h'(t) = Ah(t) + Bx(t)$$ $$y(t) = Ch(t)$$

To apply this to discrete sequences, we discretize the continuous parameters $(A, B)$ using a step size parameter $\Delta$. Under the Zero-Order Hold (ZOH) rule, this yields:

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

Interactive Widget 1: Discretization Explorer

Adjust the step size $\Delta$ and the system matrix eigenvalue $A$ to see how the discrete transition matrix $\bar{A}$ changes, and observe how the latent state decays over time.

Step Size ($\Delta$) 0.50

System Matrix ($A$) -1.00

Calculated Parameters

$\bar{A} = \exp(\Delta A) = $ 0.61

A larger $\Delta$ forces the system to focus on the current input (faster state decay), while a smaller $\Delta$ lets the state persist, maintaining longer-term memory.

State Decay Curve ($h_t = \bar{A}^t h_0$)

Timesteps ($t$)

State Magnitude ($h_t$)

2. The Power of Selection (S6)

Standard SSMs are Linear Time-Invariant (LTI): the parameters $(\Delta, A, B, C)$ are identical for every token in the sequence. While this allows efficient training via global convolutions, it prevents the model from selecting specific tokens or filtering out noise.

Mamba makes these parameters input-dependent:

$$B_t = \text{Linear}_N(x_t), \quad C_t = \text{Linear}_N(x_t)$$ $$\Delta_t = \text{softplus}(\text{Parameter} + \text{Linear}_1(x_t))$$

This simple change transforms the model into a time-varying system, allowing it to perform tasks like selective copying and associative recall perfectly.

Interactive Widget 2: Selective Copying Demo

Click on tokens to toggle them between Data (colored) and Noise (gray). Run both models to see how the selective S6 mechanism filters out noise to retain data, while the LTI S4 model gets overwhelmed by noise.

Input Sequence (Click to toggle)

Latent Memory State

Memory retention across sequence

Model Output / Recall

—

Click run to start simulation

3. Hardware-Aware Parallel Scan

By breaking LTI, selective SSMs can no longer be computed as a convolution. The naive recurrent computation requires $O(BLDN)$ operations, which is slow and memory-intensive.

To solve this, Mamba implements a Hardware-Aware Selective Scan. Instead of writing intermediate states of size $(B, L, D, N)$ to slow High-Bandwidth Memory (HBM), the algorithm loads parameters directly into fast GPU SRAM, performs the discretization and parallel scan in SRAM, and writes only the final output of size $(B, L, D)$ back to HBM.

Interactive Widget 3: Block Architecture Comparison

Hover over the blocks to inspect how the Mamba architecture simplifies previous designs (H3 and Gated MLP) into a single, homogenous recurrent block.

H3 Block

Linear Projection

Shift-SSM (Conv)

Main SSM (S4)

Multiplicative Gate

Gated MLP Block

Linear Projection

Activation ($\sigma$)

Multiplicative Gate

Mamba Block Proposed

Expanded Projection ($2D$)

1D Convolution

Selective SSM (S6)

SiLU / Activation

Mamba Block details: Combines the H3 block and the Gated MLP block into a single, unified recurrent layer. By expanding the block dimension by a factor $E=2$, it places the majority of parameters in simple linear projections while the selective SSM handles sequence dynamics.

Results & Benchmarks

Mamba achieves remarkable empirical success across multiple domains, matching or outperforming Transformers while offering significantly faster inference.

$5\times$

Higher Throughput

Compared to standard Transformers of similar size due to no KV cache requirement.

1M+

Context Length

Perfect extrapolation on synthetic tasks and monotonic improvement on real data.

Parameter Match

Mamba-3B outperforms Pythia-3B and matches Pythia-7B on common-sense reasoning.

Model	Size	LAMBADA (Acc ↑)	HellaSwag (Acc ↑)	PIQA (Acc ↑)
Pythia	160M	33.0%	30.2%	61.4%
Mamba	130M	44.3%	35.3%	64.5%
Pythia	2.8B	64.7%	59.3%	74.0%
Mamba	2.8B	69.2%	66.1%	75.2%

Limitations & Open Questions

While Mamba represents a massive leap forward for alternative architectures, the authors highlight several open questions:

Continuous-Discrete Spectrum: While selection helps immensely on discrete data (text, DNA), LTI models still excel on highly smooth, continuous modalities like raw audio waveforms.
Downstream Affordances: Transformers benefit from a mature ecosystem of tooling (quantization, fine-tuning recipes, RLHF). Adapting these directly to SSMs requires further research and engineering.
Scale Validation: The paper evaluates models up to 3B parameters. Validating Mamba's scaling behaviors at the 70B+ parameter scale remains an active area of exploration.

Glossary

State Space Models (SSM)

Mathematical formulations derived from control theory that map a 1D input sequence to a 1D output sequence via an implicit $N$-dimensional latent state.

Linear Time Invariance (LTI)

A property of systems where the rules governing transition do not change over time. In deep learning, LTI allows representing recurrent equations as global convolutions.

Parallel Scan

An algorithm that computes all prefix sums of a sequence in parallel, scaling logarithmically with sequence length $O(\log L)$, allowing recurrent models to parallelize training.

Citation

@article{gu2023mamba,
  title={Mamba: Linear-time sequence modeling with selective state spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023}
}