Mamba: Linear-Time Sequence Modeling with Selective State Spaces

In plain words

Most language models read text by letting every token look back at every earlier token. That works very well but gets expensive as the text grows. A cheaper family, state-space models, instead passes a small running memory along the sequence. The catch is that the memory update is the same for every token, so the model cannot decide, token by token, what to remember and what to ignore. It finds it hard to pick out one useful word among filler.

This paper makes the memory update depend on the current token. A "step size" decides how much of the old memory to keep and how much of the new token to write in, and the same idea controls what is read out. The catch is that this no longer lets the computation be written as one fixed filter, so the authors also wrote a GPU routine that runs the updates in parallel without storing the full memory for every step. The resulting model has no attention and no separate feed-forward block; one repeated block does the work. The authors report that it matches or beats open Transformers of the same size on language, and that it also works on audio and DNA.

TL;DR

• Mamba is a state-space sequence model whose step size, input matrix and output matrix depend on the current input (the "selection" mechanism) §3.2.
• The input-dependent recurrence cannot be computed as a convolution, so the authors give a fused, hardware-aware parallel scan with recomputation in the backward pass §3.3, app. D.
• The model is one homogeneous block, repeated, with no attention and no MLP §3.4.
• Author-reported: best zero-shot scores among open models at each size from 130M to 2.8B; about 5x inference throughput versus Transformers of similar size §4.2.2, §4.5.
• Released code is Apache-2.0 at state-spaces/mamba; the file for this paper's Mamba-1 layer is mamba_simple.py [code].

Problem

Attention keeps the whole context, so inference must store every past key and value (the KV cache), and training cost grows with the square of the length §3.1. Recurrent models keep a fixed-size state and are cheap, but their effectiveness depends on how well the state compresses the context §3.1.

Earlier structured state-space models (S4 and its relatives) are linear time-invariant: their parameters are the same at every step. The authors argue that such models cannot select content, so they fail a Selective Copying task where the positions of the tokens to remember vary, while they solve the fixed-spacing Copying task §3.1. Making the parameters input-dependent removes the convolution form, which was the reason these models were efficient §3.2, §3.3.

The assumptions are: discrete data such as text and DNA, long sequences, and a GPU that is limited by memory bandwidth rather than by matrix multiplication §3.3, app. D.

Method

Background. A structured SSM maps a scalar input sequence to an output through a latent state of size N, using four parameters (Δ, A, B, C). In continuous time, h' = A h + B x and y = C h §2. Zero-order hold (ZOH) discretisation with step Δ gives Ā = exp(ΔA) and B̄ = (ΔA)^-1 (exp(ΔA) − I) · ΔB §2, eq. 4. The discrete system can be run as a recurrence h_t = Ā h_{t−1} + B̄ x_t, y_t = C h_t, or, when Δ, A, B, C are fixed, as a convolution with a kernel built from C B̄, C Ā B̄, and so on §2.

Selection. Let B, C and Δ be functions of the input x_t, and keep A fixed §3.2, alg. 2:

• B = Linear_N(x), C = Linear_N(x) §3.2;
• Δ = softplus(Parameter + Broadcast_D(Linear_1(x))) §3.2;
• A is kept time-invariant, because it only affects the model through Δ·A in Ā and the authors expect a selective A to give similar results §3.2, §3.5.

Because B, C and Δ now carry a time index, the model is time-varying, and the convolution equivalence is lost §3.2. The authors also show that for N = 1, A = −1, B = 1 and a linear Δ, the selective recurrence becomes a gated RNN, h_t = (1 − g_t) h_{t−1} + g_t x_t with g_t = σ(Linear(x_t)); this is their explanation of why Δ is the key parameter §3.5, eq. 5, app. C.

Hardware-aware scan. The naive method would materialise Ā and B̄ of shape (B, L, D, N) and the states in high-bandwidth memory app. D. The authors instead fuse three steps into one kernel: read Δ, A, B, C from slow memory, discretise and run a parallel associative scan in fast on-chip memory, then multiply by C and write an output of shape (B, L, D) §3.3, app. D. The backward pass recomputes the intermediate states instead of storing them §3.3, app. D. Long sequences are split into chunks; the scan state carries across chunks app. D.

