Mamba-2: Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Mamba's recurrence, rewritten as a structured matrix so a GPU can compute it with matrix multiplications.

Tri Dao and Albert Gu restrict the state space model's per-step decay to one number, show that the restricted model is a form of masked attention without the softmax, and use that to compute it in chunks. The new layer is 2 to 8 times faster than Mamba's scan at state size 64 and supports states 8 times larger.

At 2.7B parameters and 300B tokens of the Pile, Mamba-2 beats Pythia-6.9B, trained on the same data, on every benchmark in the paper's appendix table. This page builds the model from the recurrence up: the matrix an SSM computes, why its blocks are low rank, the one restriction that turns it into masked attention, and the chunked algorithm that follows.

Mamba showed that a recurrent layer with a small, input-dependent state could match Transformers on language modeling. Mamba-2 keeps that layer's shape (a fixed-size state updated once per token) and changes how it is computed. The paper's argument has four steps, and the page follows them in order:

1. Any state space model (SSM) computes $y = Mx$ for a lower-triangular $T \times T$ matrix $M$, where $T$ is the sequence length.
2. That matrix is semiseparable: every block below its diagonal has rank at most $N$, the state size.
3. If the state decays by a single scalar per step, $M$ factors into a decay mask times an attention-shaped matrix $CB^{\top}$. This is the structured state space duality (SSD) of the title.
4. Cut $M$ into chunks: compute the diagonal chunks like attention, and carry everything else through a short recurrence between chunks. Nearly all of the work becomes matrix multiplication.

Why Mamba was slow on a GPU

A selective SSM reads a sequence $x_0, \ldots, x_{T-1}$ one token at a time and keeps a state vector $h_t \in \mathbb{R}^N$. At each token it decays the old state, adds the new input, and reads out one number:

This is the paper's equation (2). Here $x_t$ and $y_t$ are scalars (one channel), $B_t, C_t \in \mathbb{R}^N$ are computed from the input at token $t$, and $A_t$ is an $N \times N$ matrix that also depends on the token. “Selective” refers to that input dependence: the model can choose, token by token, how much of the past to keep (through $A_t$) and how much of the present to write (through $B_t$). In Mamba, $A_t$ is diagonal, so each of the $N$ state slots decays at its own rate, and the layer runs one such SSM per channel. Mamba-2.8B has 5120 inner channels per layer and $N = 16$, so each layer carries 5120 × 16 = 81,920 state numbers from one token to the next.

Each step of (1) is a handful of elementwise multiply-adds on vectors of length 16. A GPU does that kind of arithmetic on its general-purpose cores. Matrix multiplications run on separate tensor cores, which are far faster: an A100 is rated at 312 teraFLOP/s for 16-bit matrix multiplication on tensor cores and 19.5 teraFLOP/s for ordinary 32-bit arithmetic, a factor of 16.⁶ Mamba's scan does work proportional to $T$, but none of that work is a matrix multiplication, so it cannot use the tensor cores. It needed a custom fused CUDA kernel to be efficient, and even with that kernel the paper counts it as less efficient than models built on matrix multiplication.

Attention has the opposite profile. Its cost grows like $T^2$, but almost all of it is two large matrix multiplications, and FlashAttention showed how to run those at close to the tensor cores' peak by tiling them through on-chip memory. Mamba-2 asks whether an SSM can be rewritten so that its work is also matrix multiplication, while keeping the cost linear in $T$.

An SSM is a triangular matrix

Unroll (1) from $h_0 = B_0 x_0$. Each input $x_s$ enters the state through $B_s$ and is then multiplied by every later decay:

Multiplying by $C_t^{\top}$ gives each output as a weighted sum of the inputs so far. Renaming the output index $j$ and the input index $i$:

So $y = Mx$ with a $T \times T$ matrix $M$. It is lower triangular because output $j$ depends only on inputs $i \le j$. For three tokens it reads

