Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Problem

Framing

Linear-time sequence models had not matched Transformer quality on discrete, information-dense modalities like language. Mamba closes this gap by making SSM dynamics input-selective and pairing them with a hardware-aware recurrent scan, yielding linear scaling, 5\times inference throughput gains, and strong language results up to 3B parameters.

Currently Used Methods

Foundational

• @vaswaniAttentionAllNeed2017 — dense self-attention enables strong content-based token interactions.
  • Limitation in context: quadratic context cost and KV-cache-heavy autoregressive inference.
• @hochreiterLSTM1997 — gated recurrence with input-dependent memory updates.
  • Limitation in context: sequential execution limits throughput and long-context scaling.
• "Efficiently Modeling Long Sequences with Structured State Spaces" — structured SSMs give long-range linear recurrences.
  • Limitation in context: time-invariant dynamics weaken content-based reasoning on text.
• @pengRWKV2023 — recurrent language model with attention-like decay and token mixing.
  • Limitation in context: limited state expansion and less expressive selective SSM parameterization.
• "Hyena Hierarchy: Towards Larger Convolutional Language Models" — long convolutions offer subquadratic sequence mixing.
  • Limitation in context: convolutional kernels do not support selective token-dependent state updates.

Proposed Method

Architecture

Mamba repeats a single homogeneous block. Each block uses an input projection, a short convolution, a selective SSM branch, and a multiplicative gate from the parallel branch; it removes attention and standalone MLP blocks by folding the SSM path into a gated MLP-style design.

Loss / Objective

The paper trains Mamba with the standard next-token autoregressive objective.

Sampling Rule / Algorithm

The selective SSM updates the hidden state with input-dependent discretized parameters.

Training Procedure

• Language pretraining data: The Pile.
• Language training length for zero-shot comparison: 300B tokens.
• Context lengths in scaling experiments: up to $2^{20}$.
• Total batch size in scaling experiments: about $2^{24}$ tokens/step.
• Optimizer: AdamW.
• Weight decay: 0.1.
• Warmup: 10% of total steps.
• AdamW $\beta$: $(0.9, 0.95)$ in speech experiments.
• Mamba learning-rate sweep reported up to $8 \times 10^{-3}$.

Evaluation

Datasets

• The Pile language modeling.
• HG38 human genome pretraining.
• Great Apes genomic classification.
• SC09 speech generation.
• Synthetic selective copying.
• Synthetic induction heads.

Metrics

• Perplexity.
• Zero-shot accuracy.
• FID.
• sFID.
• FAD.
• Induction/selective-copy accuracy.
• Throughput.

Headline results

• Language modeling: Mamba-3B outperforms same-size Transformers and matches Transformers about 2\times larger.
• Language zero-shot: Mamba is reported best-in-class at each size in Table 3.
• Inference: about 5\times higher throughput than comparable Transformers.
• Induction heads: perfect generalization to sequences beyond $10^6$ tokens, about 4000\times train length.
• Speech generation (SC09): FID drops by more than half versus prior state of the art.

Ablations

• Selective parameters: $\Delta$ is the dominant selective parameter; adding selective $\mathbf{B}, \mathbf{C}$ helps further.
• SSM parameterization: real-valued S4 variants match complex variants inside the block.
• Architecture: the simpler Mamba block matches H3 while using a more uniform design.
• State size: larger recurrent state dimension $N$ improves performance.

Method Strengths and Weaknesses

Strengths

• Restores content-dependent reasoning to SSMs with input-selective dynamics.
• Linear-time recurrent inference avoids Transformer KV-cache growth.
• Strong cross-domain results span language, genomics, audio, and synthetic recall.
• Hardware-aware scan makes selective recurrence practical on modern GPUs.

Weaknesses

• Core scan algorithm is substantially more specialized than standard attention kernels.
• Evaluation emphasizes autoregressive benchmarks, not bidirectional or encoder-decoder settings.
• Reported gains rely on careful systems design, not only the mathematical model.
• Method adds custom discretization and scan machinery that complicates implementation.

Suggestions from the authors

• Apply selective SSMs to more modalities and foundation-model settings.
• Study whether models learn automatic state resets across episode boundaries.
• Explore other forms of selection beyond $\Delta$, $\mathbf{B}$, and $\mathbf{C}$.
• Extend selective mechanisms to broader RNN and CNN families.

Links

Prior Papers

• @vaswaniAttentionAllNeed2017 — Transformer attention is the main baseline Mamba targets on quality and efficiency.
• @hochreiterLSTM1997 — Mamba connects selective SSM updates to classical gated recurrence.
• @pengRWKV2023 — RWKV is a recurrent language-model baseline with overlapping efficiency goals.

Further Papers

No vault papers identified as further work yet.