Mamba: Linear-Time Sequence Modeling with Selective State Spaces
A linear-time sequence model whose state update depends on the token it reads.
Attention compares every token with every earlier token and keeps a cache of the whole past. Mamba carries a fixed-size state from token to token, like an RNN, and computes from each token how much of the old state to keep and how much of the token to write in.
Explaining the paperMamba: Linear-Time Sequence Modeling with Selective State SpacesTrained on 256-token sequences, a 74K-parameter Mamba solves an in-context recall task perfectly on sequences of a million tokens. The attention models in the same test fall below 10% accuracy at 4,096.
A Transformer processes a sequence with attention: each token computes a score against every earlier token and takes a weighted average of their values. For a sequence of length that is about scores per layer, and to generate text the model stores the keys and values of every past token (the KV cache), so memory grows with the length of the context.
Mamba, by Albert Gu and Tri Dao (December 2023), is a recurrent model. It reads the sequence left to right and carries a state of fixed size, so its cost grows linearly in and it generates each new token in constant time and memory. Recurrent models had been tried many times. Mamba builds on the structured state space models of the S4 line, which train fast because their recurrence is also a convolution, and identifies their main weakness: the rule that updates the state is the same for every token, so the model cannot keep a name and skip a filler word. Mamba makes the update rule a function of the current token, then supplies a GPU kernel that computes the resulting recurrence quickly.
The sections below cover, in order: what attention and recurrences cost; what a state space model is; how its continuous-time form becomes a per-token update with a step size ; why the time-invariant version is a convolution; a task that time-invariant models fail; how Mamba's selection mechanism fixes it, and why it reduces to an RNN gate; the parallel scan and memory layout that make it fast; the full block; and the results.
What attention costs, and what a recurrence costs
In Figure 1 the left panel is causal attention: row holds the scores of token against tokens to , so the filled triangle has cells (78 at ). The right panel is a recurrence: one state passes from row to row, and each token does a fixed amount of work to update it, updates in total.
The memory difference matters more in deployment than the compute difference. Take two models of about 2.8 billion parameters. Pythia-2.8B (one of a suite of open Transformers trained on the Pile) has 32 layers of width 2560 and stores a key and a value of 2,560 numbers per layer for every token: 163,840 numbers per token, or 335 million for a 2,048-token context. Mamba-2.8B, with 64 layers, stores per layer a state of 5,120 channels by 16 numbers plus a 4-token convolution buffer: 6.55 million numbers in total, the same as Pythia's cache for 40 tokens, and it does not grow as the text gets longer. Because each sequence occupies so little memory, Mamba can generate for many more sequences at once on the same GPU, and the paper measures 4–5× the inference throughput of a Transformer of similar size.
Recurrent networks have two classic problems. A nonlinear RNN such as an LSTM must be trained step by step, because each state depends on the previous one through a nonlinearity, and that does not parallelize across the sequence on a GPU. And a fixed-size state has to compress the past, so the model needs a good rule for what to keep. Structured SSMs solved the first problem, and Mamba targets the second.
What a state space model is
A state space model (SSM) maps an input signal to an output through a hidden state that evolves by a linear differential equation. The term comes from control theory; the paper uses it only for the structured SSMs of the S4 family:
Here is a scalar, one channel of the input. is an matrix that sets how the state evolves on its own (in Mamba it is diagonal with negative entries, so each of the state numbers decays toward zero at its own rate). sets how the input is written into the state, and how the state is read out. The HiPPO theory showed that with a specific the state holds the coefficients of an optimal polynomial approximation of the input's history; Mamba keeps a simplified real diagonal version of it, for at initialization (S4D-Real), and learns it from there.
A layer has many channels, not one. Mamba runs an independent SSM on each of its inner channels ( is the model width, an expansion factor covered with the block), each with its own -number state. With the default and a width of 1,024, one layer's state is a 2,048 × 16 grid: 32,768 numbers that summarize everything the layer has read.
Discretization: how the step size Δ turns the ODE into a recurrence
Equation (1) runs in continuous time, but tokens arrive one at a time. To get an update from token to token , assume each token's value is held constant for a length of time and solve (1) exactly over that interval. This is the zero-order hold (ZOH): "zero order" because the input is approximated by a constant, a polynomial of degree zero, during each step. With a constant input a linear ODE has a closed-form solution, so ZOH is exact under its assumption. It produces discrete parameters and :
and the recurrence and readout
Work it for one state number with and . Equation (2) gives and , so the update is a weighted average of the old state and the new token:
At the weights are 0.905 on the old state and 0.095 on the token, so the state barely moves. At they are 0.368 and 0.632. At they are 0.050 and 0.950, and the state is close to a copy of the current token. is the amount of continuous time a token lasts: a long step lets the state relax all the way to the token's value and forget what it held, a short step leaves the state almost unchanged. The paper summarizes it the same way: a large resets the state and focuses on the current input, a small persists the state and ignores it.
The released code computes exactly but replaces with the first-order approximation (an Euler step). In the scalar example the two input weights are 0.095 against 0.1 at , and 0.632 against 1.0 at . Figure 2 shows the violet Euler dots drifting above the curve as grows. Mamba initializes every channel's step between 0.001 and 0.1, where the two agree closely, and the learned can absorb part of the difference; the Provenance panel at the end gives the file and line.
The convolution view, and why it needs fixed parameters
If , and are the same at every step, the model is time-invariant (linear time-invariant, LTI). Unroll (3) from :
The coefficient on is , which depends only on the distance , not on . Collected into a sequence, these coefficients form a convolution kernel:
For the scalar example with and , the kernel is 0.393, 0.239, 0.145, 0.088, 0.053, 0.032, …: the current token gets weight 0.393, the token five steps back gets 0.032 whether it is a name the model will need later or a filler word. The same model can therefore be run two ways. As the recurrence (3) it takes one step per token, which suits generation. As the convolution (4) it computes all outputs at once from the whole input, and with the Fast Fourier Transform that costs and runs in parallel across the sequence, which suits training. S4 used this equivalence to train on long sequences.
Both views are the same linear map, which can be written as a matrix. Stack the outputs into and the inputs into ; then with a lower-triangular whose entry is the weight of input in output :
The subscripts allow parameters that change per step, which the next sections need. In the time-invariant case every entry on a diagonal has the same distance and the same factors, so each diagonal is constant, and a matrix with constant diagonals is a convolution. In Figure 3, the right panel plots every entry of diagonal above position . In time-invariant mode they collapse onto one curve, the kernel. Switch to selective mode, where the step, and are computed from each token, and the entries of each diagonal spread apart: there is no kernel any more. The readout confirms that the recurrence and give the same output in both modes. (The Mamba-2 paper builds its whole algorithm on this matrix.)
Why time-invariant models fail at selective copying
The paper tests this with selective copying. Each input is 4,096 tokens long and contains 16 data tokens at random positions; every other position holds a single noise token, and the vocabulary has 16 symbols in total. After the sequence the model must output the 16 data tokens in order. In the plain copying task the data sits at fixed positions, and a convolution solves it by having a kernel of the right length; randomizing the positions removes that shortcut, so the model has to decide what to store from what each token is.
Figure 4 isolates the mechanism with a single state number. Four content tokens are separated by runs of filler, and each bar shows how much a token contributes to the final state: token enters with weight and is then multiplied by for every later token . With a fixed and 20 filler tokens after the last content token, that token's weight shrinks by , and the earlier content fades even more. A smaller fixed slows the fade, but each content token is then written in with a tiny weight, and the filler with the same weight.
The middle setting, "masked", sets the filler's input to zero, as an input gate would. The filler bars disappear but the content still fades, because every filler token still advances the decay. The paper makes the same point: the task can be seen as arriving already masked if noise tokens embed to zero, and a model that only masks inputs does not change the spacing between the relevant ones. Only the third setting, a token-dependent , removes the dependence on the gap. Filler gets , so and : it neither writes into the state nor decays it, and 20 filler tokens cost a factor of 0.98.
On the real task, an S4 layer with no gating reaches 18.3% accuracy. The paper's diagnosis is that LTI models lack content-based reasoning: attention computes a weight from each pair of tokens, while a fixed kernel assigns weights by position alone.
How Mamba's selection mechanism works
Mamba makes three of the SSM's parameters functions of the current input. At every position it computes a fresh , and from the token's vector with small linear maps:
and are linear projections to numbers. in the paper projects to a single number and broadcasts it to all the SSM's channels; the released code uses a rank- projection instead, (64 at width 1,024), lifted back to one step per channel. is the softplus, , which keeps the step positive, and "Parameter" is a learned per-channel bias. The paper calls this selective SSM S6. The shapes for a batch of sequences (Algorithm 2 of the paper, where counts the SSM's channels, = 2,048 in the block): is and does not depend on the input; and are , one per token, shared by all channels; is , one per token per channel; so the discretized and are .
Each parameter has a distinct job. sets, per channel, how much of the state the token resets and how much of the token goes in, the dial worked out in the discretization section. chooses which of the state slots the token writes into, and which slots the output reads from, so a token can store information in one slot and a later token can retrieve it. stays a fixed parameter because it reaches the recurrence only through : once varies per token, so does the decay. The paper names as the main source of the improvement and hypothesizes that a selective would perform about the same.
The connection to RNN gates is exact, and the paper proves it as Theorem 1. Take one state number, , with , , and where is a linear function of . Since , ZOH gives
and substituting into (3):
This is the gated update of a gated RNN such as a GRU, with the new input in place of the candidate state and the gate computed from the current input. Because the gate is computed from alone and never from the state, the recurrence stays linear in , which the scan section depends on. Work three tokens through it: a data token with and , a filler token with and , a data token with and . The steps are , 0.0067 and 3.05, and the gates 0.953, 0.0067 and 0.953. Starting from zero the state goes to 0.857, then 0.853 (the filler moved it by 0.004), then 0.802, mostly the new data token.
With the third token overwrote the first, because one number holds one value. Mamba has slots per channel with different decay rates ( at initialization), input-dependent and to route writes and reads between slots, and thousands of channels, each with its own . The paper's state-size ablation shows the extra slots pay off only with selection: with selective and , perplexity improves from 9.73 at to 8.71 at for a negligible number of extra parameters.
With selection, every position has its own , and , the diagonals of stop being constant, and the convolution and the FFT no longer apply. The model has to be computed as a recurrence, and a recurrence of length seems to need sequential steps.
Computing the selective scan fast on a GPU
For a single channel and state slot, each step of the selective recurrence is an affine map, , with and . Applying an earlier map and then a later map gives another affine map:
so a pair can stand for any run of consecutive steps, and combining two runs costs one multiply and one multiply-add. Composition of functions is associative, so the runs can be combined in any grouping. That is the condition for a parallel scan (a prefix sum with a general associative operator): combine neighbors in pairs, then pairs of pairs, and after rounds every prefix is known. For that is 11 rounds of independent work instead of 2,048 dependent steps. A classic RNN, , has no such shortcut: two tanh steps do not compose into one step of the same form, so there is nothing to combine.
Figure 6 draws the simplest version (Hillis–Steele), which does more total work than a sequential loop; the paper uses a work-efficient scan (Blelloch, 1990), which takes about twice as many rounds and does work in total.
The scan removes the sequential dependency but not the memory traffic. A GPU has a large, relatively slow main memory (HBM, tens of gigabytes) and a small fast on-chip memory (SRAM) next to the compute units. The discretized and have shape , 16 times larger than the input and output of shape . A straightforward implementation, such as the reference function in the repository, builds those tensors in HBM and reads them back during the scan. For one sequence of 2,048 tokens with 2,048 channels and , each of them is 67 million numbers.
Mamba's kernel is fused, in the manner of FlashAttention: it loads , , , and from HBM into SRAM, discretizes and runs the scan there, and writes only the output back. Sequences too long for SRAM are split into chunks, and the scan carries its state from one chunk to the next. The paper states the saving as a factor of in memory traffic and measures the scan at up to 20–40× faster than a standard scan implementation in PyTorch, and faster than FlashAttention-2 beyond sequence length 2,048.
Training also needs the intermediate states for the backward pass. Storing them would bring the tensor back into HBM, so the kernel recomputes them in the backward pass from the inputs, which costs less time than reading them from memory. With recomputation applied to the rest of the block as well, a selective SSM layer stores about 16 bytes of activations per token, and two Mamba layers use about the same activation memory as one FlashAttention layer plus one MLP layer of a Transformer (about 32 bytes).
The computation the kernel performs, written as a plain loop with shapes:
# selective SSM forward pass: the reference loop the fused kernel replaces
# x: (B, L, ED) after conv + SiLU. Learned, all receive gradients:
# A_log (ED, N), D (ED,), W_x (ED, R + 2N), W_dt (R, ED), b_dt (ED,)
A = -exp(A_log) # (ED, N), negative: the state decays
dt, Bm, Cm = split(x @ W_x, [R, N, N]) # (B,L,R) (B,L,N) (B,L,N)
dt = softplus(dt @ W_dt + b_dt) # (B, L, ED): one step per channel
h = zeros(B, ED, N) # the layer's whole memory
for t in range(L):
dA = exp(dt[:, t, :, None] * A) # (B, ED, N) exact
dB = dt[:, t, :, None] * Bm[:, t, None, :] # (B, ED, N) Euler, as shipped
# paper's ZOH: dB = (dA - 1) / A * Bm[:, t, None, :]
h = dA * h + dB * x[:, t, :, None] # (B, ED, N)
y[:, t] = (h * Cm[:, t, None, :]).sum(-1) # (B, ED)
y = y + D * x # skip connection, then y * silu(z)The Mamba block architecture
Earlier SSM architectures such as H3 (an S4 layer placed between two multiplicative gates) alternate an SSM block with an MLP block, the way a Transformer alternates attention with an MLP. Mamba merges the two into one block and stacks that block with normalization and residual connections. A block does five things:
- Expand and split.
in_projmaps width to and splits the result into two branches of width , called and . The paper fixes . - Mix nearby tokens. On the branch, a causal depthwise convolution of width 4 lets each channel combine the current token with the three before it, followed by a SiLU, . This is the short convolution H3 places before its SSM.
- Select. The selective SSM runs on that branch, computing , and from it by (5), plus a learned skip term (a per-channel parameter, initialized to 1).
- Gate. The branch goes through a SiLU and multiplies the SSM output elementwise, . With the SiLU, this gated half has the form of a SwiGLU MLP, the variant used in PaLM and LLaMA.
- Project back.
out_projmaps width back to , and the residual connection adds the block's input.
out_proj returns to width D inside a residual. Hover or tap any box for its shapes and weight count.At width the shapes for one 2,048-token sequence are below. and are per channel, so there are 2,048 independent SSMs with 16 state numbers each, while and are computed once per token and shared by all 2,048 channels.
# one Mamba block, batch 1, length L=2048, width D=1024, E=2, N=16
in_proj : (1, 2048, 1024) -> (1, 2048, 4096) # then split in two
x, z : each (1, 2048, 2048) # SSM branch / gate branch
conv1d : (1, 2048, 2048) width 4, causal, depthwise, then SiLU
x_proj : (1, 2048, 2048) -> (1, 2048, 96) # 64 for dt, 16 B, 16 C
dt : (1, 2048, 2048) # dt_proj 64 -> 2048, softplus: per channel
B, C : each (1, 2048, 16) # shared by all 2048 channels
A : (2048, 16) # fixed per channel, not input-dependent
state h : (1, 2048, 16) # carried across tokens, never grows with L
y : (1, 2048, 2048) -> y + D*x, then y * silu(z)
out_proj: (1, 2048, 2048) -> (1, 2048, 1024)Most of the weights sit in the two projections: in and out, so per block, which is 6.29 million at . The convolution and the SSM (the , , projections, and ) add about 0.37 million, for 6.67 million in total. A Transformer layer (attention plus MLP) has about weights, so the paper stacks two Mamba blocks for every Transformer layer it compares against. The language model is a token embedding, the stack of blocks, and a final normalization, with no attention, no separate MLP blocks and no positional encoding.
Results: synthetic tasks, language, DNA, audio, speed
Selective copying. Table 1 of the paper crosses three block designs (no gate, H3, Mamba) with three inner layers (S4; Hyena, a long convolution whose kernel an MLP generates; and S6). Most of the gain comes from changing the layer to S6: with no gating at all, S4 scores 18.3% and S6 97.0%; inside the H3 block, 57.0% becomes 99.7%; inside the Mamba block, 56.4% becomes 99.8%. Gated blocks help the non-selective layers (18.3% to 57.0% for S4), but none of them passes 57.0% without S6.
Induction heads. This task tests associative recall: if the sequence contained "Harry Potter" earlier, the next "Harry" should be followed by "Potter". Two-layer models are trained at length 256 with a vocabulary of 16 and tested from 64 up to 1,048,576 tokens. Mamba, with 74K parameters, scores 100% at every length, 4,000× its training length. No other model stays accurate beyond twice its training length: the best attention variant (xPos, a positional encoding designed for length extrapolation) gets 99.6% at 512 and 67.6% at 1,024, attention runs out of memory beyond 16,384, and H3 and Hyena, which are time-invariant, fall below 10% by 8,192.
Language modeling. On the Pile, with model sizes from about 125M to 1.3B parameters trained under the Chinchilla compute-optimal protocol, Mamba is the first attention-free model to match Transformer++, a strong modern recipe with rotary embeddings (RoPE), a SwiGLU MLP, RMSNorm, no linear biases and higher learning rates; the paper notes the match holds particularly as the sequence length grows. At 2.8B parameters, trained on 300B tokens with the same data and tokenizer as Pythia, Mamba-2.8B averages 63.3 on the paper's zero-shot suite (LAMBADA, HellaSwag, PIQA, ARC-E, ARC-C, WinoGrande), against 59.1 for Pythia-2.8B and 61.7 for Pythia-6.9B, and reaches Pile perplexity 6.22 against Pythia-2.8B's 6.73, which the abstract summarizes as matching Transformers twice its size.
DNA. Pretrained on the HG38 human genome, about 4.5 billion base pairs with one token per base, Mamba matches Transformer++ and HyenaDNA with roughly 3–4× fewer parameters at the largest size tested (about 40M), and its pretraining perplexity keeps improving as the context grows to a million bases. The paper attributes the long-context behavior to selection: a model that can reset or ignore can use more context without being confused by irrelevant parts of it.
Audio. On SC09, a benchmark of one-second spoken digits generated unconditionally, lower FID (Fréchet Inception Distance, a distance between generated and real samples) is better. A 6.1M-parameter Mamba scores 0.94, against 1.99 for the 5.8M SaShiMi (an S4-based audio model), 2.03 for the 19.1M GAN WaveGAN and 1.92 for the 24.1M diffusion model DiffWave. A 24.3M Mamba scores 0.67, the best in the table.
Speed. The fused scan with is faster than FlashAttention-2 beyond 2,048 tokens and up to 20–40× faster than a PyTorch scan, and end-to-end generation throughput is 4–5× that of a Transformer of similar size; the paper notes that an untrained Mamba-6.9B would have higher inference throughput than a Transformer-1.3B, five times smaller.
Limitations
The paper's discussion names three. First, selection gives up some of the bias toward smooth, uniformly sampled signals that made LTI SSMs strong on audio: in its ablation on long-form audio waveforms, replacing S4 with S6 inside the same block significantly hurt performance. Second, its language experiments stop at about 3B parameters, below the 7B scale at which open LLMs such as LLaMA and other recurrent models such as RWKV and RetNet had been evaluated, so whether Mamba stays competitive larger was left open. Third, it is not yet known whether Mamba supports the things Transformers are used for after pretraining, such as fine-tuning, prompting, in-context learning, RLHF and quantization, as well as Transformers do.
Two more follow from the design. The model depends on a custom CUDA kernel, since the FFT path that made S4 easy to implement is gone. And a recurrent state is a fixed amount of memory: per layer of Mamba-2.8B, 5,120 × 16 = 81,920 numbers, so a task that needs many arbitrary earlier tokens reproduced exactly can exceed it, where attention can look any token up in its cache. The paper presents this as the compression trade-off of all recurrent models (its Section 3.1); the hybrid experiments that measure it came later, in Mamba-2, where a 2.7B model with 6 attention layers among 64 reached lower perplexity than either pure model. Mamba-3 revisits the discretization and the state update itself.
Questions you might still have
If the state has a fixed size, does Mamba forget the distant past?
It keeps whatever fits in its state, which per layer is E·D channels times N = 16 numbers (81,920 numbers per layer for Mamba-2.8B). Selection decides what goes in: a token with a near-zero step neither writes nor decays the state, so filler costs nothing. On induction heads that was enough to recall the right token at a million positions. A task that needs many arbitrary past tokens verbatim can exceed that capacity; the paper frames this as the compression trade-off of any recurrent model, and later work, covered in the Mamba-2 explainer, found that mixing a few attention layers back in helps.
Why does making Δ, B and C depend on the input rule out the convolution and the FFT?
A convolution applies the same kernel at every position. The SSM is a convolution only when Ā, B̄ and C are identical at every step, because then the weight from token j to output t depends only on t minus j. Once they are computed from each token, the weight depends on every step between j and t, so each diagonal of the input-output matrix holds different values and there is no single kernel for the FFT to apply.
Why is A not computed from the input like Δ, B and C?
A reaches the recurrence only through the product ΔA inside Ā = exp(ΔA). Since Δ already changes per token, Ā already changes per token. The paper hypothesizes that a selective A would perform about the same and leaves it fixed for simplicity.
Does Mamba use positional encodings?
No. The official language model is a token embedding, a stack of Mamba blocks with normalization and residuals, and a final norm. Order enters through the recurrence and the causal convolution, which both process tokens left to right.
How does generation work, one token at a time?
Each layer keeps two small buffers per sequence: the last 4 inputs for the causal convolution and the (E·D, N) SSM state. A new token runs through the projections, updates the convolution buffer, computes its Δ, B and C, applies one step h = Ā·h + B̄·x and reads out y = C·h. The cost per token is the same at position 10 and at position 100,000.
How is Mamba-2 different?
Mamba-2 restricts Ā to one scalar per head instead of a diagonal per channel. That lets the whole sequence computation be written as a structured matrix and run mostly with matrix multiplications, which GPUs execute far faster than a scan, and it allows a larger state (N = 128 against 16). The Mamba-2 explainer covers that derivation.
Footnotes & further reading
- The paper: Gu, Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Carnegie Mellon & Princeton, COLM 2024). Code.
- The structured SSM Mamba builds on, with the convolution view and the diagonal initialization: Gu, Goel, Ré, Efficiently Modeling Long Sequences with Structured State Spaces (S4) (see the S4 explainer), and Gu et al., On the Parameterization and Initialization of Diagonal State Space Models (S4D).
- The argument that an SSM state can optimally compress history: Gu et al., HiPPO: Recurrent Memory with Optimal Polynomial Projections.
- The IO-aware kernel design Mamba's scan follows: Dao et al., FlashAttention (see the FlashAttention explainer).
- The block fuses the SSM layer of Fu et al., Hungry Hungry Hippos (H3), with a gated MLP of the SwiGLU family.
- The parallel prefix sum: Blelloch, Prefix Sums and Their Applications (CMU, 1990).
- The reference selective scan and block defaults are in the official repository, mamba_simple.py and selective_scan_interface.py.
- The follow-ups: Dao & Gu, Transformers are SSMs (Mamba-2) (see the Mamba-2 explainer), and the Mamba-3 explainer.
How could this explainer be improved? Found an error, or something unclear? I read every message.