The diagonal entry $C_j^{\top} B_j$ is how much token $j$ contributes to its own output. The entry $M_{ji}$ below the diagonal (row $j$, column $i$) routes input $i$ through the state and through the decays $A_{i+1}, \ldots, A_j$ before it is read out. The recurrence and the product $Mx$ compute the same numbers (the paper's Theorem 3.5); they are two algorithms for one linear map. The recurrence costs work proportional to $T$ and never forms $M$. Forming $M$ and multiplying costs work proportional to $T^2$, which is what attention pays.

When a layer has more than one channel, the paper shares one $M$ across the $P$ channels of a head: the input is a $T \times P$ matrix $X$, and the output is $Y = MX$. In Mamba each head is one channel ($P = 1$); Mamba-2 uses $P = 64$ or $128$, like a Transformer's head dimension. The larger $P$ turns the vector operations in the recurrence into small matrix operations, which the chunked algorithm below relies on.

Why the off-diagonal blocks of an SSM matrix are low rank

Take any block of $M$ that sits below the diagonal: rows $r \ge j$ and columns $c \le i - 1$, for some split with $j \ge i$. Every entry in it is $C_r^{\top} A_r \cdots A_{c+1} B_c$, and the chain of decays passes through the block's corner. Cut the chain there:

The left factor $u_r$ is an $N$-vector that depends only on the row, the right factor $v_c$ is an $N$-vector that depends only on the column, and $W$ is the same $N \times N$ matrix for the whole block. Stacking the rows gives the block as $U W V$, with $U$ of width $N$ and $V$ of height $N$, so its rank is at most $N$ however large the block is. In words: every contribution of input $c$ to output $r$ passes through the $N$-dimensional state at the block's boundary.

A lower-triangular matrix in which every block on or below the diagonal has rank at most $N$ is called N-semiseparable (Definition 3.1), and the form $M_{ji} = C_j^{\top} A_j \cdots A_{i+1} B_i$ is its sequentially semiseparable (SSS) representation (Definition 3.2). Equation (3) proves one direction (Lemma 3.3). The converse is a known result (Proposition 3.4): every $N$-semiseparable matrix can be written this way. So SSMs with state size $N$ and $N$-semiseparable matrices are the same set of linear maps. The SSS form stores $O(N^2 T)$ numbers (one $N \times N$ matrix per token), but these matrices can be stored in $O(NT)$ numbers and multiplied by a vector in $O(NT)$ time (Proposition 3.6, from Pernet, Signargout and Villard).⁴

Because of this equivalence, any method for multiplying a vector by a semiseparable matrix is also a way to run an SSM. The recurrence is one such method. The chunked method of the SSD section is another, and it uses the factorization (3) directly.

Mamba-2's restriction: one decay per step

With a diagonal $A_t$, the chain $A_j \cdots A_{i+1}$ in (2) is a diagonal matrix whose $n$-th entry is the product of slot $n$'s decays. In Mamba those decays are $a_t[n] = \exp(\Delta_t A_n)$, with a step size $\Delta_t$ computed from the token and one learned negative rate $A_n$ per slot, so each slot has its own cumulative decay between any two tokens. Writing (2) out slot by slot:

Each slot contributes its own decay pattern, so $M$ is a sum of $N$ separately masked rank-one matrices, and no single mask factors out of the sum.

Mamba-2 requires $A_t = a_t I$: a scalar times the identity, the same decay for every slot at token $t$. Then $A_j \cdots A_{i+1} = (a_j \cdots a_{i+1}) I$, the scalar comes out of the dot product, and

or, for the whole matrix, $M = L \circ (C B^{\top})$, where $\circ$ multiplies entry by entry, $C$ and $B$ are $T \times N$ with rows $C_j^{\top}$ and $B_i^{\top}$, and the diagonal of $L$ is 1 (an empty product). The mask $L$ is itself 1-semiseparable: it is $M$ for a state of size 1. With decays $a_1, a_2, a_3 = 0.9, 0.5, 0.8$ on four tokens:

The bottom-left entry,

is the fraction of whatever token 0 wrote into the state that is still there at token 3. Multiplying a vector by $L$ is the scalar recurrence $y_t = a_t y_{t-1} + x_t$ (the paper's equation (7)).

In the official code, each head has one learned rate $A = -\exp(A_{\log}) < 0$ and each token a step size $\Delta_t > 0$, so $a_t = \exp(\Delta_t A)$ lies strictly between 0 and 1. The paper describes the $a_t$ as “input-dependent scalars bounded in [0, 1]”. A value near 1 keeps the state; a value near 0 clears it.

The last control in Figure 2 is how Mamba-2 handles several short sequences packed into one long row for training (Section 8.3). The decay that carries the state from token $t$ into token $t+1$ is $a_{t+1}$, and every $L_{ji}$ with $i \le t < j$ contains it as a factor. Setting $a_{t+1} = 0$ on the first token of the next sequence therefore zeroes the whole rectangle of $L$ that would leak one sequence into the next, with no padding and no special mask.

Why not keep the diagonal $A_t$? The paper notes that a diagonal $A$ has the same theoretical efficiency but is “less hardware-friendly”: with one mask per slot, the masked-attention computation of the next section would need $N$ separate masked $T \times T$ products instead of one masked matrix multiplication. The restriction costs expressiveness: the paper describes it as “slightly decreasing the expressive power in return for significant training efficiency improvements”, and it does not run an experiment that isolates the cost (the section on limits returns to this).

The same matrix as masked attention

Softmax attention computes $Y = \operatorname{softmax}(QK^{\top})\,V$ with a query, key and value per token. Linear attention (Katharopoulos et al., 2020) drops the softmax and applies the causal mask directly, $Y = (L \circ QK^{\top})V$ with $L$ the all-ones lower-triangular matrix, and shows that this product can be computed as a recurrence.³ Compare with (4). Renaming $Q = C$, $K = B$ and $V = X$:

So the scalar-decay SSM is masked attention with no softmax and with the causal mask replaced by the decay mask $L$. Linear attention is the case $a_t = 1$ for every token (nothing decays). RetNet uses a fixed decay per head, and GateLoop makes the decay input-dependent; the paper places both in the same family. It calls the general form $(L \circ QK^{\top})V$ with a structured $L$ structured masked attention, and proves the converse direction too (Theorem 5.2): any structured masked attention that can be computed as a recurrence with a bounded state must have a semiseparable mask. The attention analogy also fixes the roles of the sizes: the state size $N$ plays the part of the query-key dimension, and the head dimension $P$ the value dimension.

The two sides of (5) are two ways to compute the same outputs, with opposite costs. With $P = N$:

• Linear (recurrent) mode. Keep an $N \times P$ state per head and update it once per token: $h_t = a_t h_{t-1} + B_t x_t^{\top}$, $y_t = h_t^{\top} C_t$. Work is proportional to $TN^2$, the state is $N^2$ numbers, and each step is an outer product and a small matrix-vector product, done one token after another.
• Quadratic (attention) mode. Form $CB^{\top}$ (work $T^2 N$), mask it with $L$, multiply by $X$ (work $T^2 P$). All matrix multiplication, all tokens in parallel, and $T^2$ memory for the mask and scores.

Figure 3 runs a six-token example both ways. At token 3 the recurrent mode takes the previous state $h_2$, the decay 0.8, the input 0.5 with $B_3$ and $C_3$, and the attention mode takes row 3 of $M$ against the first four inputs:

The first entry of that row is $L_{30}$ times $C_3^{\top} B_0$, 0.36 × 0.25, with the same $L_{30}$ as in the example above.

Neither mode suits a GPU at long sequence lengths. The recurrent mode has the right total work but runs as a long chain of small operations. The attention mode is all tensor-core work, but its work grows like $T^2$. The SSD algorithm uses attention mode inside short chunks and recurrent mode between them.

How the SSD algorithm computes Mamba-2

Cut the sequence into chunks of $Q$ tokens, and cut $M$ into a grid of $Q \times Q$ blocks to match: $T/Q$ blocks on the diagonal and the rest below it. A diagonal block describes how the tokens inside one chunk affect each other. A block below the diagonal describes how an earlier chunk affects a later one, and by (3) it factors through the $N \times P$ state at the chunk boundary. The algorithm (Section 6 and Listing 1) computes the output as the sum of these two kinds of contribution, in four steps:

1. Diagonal blocks (inside each chunk). For each chunk, compute $(L_c \circ C_c B_c^{\top}) X_c$ in attention mode, as if the state entering the chunk were zero. The chunks are independent, so this is one batched matrix multiplication over all of them. Result: $Y_{\text{diag}}$.
2. Right factors (inputs to chunk state). For each chunk, compute the state at its last token, again starting from zero: the sum of $B_i x_i^{\top}$ over the chunk's tokens, each decayed from token $i$ to the chunk's end. That is an $(N \times Q)$ by $(Q \times P)$ matrix multiplication per chunk.
3. Center factors (chunk state to chunk state). The states from step 2 ignore everything before their chunk. Run the scalar recurrence across chunks: the true state entering chunk $c+1$ is the true state entering chunk $c$, decayed by the product of chunk $c$'s decays, plus chunk $c$'s own state from step 2. This is a scan of length $T/Q$, not $T$.
4. Left factors (chunk state to outputs). For each token $j$ in a chunk, read the true incoming state through $C_j$, decayed from the chunk's start to $j$. Another batched matrix multiplication. Result: $Y_{\text{off}}$.

The output is $Y = Y_{\text{diag}} + Y_{\text{off}}$. Steps 2 to 4 are the three factors of (3) applied chunk by chunk: $v$ (right), $W$ (center), $u$ (left). They are never multiplied out into the off-diagonal blocks of $M$.

The cost follows from the shapes (Section 6.3). Counted as multiply-adds per head, step 1 is $TQ(N + P)$, steps 2 and 4 are $TNP$ each, and step 3 is $TNP/Q$. With $N = P = Q$ every matrix multiplication is $N \times N$ by $N \times N$, the total is $O(TN^2)$ (the same order as the recurrence), memory is $O(TN)$, and the only work that is not matrix multiplication is the step-3 scan over $T/Q$ chunks, which the paper reports as negligible in its GPU implementation (Theorem 6.1). The chunk length is a free parameter: the paper's Listing 1 defaults to 64, the official Mamba-2 module to 256.

At $T = 1024$ and $N = P = 64$, per head: chunks of $Q = 64$ give 8.4M multiply-adds in the diagonal blocks, 4.2M each in steps 2 and 4, and a 16-step scan with 65.5k multiply-adds, 16.8M in total. Pure attention mode ($Q = T$) needs 134M, all in the diagonal block. Pure recurrence ($Q = 1$) has the fewest multiply-adds, but it spends 4.2M of them in a scan of 1024 sequential steps, and its “matrix” multiplications are between 1 × 1 blocks.

The algorithm in PyTorch

The paper's Listing 1 implements the four steps in about 30 lines; the copy here differs from it only in how long comment lines wrap. In this code A is the log of the decay, one value per token and head, so a cumulative sum of A is the log of a cumulative product of decays. segsum turns it into all the partial sums $\sum_{k=i+1}^{j} A_k$ and sets the entries above the diagonal to $-\infty$, so exp(segsum(A)) is exactly the mask $L$ of (4), with $L_{ji} = a_j \cdots a_{i+1}$ below the diagonal, 1 on it and 0 above it.

def segsum(x):
    """Naive segment sum calculation. exp(segsum(A)) produces a 1-SS matrix,
       which is equivalent to a scalar SSM."""
    T = x.size(-1)
    x_cumsum = torch.cumsum(x, dim=-1)
    x_segsum = x_cumsum[..., :, None] - x_cumsum[..., None, :]
    mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
    x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
    return x_segsum

def ssd(X, A, B, C, block_len=64, initial_states=None):
    """
    Arguments:
        X: (batch, length, n_heads, d_head)
        A: (batch, length, n_heads)
        B: (batch, length, n_heads, d_state)
        C: (batch, length, n_heads, d_state)
    Return:
        Y: (batch, length, n_heads, d_head)
    """
    assert X.dtype == A.dtype == B.dtype == C.dtype
    assert X.shape[1] % block_len == 0

    # Rearrange into blocks/chunks
    X, A, B, C = [rearrange(x, "b (c l) ... -> b c l ...", l=block_len)
                  for x in (X, A, B, C)]

    A = rearrange(A, "b c l h -> b h c l")
    A_cumsum = torch.cumsum(A, dim=-1)

    # 1. Compute the output for each intra-chunk (diagonal blocks)
    L = torch.exp(segsum(A))
    Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)

    # 2. Compute the state for each intra-chunk
    # (right term of low-rank factorization of off-diagonal blocks; B terms)
    decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
    states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)

    # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states
    # at chunk boundaries (middle term of factorization of off-diag blocks;
    # A terms)
    if initial_states is None:
        initial_states = torch.zeros_like(states[:, :1])
    states = torch.cat([initial_states, states], dim=1)
    decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
    new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
    states, final_state = new_states[:, :-1], new_states[:, -1]

    # 4. Compute state -> output conversion per chunk
    # (left term of low-rank factorization of off-diagonal blocks; C terms)
    state_decay_out = torch.exp(A_cumsum)
    Y_off = torch.einsum('bclhn,bchpn,bhcl->bclhp', C, states, state_decay_out)

    # Add output of intra-chunk and inter-chunk terms
    # (diagonal and off-diagonal blocks)
    Y = rearrange(Y_diag+Y_off, "b c l h p -> b (c l) h p")
    return Y, final_state