Block. Each block expands the model width D by E = 2, projects the input into two branches, applies a short causal convolution and SiLU to one branch, runs the selective SSM on it, multiplies by SiLU of the other branch, and projects back to D §3.4, app. E.2.2. A normalisation layer and a residual connection wrap the block, and the blocks are stacked with no separate attention or MLP §3.4. Relative to H3, the block replaces the first multiplicative gate with an activation; relative to the gated MLP, it adds an SSM to the main branch §3.4.

# Selective SSM (S6) over one sequence  [alg. 2] [§3.2]
# Input:  x (B, L, D)
# Output: y (B, L, D)
# A is a fixed (D, N) parameter, with negative real values [§3.6]

def selective_ssm(x, A):
    # Per-token SSM inputs; B and C are linear maps of x  [§3.2]
    B_sel = linear_N(x)                            # (B, L, N)
    C_sel = linear_N(x)                            # (B, L, N)
    # Per-token step size: 1-dim projection, broadcast to D channels,
    # plus a learned per-channel bias, then softplus  [§3.2]
    delta = softplus(delta_bias + broadcast_D(linear_1(x)))  # (B, L, D)
    # Zero-order hold: A_bar = exp(delta A) and B_bar = (A_bar - 1) B / A
    # (equal to (delta A)^-1 (A_bar - 1) delta B)  [§2] [eq. 4]
    A_bar = exp(delta[..., None] * A)              # (B, L, D, N)
    B_bar = (A_bar - 1) / A * B_sel[:, :, None, :]  # (B, L, D, N)
    # Recurrence over time; the GPU version is a parallel scan  [§3.3]
    h = zeros(B, D, N)
    y = zeros(B, L, D)
    for t in range(L):
        h = A_bar[:, t] * h + B_bar[:, t] * x[:, t, :, None]
        # Read out with C at this step  [§2]
        y[:, t] = sum_n(h * C_sel[:, t, None, :])  # (B, D)
    return y

# Mamba block  [§3.4] [app. E.2.2]
# Input:  hidden (B, L, D); expansion E = 2
# Output: (B, L, D)

def mamba_block(hidden):
    # Two branches, each of width E * D  [§3.4]
    x, gate = linear_in(hidden).split(2)           # (B, L, E*D) each
    # Short causal convolution and SiLU on the SSM branch  [app. E.2.2]
    x = silu(causal_conv(x))                       # (B, L, E*D)
    y = selective_ssm(x, A)                        # (B, L, E*D)
    # Multiplicative SiLU gate from the other branch  [§3.4]
    y = y * silu(gate)                             # (B, L, E*D)
    return linear_out(y)                           # (B, L, D)

Hyperparameters

Values in the first block are from the paper. Rows marked "code" are read from the released code at code_commit and are not stated in the paper.

Results

Author-reported. Setup: Mamba models of 130M to 2.8B parameters, with baselines up to 7.4B; zero-shot models trained on 300B tokens, scaling runs on The Pile. Baselines are GPT-3 style Transformers, Transformer++ (rotary, SwiGLU, improved recipe), H3, Hyena, RWKV, RetNet, and Pythia §4.2, app. E.2. Hardware for speed numbers is one A100 80GB app. E.5.

Synthetic tasks §4.1. Selective Copying accuracy:

Induction Heads, trained at length 256: Mamba reaches perfect accuracy at all tested lengths up to 2^20 §4.1, app. E.1. The authors report that no other method goes beyond about 2x the training length §4.1. The appendix table lists the MHA variants as running out of memory from 2^15 onward, and H3 at 4.7–8.2% from 2^15 to 2^20 app. E.1.

Language modelling, scaling §4.2.1. Pile validation perplexity across 125M to 1.3B at the Chinchilla protocol and 2k and 8k context. The paper's figure shows Mamba as the first attention-free model to match Transformer++, mainly at longer context §4.2.1. No numbers are tabulated in the text.

