VerifiedarXiv:2312.0075233 min
Architecture · Sequence models

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 SpacesGu, Dao · CMU & Princeton · COLM 2024 · arXiv:2312.00752 ↗

Trained 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 LL that is about L2/2L^2/2 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 LL 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 Δ\Delta; 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 ii holds the scores of token ii against tokens 11 to ii, so the filled triangle has L(L+1)/2L(L+1)/2 cells (78 at L=12L = 12). 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, LL updates in total.

Figure 1 · two ways to read a sequence
L = 12
Attention fills a triangle of token-to-token scores and keeps one cache entry per token. A recurrence passes one fixed-size state down the column. Drag L to see the triangle grow quadratically while the column grows linearly.

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 x(t)x(t) to an output y(t)y(t) through a hidden state h(t)∈RNh(t) \in \mathbb{R}^N 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:

h′(t)=A h(t)+B x(t),y(t)=C h(t)h'(t) = \mathbf{A}\,h(t) + \mathbf{B}\,x(t), \qquad y(t) = \mathbf{C}\,h(t)
(1)

Here x(t)x(t) is a scalar, one channel of the input. A\mathbf{A} is an N×NN \times N matrix that sets how the state evolves on its own (in Mamba it is diagonal with negative entries, so each of the NN state numbers decays toward zero at its own rate). B∈RN×1\mathbf{B} \in \mathbb{R}^{N \times 1} sets how the input is written into the state, and C∈R1×N\mathbf{C} \in \mathbb{R}^{1 \times N} how the state is read out. The HiPPO theory showed that with a specific A\mathbf{A} the state holds the coefficients of an optimal polynomial approximation of the input's history; Mamba keeps a simplified real diagonal version of it, An=−n\mathbf{A}_n = -n for n=1,…,Nn = 1, \dots, N 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 E⋅DE \cdot D inner channels (DD is the model width, E=2E = 2 an expansion factor covered with the block), each with its own NN-number state. With the default N=16N = 16 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 t−1t-1 to token tt, assume each token's value is held constant for a length of time Δ\Delta 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 Aˉ\bar{\mathbf{A}} and Bˉ\bar{\mathbf{B}}:

Aˉ=exp⁡(ΔA),Bˉ=(ΔA)−1(exp⁡(ΔA)−I)⋅ΔB\bar{\mathbf{A}} = \exp(\Delta\mathbf{A}), \qquad \bar{\mathbf{B}} = (\Delta\mathbf{A})^{-1}\big(\exp(\Delta\mathbf{A}) - \mathbf{I}\big)\cdot\Delta\mathbf{B}
(2)

and the recurrence and readout

ht=Aˉ ht−1+Bˉ xt,yt=C hth_t = \bar{\mathbf{A}}\,h_{t-1} + \bar{\mathbf{B}}\,x_t, \qquad y_t = \mathbf{C}\,h_t
(3)

Work it for one state number with A=−1\mathbf{A} = -1 and B=1\mathbf{B} = 1. Equation (2) gives Aˉ=e−Δ\bar{A} = e^{-\Delta} and Bˉ=1−e−Δ\bar{B} = 1 - e^{-\Delta}, so the update is a weighted average of the old state and the new token:

ht=e−Δ ht−1+(1−e−Δ) xth_t = e^{-\Delta}\, h_{t-1} + \big(1 - e^{-\Delta}\big)\, x_t

At Δ=0.1\Delta = 0.1 the weights are 0.905 on the old state and 0.095 on the token, so the state barely moves. At Δ=1\Delta = 1 they are 0.368 and 0.632. At Δ=3\Delta = 3 they are 0.050 and 0.950, and the state is close to a copy of the current token. Δ\Delta 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 Δ\Delta resets the state and focuses on the current input, a small Δ\Delta persists the state and ignores it.

Figure 2 · zero-order hold
Δ = 0.50
Ā = 0.606 · B̄ = 0.394
Ten tokens, each held for Δ units of time. The exact solution of (1) relaxes toward each held value, and the ZOH recurrence (3) lands on it at every token boundary. Slide Δ from 0.01 (state frozen) to 5 (state snaps to each token). Tick the box to overlay the code's Euler B̄ = Δ·B, which overshoots once Δ is large.