Line by line: Y_diag is step 1, contracting $C$ with $B$ over the state index n, masking by L, and contracting with X over the in-chunk position s. decay_states is the decay from each token to its chunk's end, and states is step 2. Step 3 prepends the initial state (zeros unless one is passed in), builds a 1-semiseparable mask over chunks from each chunk's total log-decay, and applies it with one small einsum; after this, states[:, c] is the true state entering chunk $c$, and final_state is the state after the last token. Y_off is step 4. The paper calls this version “relatively efficient even in native PyTorch without requiring special low-level kernels”; the speed numbers later on are for a dedicated implementation, not for this listing.

The model calls the same function with concrete shapes. The listing below shows one Mamba-2 layer at the 2.7B model's width, with the official module's defaults:

# One Mamba-2 layer at d_model = 2560 (code defaults: expand 2, headdim 64,
# d_state 128, chunk_size 256), one sequence of T = 1024 tokens:
#   heads H = 2 * 2560 / 64 = 80,  P = 64,  N = 128,  Q = 256,  chunks c = 4
dt = softplus(dt_raw + dt_bias)        # (1, 1024, 80)    one step size per head
A  = -exp(A_log)                        # (80,)            learned, negative
ssd(X * dt[..., None], A * dt, B, C, block_len=256)   # log a_t = dt_t * A
#   X  (1, 4, 256, 80, 64)   after the chunk rearrange
#   L  (1, 80, 4, 256, 256)  one 256 x 256 mask per head per chunk
#   states      (1, 4, 80, 64, 128)  -> after the cat: (1, 5, 80, 64, 128)
#   decay_chunk (1, 80, 5, 5)
#   final_state (1, 80, 64, 128)     the state you would keep for generation