Zero-shot evaluation §4.2.2. Pile perplexity is on the validation split and only compares with models using the same tokenizer and dataset (Pythia, RWKV). Accuracy on six tasks; the Average column is the mean of those six accuracies. Bold values are the best in each size class in the source.

HellaSwag and Arc-C use length-normalised accuracy app. E.2.3. The text's "Mamba-3B" claim is not a row in this table (see Pitfalls).

DNA §4.3, app. E.3. HG38 pretraining: Mamba scales better than HyenaDNA and Transformer++ with model size at context 1024, and improves with longer context while the baselines do not. Great Apes classification after fine-tuning (chance is 20%) app. E.3.4:

Audio §4.4. YouTubeMix (piano, 16 kHz) pretraining: Mamba beats the SaShiMi S4+MLP baseline at every context from 8192 to about 10^6, and the gap widens with length §4.4.1. SC09 speech generation (unconditional, 1-second clips, digits 0-9) §4.4.2:

The SC09 architecture ablation at 6M parameters, outer and centre blocks §4.4.2:

Speed and memory §4.5, app. E.5. The fused scan is up to 20–40x faster than a standard PyTorch scan and faster than FlashAttention-2 beyond sequence length 2K §4.5. Inference throughput: 4–5x that of a Transformer of similar size; the authors' abstract figure is 5x §4.5, §1. Training memory for 125M models, batch size vs GB:

The authors say Mamba's memory is comparable to a highly optimised Transformer, not lower app. E.5.

Ablations §4.6, 350M models at Chinchilla tokens, perplexity on the Pile:

Which parameters are selective (350M models) §4.6.2:

State dimension N, with B and C selective (bottom) or constant (top) §4.6.2:

Initialisation of A with a selective SSM §4.6.2: A_n = −(n+1) and a random exp-normal init both give 8.71; A_n = −1/2 gives 8.85; the complex S4D-Lin init gives 9.16.

Paper vs code

Code: official repository https://github.com/state-spaces/mamba, commit e9594ce1c732d97440f0332fdc43170a2294dbfa, Apache-2.0 LICENSE. The paper links this repository in its introduction §1. The repository's HEAD also contains later Mamba-2 and Mamba-3 modules (mamba2.py, mamba3.py); the Mamba-1 layer described in this card is mamba_ssm/modules/mamba_simple.py mamba_simple.py.

1. Δ initialisation. Paper: Δ is initialised as τ⁻¹(Uniform([0.001, 0.1])) §3.6. Code: Δ is sampled log-uniformly between 0.001 and 0.1 and floored at 1e-4, then the inverse softplus is applied mamba_simple.py. The two give different distributions. The paper's version is uniform in Δ itself; the code's is uniform in log Δ.

2. B discretisation. Paper: B̄ is the ZOH formula (ΔA)⁻¹(exp(ΔA) − I)·ΔB §2, eq. 4. Code: Ā = exp(ΔA) is used, but the input term is Δ·B·u, with no ZOH factor for B selective_scan_interface.py. For N = 1, A = −1, the paper's gate form is g = 1 − e^(−Δ) app. C, while the code's input weight is Δ. The two agree to first order in Δ only. The gating interpretation in §3.5 therefore describes the paper's discretisation, not the code path.

3. Δ projection. Paper's default: s_Δ(x) = Broadcast_D(Linear_1(x)) §3.2; the low-rank Linear_D(Linear_R(x)) is an alternative §3.6. Code: x_proj produces a rank-R vector (R = ceil(D/16)) that is projected to D by dt_proj, i.e. the low-rank form mamba_simple.py. The paper's ablation with Δ projection size 64 §4.6.2 is not the code default.

4. Skip term D. Paper: the algorithm and equations have no skip term §3.2, alg. 2, §2. Code: adds u·D inside the scan and multiplies by SiLU(z) after selective_scan_interface.py. The model in the code is therefore y = SSM(x) + D·x, gated.

5. Causal convolution. The paper mentions a short convolution on the SSM branch app. C, app. E.2.2 but does not state its width. Code: depthwise causal conv with width 4 and bias mamba_simple.py.