The released code computes Aˉ=exp⁡(ΔA)\bar{\mathbf{A}} = \exp(\Delta\mathbf{A}) exactly but replaces Bˉ\bar{\mathbf{B}} with the first-order approximation ΔB\Delta\mathbf{B} (an Euler step). In the scalar example the two input weights are 0.095 against 0.1 at Δ=0.1\Delta = 0.1, and 0.632 against 1.0 at Δ=1\Delta = 1. Figure 2 shows the violet Euler dots drifting above the curve as Δ\Delta grows. Mamba initializes every channel's step between 0.001 and 0.1, where the two agree closely, and the learned B\mathbf{B} 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 Aˉ\bar{\mathbf{A}}, Bˉ\bar{\mathbf{B}} and C\mathbf{C} are the same at every step, the model is time-invariant (linear time-invariant, LTI). Unroll (3) from h0=0h_0 = 0:

yt=CBˉ xt+CAˉBˉ xt−1+CAˉ2Bˉ xt−2+⋯y_t = \mathbf{C}\bar{\mathbf{B}}\,x_t + \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}}\,x_{t-1} + \mathbf{C}\bar{\mathbf{A}}^2\bar{\mathbf{B}}\,x_{t-2} + \cdots

The coefficient on xt−kx_{t-k} is CAˉkBˉ\mathbf{C}\bar{\mathbf{A}}^k\bar{\mathbf{B}}, which depends only on the distance kk, not on tt. Collected into a sequence, these coefficients form a convolution kernel:

Kˉ=(CBˉ, CAˉBˉ, …, CAˉkBˉ, …),y=x∗Kˉ\bar{\mathbf{K}} = \big(\mathbf{C}\bar{\mathbf{B}},\ \mathbf{C}\bar{\mathbf{A}}\bar{\mathbf{B}},\ \dots,\ \mathbf{C}\bar{\mathbf{A}}^k\bar{\mathbf{B}},\ \dots\big), \qquad y = x * \bar{\mathbf{K}}
(4)

For the scalar example with Δ=0.5\Delta = 0.5 and C=1\mathbf{C} = 1, 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 O(Llog⁡L)O(L \log L) 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 yy and the inputs into xx; then y=Mxy = Mx with a lower-triangular MM whose entry (t,j)(t, j) is the weight of input jj in output tt:

Mtj=Ct AˉtAˉt−1⋯Aˉj+1 Bˉj(j≤t)M_{tj} = \mathbf{C}_t\,\bar{\mathbf{A}}_t \bar{\mathbf{A}}_{t-1} \cdots \bar{\mathbf{A}}_{j+1}\,\bar{\mathbf{B}}_j \quad (j \le t)

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 t−jt - j 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 kk above position kk. In time-invariant mode they collapse onto one curve, the kernel. Switch to selective mode, where the step, B\mathbf{B} and C\mathbf{C} 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 MxMx give the same output in both modes. (The Mamba-2 paper builds its whole algorithm on this matrix.)

Figure 3 · the SSM as a matrix
Δ = 0.50
hover or tap a cell of Mwidest diagonal spread 0 · max |y_rec − Mx| < 1e-15
Left: M for a 12-token input, brighter teal for larger weights. Right: the entries of each diagonal, with the kernel K̄ drawn through them in time-invariant mode. Hover a cell to highlight its diagonal; toggle the model and drag Δ.

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 ii enters with weight Bˉi\bar{B}_i and is then multiplied by Aˉj=e−Δj\bar{A}_j = e^{-\Delta_j} for every later token jj. With a fixed Δ=0.5\Delta = 0.5 and 20 filler tokens after the last content token, that token's weight shrinks by e−0.5⋅20≈4.5×10−5e^{-0.5 \cdot 20} \approx 4.5 \times 10^{-5}, and the earlier content fades even more. A smaller fixed Δ\Delta 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 Δ\Delta, removes the dependence on the gap. Filler gets Δ≈0.001\Delta \approx 0.001, so Aˉ≈0.999\bar{A} \approx 0.999 and Bˉ≈0.001\bar{B} \approx 0.001: it neither writes into the state nor decays it, and 20 filler tokens cost a factor of 0.98.

Figure 4 · what survives in the state
20 tokens
Δ = 0.50
content weights 1–4: 4e‑19 · 3e‑13 · 2e‑10 · 2e‑5content total 2e‑5 · share 0.0%
Each bar is a token's weight in the final state (log scale), content or filler. Switch how filler is handled, then drag the gap. With a fixed Δ the content total and share collapse as the gap grows; masking keeps the share at 100% but the total still collapses; a selective Δ keeps both.

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 Δt\Delta_t, Bt\mathbf{B}_t and Ct\mathbf{C}_t from the token's vector xtx_t with small linear maps:

Bt=sB(xt),Ct=sC(xt),Δt=τΔ(Parameter+sΔ(xt))\mathbf{B}_t = s_B(x_t),\quad \mathbf{C}_t = s_C(x_t),\quad \Delta_t = \tau_\Delta\big(\text{Parameter} + s_\Delta(x_t)\big)Bt=sB(xt),Ct=sC(xt),Δt=τΔ(Parameter+sΔ(xt))\begin{gathered} \mathbf{B}_t = s_B(x_t),\quad \mathbf{C}_t = s_C(x_t), \\ \Delta_t = \tau_\Delta\big(\text{Parameter} + s_\Delta(x_t)\big) \end{gathered}
(5)

sBs_B and sCs_C are linear projections to NN numbers. sΔs_\Delta in the paper projects to a single number and broadcasts it to all the SSM's channels; the released code uses a rank-RR projection instead, R=⌈D/16⌉R = \lceil D/16 \rceil (64 at width 1,024), lifted back to one step per channel. τΔ\tau_\Delta is the softplus, log⁡(1+ez)\log(1 + e^z), 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 BB sequences (Algorithm 2 of the paper, where DD counts the SSM's channels, E⋅DE \cdot D = 2,048 in the block): A\mathbf{A} is (D,N)(D, N) and does not depend on the input; B\mathbf{B} and C\mathbf{C} are (B,L,N)(B, L, N), one per token, shared by all channels; Δ\Delta is (B,L,D)(B, L, D), one per token per channel; so the discretized Aˉ\bar{\mathbf{A}} and Bˉ\bar{\mathbf{B}} are (B,L,D,N)(B, L, D, N).

Each parameter has a distinct job. Δt\Delta_t 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. Bt\mathbf{B}_t chooses which of the NN state slots the token writes into, and Ct\mathbf{C}_t which slots the output reads from, so a token can store information in one slot and a later token can retrieve it. A\mathbf{A} stays a fixed parameter because it reaches the recurrence only through Aˉt=exp⁡(ΔtA)\bar{\mathbf{A}}_t = \exp(\Delta_t\mathbf{A}): once Δt\Delta_t varies per token, so does the decay. The paper names Δ\Delta as the main source of the improvement and hypothesizes that a selective A\mathbf{A} would perform about the same.

The connection to RNN gates is exact, and the paper proves it as Theorem 1. Take one state number, N=1N = 1, with A=−1\mathbf{A} = -1, B=1\mathbf{B} = 1, and Δt=softplus(zt)\Delta_t = \text{softplus}(z_t) where ztz_t is a linear function of xtx_t. Since e−softplus(z)=1/(1+ez)e^{-\text{softplus}(z)} = 1/(1 + e^{z}), ZOH gives

Aˉt=11+ezt=1−σ(zt),Bˉt=e−Δt−1−1=1−Aˉt=σ(zt)\bar{A}_t = \frac{1}{1 + e^{z_t}} = 1 - \sigma(z_t), \qquad \bar{B}_t = \frac{e^{-\Delta_t} - 1}{-1} = 1 - \bar{A}_t = \sigma(z_t)Aˉt=11+ezt=1−σ(zt),Bˉt=e−Δt−1−1=1−Aˉt=σ(zt)\begin{aligned} \bar{A}_t &= \frac{1}{1 + e^{z_t}} = 1 - \sigma(z_t), \\ \bar{B}_t &= \frac{e^{-\Delta_t} - 1}{-1} = 1 - \bar{A}_t = \sigma(z_t) \end{aligned}

and substituting into (3):

gt=σ(Linear(xt)),ht=(1−gt) ht−1+gt xtg_t = \sigma\big(\text{Linear}(x_t)\big), \qquad h_t = (1 - g_t)\,h_{t-1} + g_t\,x_t
(6)

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 xtx_t alone and never from the state, the recurrence stays linear in hh, which the scan section depends on. Work three tokens through it: a data token with x1=0.9x_1 = 0.9 and z1=3z_1 = 3, a filler token with x2=0.2x_2 = 0.2 and z2=−5z_2 = -5, a data token with x3=0.8x_3 = 0.8 and z3=3z_3 = 3. The steps are Δ=3.05\Delta = 3.05, 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.