The step size $\Delta_t$ enters twice (this is the official test in ssd_minimal.py): it scales the input $x_t$, and it scales the log-decay, so $\log a_t = \Delta_t A$. The model has one $B$ and one $C$ for all 80 heads (ngroups=1), so they reach the listing with a head axis of size 1, which the einsums broadcast; the official test calls it the same way. The SSD function has no weights of its own. In training, the gradient flows through it into the input projection that produced $x$, $B$, $C$ and $\Delta$, into A_log and the step-size bias, and from there into everything upstream.

The Mamba-2 block

The SSD layer sits inside a block that also changed. In Mamba-1, an input projection makes the SSM input $x$ and a gate $z$; $x$ passes through a short convolution; and then a second projection, x_proj, computes $\Delta$, $B$ and $C$ from the convolved $x$. Mamba-2 views the layer as a function of $(A, X, B, C)$, the way attention is a function of $(Q, K, V)$, and produces all of its inputs with one projection at the start of the block (Section 7.1, Figure 6). In the official module that projection outputs $z$, $x$, $B$, $C$ and $\Delta$ together, and the convolution runs over $x$, $B$ and $C$. As in Mamba-1, the module also adds a skip term $D\,x_t$ to the SSD output, with one learned $D$ per head.

Mamba-2 does this for tensor parallelism, which splits each large weight matrix of a layer across several GPUs. A Transformer's attention and MLP blocks need one all-reduce each under this scheme: the first projection is split by columns, so each GPU computes its own heads, and the last is split by rows, with one sum at the end. In Mamba-1, x_proj reads all channels of $x$, and after the first projection those channels are spread across GPUs, so computing $B$ and $C$ needs communication in the middle of the block. In Mamba-2 every SSM input comes from the block's input, which each GPU already has, so the block needs one all-reduce, like attention (Figure 7 of the paper). The parallel projection also saves parameters: in the paper's block ablation (Table 4) it takes a 129.3M model to 126.5M and perplexity from 11.76 to 11.66.