6. Gate. Paper: SiLU applied to the second branch, multiplied with the SSM output §3.4. Code: same, y·SiLU(z) mamba_simple.py. No discrepancy found.

7. A initialisation. Paper: S4D-Real, A_n = −(n+1) §3.6. Code: A = 1..N, stored as log, and A = −exp(A_log) mamba_simple.py. No discrepancy found.

8. State size. Paper's efficiency benchmark uses N = 16 §4.5, app. E.5. Code default d_state = 16 mamba_simple.py. Consistent.

Pitfalls

• ZOH versus the code's B approximation. Implementing the paper's eq. 4 for B̄ will not match the released kernel. If the goal is the paper's math, implement eq. 4; if the goal is to match the released code, use Δ·B·u and the code's Δ init (see Paper vs code). Say which one the implementation follows.

• Model size naming. The abstract and §1 say "Mamba-3B"; the zero-shot table has Mamba at 130M, 370M, 790M, 1.4B and 2.8B, with no Mamba row at 3B or 6.9B §4.2.2. The "matches a Transformer twice its size" claim rests on the text, not on a table row. Do not infer a 3B result from the table.

• Tie on one task. The table bolds WinoGrande for both Pythia-160M and Mamba-130M (51.9). The "best in class on every result" claim §4.2.2 therefore has a tie at 130M.

• Throughput test uses untrained weights. The 6.9B inference throughput comparison uses an untrained Mamba model; the Transformer baseline is the Hugging Face Transformers implementation app. E.5. Throughput also depends on the batch size that fits in memory §4.5.

• No variance. The paper does not report seeds, error bars or the number of runs for the training results (only three runs for inference timing) app. E.5. Treat small gaps in ablations (for example 8.69 vs 8.71) as within noise until replicated. This is a gap in the paper, not something the card can fill.

• Selection hurts raw audio. On YouTubeMix, Mamba-S6 is worse than Mamba-S4 (LTI) in every block; keeping LTI layers in the outer U-Net stages and using S6 only in the centre closes most of the gap app. E.4.1. The paper's advice is to keep layers next to raw signal LTI. Use S6 for discrete tokens such as text and DNA §5, app. E.4.1.

• Real vs complex. The paper uses real A everywhere except the audio experiment, where it switches to complex S4-style A §3.6, §4.4.1. The complex audio model is still better than the real one in that setting app. E.4.1.

• Learning rates. For DNA the swept LR is at the top of its range (8e-3 for Mamba), which the authors say may leave results suboptimal app. E.3.2. For the species task, Mamba uses 1e-4 and Hyena 4e-5 app. E.3.4.

• Hardware dependence. All speed and memory numbers are on A100 80GB GPUs with BF16 app. E.5. The fused scan is a CUDA kernel; the paper's own reference (selective_scan_ref) is a Python loop over the sequence length selective_scan_interface.py and is for testing, not speed.

• Memory cost. The paper estimates about 16 bytes of activation per token per selective SSM layer, against 12 for attention and 20 for an MLP, assuming mixed-precision training app. D. The memory table in app. E.5 shows Mamba using more memory than the FlashAttention-2 Transformer at every tested batch size.

• Inference state. Decoding keeps a convolution state of size D_inner x d_conv and an SSM state of size D_inner x N per layer, updated in place mamba_simple.py. Memory per sequence does not grow with the context; the paper credits this for the larger batch sizes behind its throughput claim §4.5.

• Sequence length in DNA. Sequence-length warmup and batch-size scheduling are part of the DNA recipe; the paper did not test turning warmup off app. E.3.3.

Lineage

• Builds on: structured state-space models S4 and its diagonal variants (S4D, S5), HiPPO initialisation, and the H3 architecture, which Mamba simplifies §2, §3.4, app. B.
• Compared against: GPT-3 style Transformer, Transformer++ (rotary, SwiGLU), H3, Hyena, RetNet, RWKV, Pythia, SaShiMi (audio), HyenaDNA (DNA), FlashAttention-2 (speed) §4, app. E.
• Follow-up in this wiki: Mamba-2 (Structured State Space Duality), which relates SSMs to masked attention and gives a chunked algorithm for the same model family.