Figure 5 · Δ as an RNN gate (Theorem 1)
w = 8.0
b = -4.00
token 3: x = 0.90, z = 3.20, Δ = softplus(z) = 3.240, g = 1 − e^−Δ = 0.961
Tokens and the state of (6), with each token's gate as a bar. The gate is σ(w·x + b), computed from the token alone. At w = 0 every token gets the same gate, which is the time-invariant case; raise w and the state jumps to large tokens and holds through small ones. Hover a token for its z, Δ and g.

With N=1N = 1 the third token overwrote the first, because one number holds one value. Mamba has N=16N = 16 slots per channel with different decay rates (An=−n\mathbf{A}_n = -n at initialization), input-dependent Bt\mathbf{B}_t and Ct\mathbf{C}_t to route writes and reads between slots, and thousands of channels, each with its own Δ\Delta. The paper's state-size ablation shows the extra slots pay off only with selection: with selective B\mathbf{B} and C\mathbf{C}, perplexity improves from 9.73 at N=1N = 1 to 8.71 at N=16N = 16 for a negligible number of extra parameters.

With selection, every position has its own Aˉt\bar{\mathbf{A}}_t, Bˉt\bar{\mathbf{B}}_t and Ct\mathbf{C}_t, the diagonals of MM 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 LL seems to need LL 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, h↦ath+bth \mapsto a_t h + b_t, with at=Aˉta_t = \bar{A}_t and bt=Bˉtxtb_t = \bar{B}_t x_t. Applying an earlier map (a1,b1)(a_1, b_1) and then a later map (a2,b2)(a_2, b_2) gives another affine map:

h  ↦  a2(a1h+b1)+b2=(a2a1) h+(a2b1+b2)h \;\mapsto\; a_2(a_1 h + b_1) + b_2 = (a_2 a_1)\,h + (a_2 b_1 + b_2)

so a pair (a,b)(a, b) 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 log⁡2L\log_2 L rounds every prefix (a1⋯at, ht)(a_1 \cdots a_t,\ h_t) is known. For L=2048L = 2048 that is 11 rounds of independent work instead of 2,048 dependent steps. A classic RNN, ht=tanh⁡(Wht−1+Uxt)h_t = \tanh(W h_{t-1} + U x_t), 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 · the parallel scan
round 3/3
round 3 of 3 · combines so far 17 · finished h_t: 8 of 8
Eight tokens with their (a, b) pairs. Each node shows the b of the run it covers; teal nodes cover a run starting at token 1, so their b is the finished state h_t. Play or scrub through the 3 rounds, then switch to the 8 sequential steps; both end with the same states. Hover a node for its pair.

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 O(L)O(L) 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 Aˉ\bar{\mathbf{A}} and Bˉx\bar{\mathbf{B}}x have shape (B,L,D,N)(B, L, D, N), 16 times larger than the input and output of shape (B,L,D)(B, L, D). 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 N=16N = 16, each of them is 67 million numbers.

Mamba's kernel is fused, in the manner of FlashAttention: it loads Δ\Delta, A\mathbf{A}, B\mathbf{B}, C\mathbf{C} and xx from HBM into SRAM, discretizes and runs the scan there, and writes only the (B,L,D)(B, L, D) 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 O(N)O(N) 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.

Figure 7 · memory traffic of one scan
L = 2,048
N = 16
one (B, L, D, N) tensor: 67.1M numbers (0.27 GB at 4 bytes)materialized / fused = 22.2×
Numbers moved through GPU main memory for one selective scan with 2,048 channels. Materialized: write and read back Ā and B̄x. Fused: read the inputs, write the output. Drag N and watch the ratio track it.

Training also needs the intermediate states for the backward pass. Storing them would bring the (B,L,D,N)(B, L, D, N) 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:

Figure 8 · the Mamba block
hover or tap a box for its shapes and weight count (D = 1024, L = 2048)
The SSM branch (Conv1d, SiLU, selective SSM) and the gate branch multiply, then out_proj returns to width D inside a residual. Hover or tap any box for its shapes and weight count.