The second change is an extra normalization after the output gate and before the output projection. The paper reports that preliminary experiments at larger scale were prone to instability and that this norm helps; in Table 4 it also improves perplexity (11.76 to 11.54 alone, 11.49 with both changes). TransNormerLLM and RetNet add a similar norm, but before the gate.⁵ The paper allows LayerNorm, GroupNorm or RMSNorm and uses GroupNorm under tensor parallelism, so that each GPU can normalize its own group of channels with no extra communication; the official module uses a gated RMSNorm applied after the gate (norm_before_gate=False). All three normalize each token's channels separately, as in layer normalization.

Those state sizes differ by a factor of 8: 655,360 numbers per layer for Mamba-2 against 81,920 for Mamba-1, at the same width. The paper's introduction describes this as states “8× the size of Mamba or even higher, with minimal slowdown”. The size of each head's state follows the attention analogy: $B$ and $C$ are shared by all heads (ngroups=1) while each head has its own $x$ and its own decay. The paper calls this pattern multi-value attention (MVA), or multi-input SSM, by analogy with multi-query attention, which shares keys and values instead. Mamba already used it, and it ablates best: at 125M parameters, perplexity 11.66 against 12.62 for sharing $B$ and $X$ (multi-query, one $C$ per head) and 12.59 for sharing $C$ and $X$ (multi-key, one $B$ per head), with the same total state size (Table 5).

State size and associative recall

A recurrent model keeps only what fits in its state; it cannot reread the input. Tasks that require storing many unrelated facts test that limit. Multi-query associative recall (MQAR) is the synthetic version: the model sees a list of key-value pairs, then a series of keys, and must output each key's value.⁷ The paper uses a harder setup than the task's original: sequence length 256, 512 or 1024 with a quarter as many key-value pairs (so 256 pairs at length 1024), a vocabulary of 8192, two-layer models, and model dimension 32 to 256. The baselines are softmax attention and Based, a linear-attention hybrid.

Two results (Figure 8 of the paper). First, larger state helps: the paper reports that raising $N$ from 16 to 64 to 256 consistently improves accuracy. The one exception in its plot is length 256 at model dimension 32, where $N = 64$ reads 0. Second, at equal state size Mamba-2 is far ahead of Mamba-1. At $N = 16$ and length 1024, Mamba-1 is at 0 for every model size while Mamba-2 climbs to 0.92 at dimension 256. The paper says it is not sure which part of the architecture causes this. Attention scores below Mamba-2's larger states on this task: at length 1024 it reaches about 0.27 at dimension 64 and 0.52 at 128, while Mamba-2 with $N = 64$ or 256 is near 0.9 or above from dimension 64. The paper's caption summarizes Mamba-2 as performing “significantly better than Mamba-1 and even vanilla attention” on this task.

MQAR is a controlled test with tiny models. It shows that state size limits recall on this task and that Mamba-2 can afford a much larger state, but it does not settle how pure SSMs compare to attention on recall in large language models, which is one motivation for the hybrids tested below.

Speed: SSD against Mamba's scan and FlashAttention-2

The paper benchmarks the SSD layer on an A100 80GB PCIe (Figure 10). Against Mamba's fused scan at state size $N = 64$, SSD is 2 to 8 times faster across sequence lengths. Both do work proportional to $T$; SSD does it as matrix multiplication. Against FlashAttention-2, SSD is faster from sequence length 2K onward, and about 6 times faster at 16K, because its work grows like $T$ and attention's like $T^2$ while both run on tensor cores.