At width D=1024D = 1024 the shapes for one 2,048-token sequence are below. A\mathbf{A} and Δ\Delta are per channel, so there are 2,048 independent SSMs with 16 state numbers each, while B\mathbf{B} and C\mathbf{C} 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: 2ED22ED^2 in and ED2ED^2 out, so 3ED2=6D23ED^2 = 6D^2 per block, which is 6.29 million at D=1024D = 1024. The convolution and the SSM (the Δ\Delta, B\mathbf{B}, C\mathbf{C} projections, A\mathbf{A} and DD) add about 0.37 million, for 6.67 million in total. A Transformer layer (attention plus MLP) has about 12D212D^2 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.

Figure 9 · the synthetic tasks
No gate block + S4: 18.3% · same block with S6: 97.0 (+78.7 from the layer alone)
Numbers from the paper's tables. In selective copying, click a cell to compare it with the same block using S6. In induction heads, hover a test length to read every model's accuracy; the dashed line marks the training length 256.

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 N=16N = 16 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.

Provenance Verified against primary literatureHow we verify
Gu & Dao (2023), COLM 2024The selective SSM (S6) with Δ, B and C computed from the input, the hardware-aware fused scan, and the Mamba block. This page's equations (1) to (6) are the paper's (1), (4), (2), (3), the S6 definition in Sec. 3.2 / Algorithm 2, and Theorem 1.
mamba_simple.py (official code)d_state=16, d_conv=4 (depthwise, causal), expand=2, dt_rank=ceil(d_model/16), A = -exp(A_log) with A_log = log(1..N) (S4D-Real), D skip initialized to ones, output y * silu(z).
Delta initialization, paper vs codeThe paper sets the Δ bias to softplus⁻¹ of Uniform([0.001, 0.1]); the code samples Δ log-uniformly in [0.001, 0.1] (mamba_simple.py:92-95) before inverting the softplus. Both start every channel with a step between 0.001 and 0.1.
selective_scan_ref and the CUDA kernelĀ = exp(ΔA) computed exactly; B̄ = Δ·B, the Euler form (selective_scan_interface.py:162 and :167, selective_scan_fwd_kernel.cuh:221-222). See the flag below.
S4 / S4D (Gu et al., 2022)The structured SSM, the convolution view trained with the FFT, and the diagonal real initialization A_n = -(n+1) (paper indexing from 0) that Mamba keeps.
HiPPO (Gu et al., 2020)The theory that a particular A makes the SSM state an optimal compression of the input history.
Blelloch (1990)The work-efficient parallel prefix sum the paper cites for its scan. Figure 6 on this page draws the simpler Hillis–Steele scan, which takes the same log₂L rounds with more total work.
Mamba-2 (Dao & Gu, 2024)Later work, not this paper: at 2.7B parameters a hybrid with 6 attention layers out of 64 reached Pile perplexity 5.95, against 6.09 for pure Mamba-2 and 6.13 for Transformer++.
correctionThe paper defines B̄ by the zero-order-hold formula B̄ = (ΔA)⁻¹(exp(ΔA) − I)·ΔB (its Eq. 4) and never mentions an approximation. The released reference scan and the CUDA kernel both compute B̄ = Δ·B, the first-order Euler form (deltaB_u = einsum(delta, B, u) in selective_scan_ref); Ā = exp(ΔA) is exact in both. The two agree to first order in Δ: with A = −1 they give B̄ = 0.0952 vs 0.1 at Δ = 0.1, and 0.632 vs 1.0 at Δ = 1. Theorem 1, which turns the recurrence into an RNN gate with g = σ(z), relies on the ZOH form; with the Euler form the input coefficient is softplus(z) instead of σ(z), which matches only while Δ is small. This page teaches ZOH and shows the code's Euler form in Figure 2 and the listing.

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

  1. The paper: Gu, Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Carnegie Mellon & Princeton, COLM 2024). Code.
  2. 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).
  3. The argument that an SSM state can optimally compress history: Gu et al., HiPPO: Recurrent Memory with Optimal Polynomial Projections.
  4. The IO-aware kernel design Mamba's scan follows: Dao et al., FlashAttention (see the FlashAttention explainer).
  5. The block fuses the SSM layer of Fu et al., Hungry Hungry Hippos (H3), with a gated MLP of the SwiGLU family.
  6. The parallel prefix sum: Blelloch, Prefix Sums and Their Applications (CMU, 1990).
  7. The reference selective scan and block defaults are in the official repository, mamba_simple.py and selective_scan_interface.py.
  8. The follow-ups: Dao & Gu, Transformers are SSMs (Mamba-2) (see the Mamba-2 explainer), and the Mamba-3 explainer.