The second panel fixes the sequence length at 4K and varies $N$. The scan's time grows linearly with $N$, from about 0.45 ms at 16 to about 7.6 ms at 256. SSD goes from about 0.22 ms to about 0.68 ms over the same range, because a larger $N$ makes its matrix multiplications larger, which tensor cores absorb with little extra time. At $N = 64$, SSD takes about 0.28 ms, less than the scan needs at $N = 16$.

A faster layer does not make a faster model at every length. The paper notes that at short sequences (about 2K) a Mamba-2 model may train more slowly than a Transformer of the same size, because the Transformer spends half its layers on MLPs, which are pure matrix multiplication, while Mamba-2 uses SSD layers throughout. Mixing SSD and MLP layers is its suggested remedy.

Language modeling and hybrids

Scaling. The paper trains models from about 125M to 1.3B parameters on the Pile with the GPT-2 tokenizer, at the token counts the Chinchilla rules prescribe for each size, and compares them with Mamba and with Transformer++. Transformer++ is the strong Transformer recipe from the Mamba paper, based on PaLM and LLaMA: rotary position embeddings, a SwiGLU MLP, RMSNorm instead of LayerNorm, no bias terms, and a peak learning rate 5 times the GPT-3 value. Mamba-2 matches or exceeds both (Figure 9), and against the Transformer baseline it is better on perplexity, on theoretical FLOPs and on wall-clock time at once.

Downstream. At 2.7B parameters, trained for 300B tokens on the Pile with the GPT-NeoX tokenizer (the same data, tokenizer and length as Pythia), Mamba-2 reaches Pile perplexity 6.09, against 6.22 for Mamba-2.8B and 6.73 for Pythia-2.8B, and an average of 60.2 over seven zero-shot tasks (LAMBADA, HellaSwag, PIQA, ARC-easy, ARC-challenge, WinoGrande, OpenBookQA) against 59.9 and 55.7. It beats Pythia-2.8B on every task. Against Mamba-2.8B it is not a sweep: Mamba is ahead on ARC-easy (69.7 against 69.6) and OpenBookQA (39.6 against 38.8). The appendix's full table adds Pythia-6.9B, which Mamba-2-2.7B beats on every column, with average accuracy 60.2 against 58.3. The paper's summary for each size is that Mamba-2 outperforms Mamba and generally matches Pythia at twice the size.

Hybrids. In a 350M model with 48 layers trained to 7B tokens, replacing some SSD layers with attention layers lowers perplexity from 8.60 (pure Mamba-2) to 8.38 with a single attention layer, and to a minimum of 8.26 with six. Past that it rises again, to 8.50 at 24. Transformer++ gets 8.68. The paper reads this as around 10% of layers being attention (6 of 48 is 12.5%) and hypothesizes that attention acts as a retrieval mechanism, letting the model refer back to an exact earlier token instead of relying on what the SSM state kept.

At 2.7B and 300B tokens (all variants 64 layers deep, same data order), pure Mamba-2 gets Pile perplexity 6.09, Transformer++ (32 attention and 32 MLP layers) 6.13, and Mamba-2-Attention, with 6 attention layers at positions 9, 18, 27, 36, 45 and 56 among 58 SSD layers, gets 5.95 and the best average accuracy of the five variants, 61.0.

What the paper does not show

The cost of the scalar decay. No experiment compares a scalar $A_t$ with a diagonal one under otherwise equal conditions. Every Mamba-2 versus Mamba comparison changes the state size, the head dimension, the block and the kernel together, so the gains cannot be credited to any one of them, and a task that needs different memory lengths in different state slots at the same token could lose from the restriction. The paper hypothesizes that its structured-matrix algorithms could be extended to the general diagonal case.

Recall at long range. The longest MQAR sequence is 1024 tokens. The paper makes no claim either way about recall at much longer lengths, where a fixed-size state has to hold more pairs.

Model-level speed. The 2 to 8 times and the 6 times are per-layer measurements. As noted above, the paper cautions that a whole Mamba-2 model might train less efficiently than a Transformer at short sequence lengths.

Systems work. Tensor parallelism, sequence parallelism and variable-length batching all had to be adapted to SSMs (Section 8). For sequence parallelism, each GPU takes a stretch of the sequence, runs it from an incoming state, and passes its final state to the next GPU, which is the chunked algorithm at the scale of devices and needs communication linear in the number of GPUs.