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

In plain words

Transformers handle sequences by letting every token look at every earlier token, which gets expensive as the text grows. A newer family, state-space models, instead carries a compact running memory from token to token, which scales well but has been harder to train quickly on modern GPUs. This paper shows that the two families are closely related. Once you write a state-space model as one big structured matrix, you can compute it either token by token (cheap in memory, slow to train) or as a set of matrix products (fast on GPUs, costly for long sequences).

The paper's answer is to split the sequence into chunks. Inside a chunk it uses the matrix-product form. Between chunks it passes a small state along with the recurrent form. The result is a layer, SSD, that is fast on GPUs, and a block design, Mamba-2, that is built around it. The authors report that it trains faster than the earlier Mamba layer and remains competitive on language modelling. These numbers come from the authors' own experiments.

TL;DR

• The paper links state-space models (SSMs) and a masked form of linear attention through "structured matrices", and calls the link state space duality (SSD) §1, §5.
• SSD restricts the SSM so that its state decay is one scalar per head. Under that restriction, the recurrent form and the quadratic form are dual §5.1–§5.3.
• The chunked SSD algorithm uses matrix multiplications inside chunks and a small recurrence across chunks. The authors state linear training cost, linear-in-length inference and matmul-dominated work §6, Theorem 6.1.
• Author-reported: SSD is 2–8× faster than Mamba's fused scan at state size 64, and faster than FlashAttention-2 from sequence length 2K §9.3, fig. 10.
• Use this card if you implement Mamba-2 or a scalar-decay SSM. Read the Pitfalls section first: the paper omits discretisation and several implementation details that the code fills in.

Problem

Softmax attention costs quadratic time in sequence length during training and needs a cache that grows with length during generation §1. Structured SSMs cost linear time in length and keep a constant-size state during generation §1, but the authors say their development had become disconnected from the work on Transformers. Training SSMs as efficiently as Transformers was described as an open problem from both the algorithmic and the systems side §1.

The paper's own limit is that SSD only covers the case of a scalar state decay per head. General diagonal SSMs such as Mamba's selective scan are not covered by the SSD algorithm §5, §7.2. The main results are stated for the scalar case.

Method

1. SSMs as matrices. A selective SSM with recurrence h_t = A_t h_{t-1} + B_t x_t and output y_t = C_t^T h_t is a linear map y = M x with M_ji = C_j^T A_{j:i} B_i §2, §3. Here A_{j:i} is the product of the state matrices from i+1 to j. The matrix M is lower-triangular and N-semiseparable, so it has a sequential semiseparable (SSS) representation §3. The paper omits discretisation from these equations and says the main models use the same parameterisation as Mamba §2, remark.

2. Scalar-identity SSMs. If each A_t is a scalar times the identity, A_t = a_t I, the matrix becomes M = L ∘ (C B^T). Here L is a 1-semiseparable mask built from cumulative products of the scalars a_t §5.1. Written as tensor contractions, this is the same computation as masked kernel attention, with C playing the role of queries, B of keys and X of values §5.1, §5.3, fig. 4.

3. 1-semiseparable masks. A 1-SS matrix 1SS(a) has entries a_{j:i} for j ≥ i. Multiplying by it is the scalar recurrence y_t = a_t y_{t-1} + x_t [§3, eq. for 1SS]. Standard causal linear attention is the special case a_t = 1 §5.2. Structured masked attention with a 1-SS mask is therefore a diagonal SSM with one scalar decay §5.2, corollary. The paper also proves that any autoregressive attention with bounded order must have a semiseparable mask §5.2, theorem; proof in app. C.

4. Structured state space duality. The paper identifies the two forms. The SSM view computes the linear recurrence with state size N. The attention view computes the quadratic masked form with kernel feature dimension N §5.3, fig. 4. The two are exact duals on the scalar-identity case §5.3.

5. Chunked SSD algorithm. Split the length-T sequence into blocks of length Q and write M as a grid of (T/Q) × (T/Q) blocks §6. The four pieces are:

• Diagonal blocks: the outputs within each chunk, computed with the quadratic form and a zero starting state §6, "Diagonal Blocks".
• Right factors: the state at the end of each chunk, assuming zero state at its start. This is one batched matrix product per chunk §6, "Right Factors".
• Centre factors: a scalar recurrence over the T/Q chunk states, which gives the true state at each chunk boundary §6, "Center Factors".
• Left factors: the contribution of the carried-in state to the outputs of each chunk, again one batched matrix product §6, "Left Factors".

The reference listing for the whole layer is [listing 1] in §6 (PyTorch, about 40 lines). The following pseudocode restates it with the chunk recurrence written as a loop. Shapes are given in comments.

# Chunked SSD for scalar-decay SSMs  [§6] [listing 1]
# Input:  X (T, H, P); log_decay (T, H), values <= 0
#         B, C (T, H, N); Q chunk length, T divisible by Q
# Output: Y (T, H, P); final state (H, P, N)

def ssd_chunked(X, log_decay, B, C, Q):
    # Cut the sequence into K = T / Q chunks  [§6] [listing 1]
    X, log_decay, B, C = chunk(X, log_decay, B, C, Q)  # (K, Q, ...)
    # Cumulative log decay, restarted in each chunk  [§5.1] [listing 1]
    decay_cumsum = cumsum(log_decay, over_axis=1)      # (K, Q, H)
    # Diagonal block: outputs from inputs in the same chunk  [§6]
    within_mask = exp(segsum(transpose(log_decay)))    # (K, H, Q, Q)
    Y_within = einsum(C, B, within_mask, X)            # (K, Q, H, P)
    # Right factor: state at each chunk end, zero state at chunk start  [§6]
    decay_to_end = exp(decay_cumsum[:, -1] - decay_cumsum)  # (K, Q, H)
    chunk_states = einsum(B, decay_to_end, X)          # (K, H, P, N)
    # Centre factor: carry states across chunks  [§6]
    # Each chunk decays the carried state by exp of its total log decay
    chunk_decay = exp(decay_cumsum[:, -1])             # (K, H)
    state = zeros(H, P, N)
    for chunk_index in range(K):
        # Save the state that enters this chunk  [§6]
        start_states[chunk_index] = state              # (H, P, N)
        state = chunk_decay[chunk_index] * state + chunk_states[chunk_index]
    # Left factor: effect of earlier chunks on each output  [§6]
    decay_from_start = exp(decay_cumsum)               # (K, Q, H)
    Y_carry = einsum(C, start_states, decay_from_start)  # (K, Q, H, P)
    Y = reshape(Y_within + Y_carry, (T, H, P))
    return Y, state

Two notes on the sketch. The loop over chunks is the paper's centre-factor step, which can be replaced by any scan algorithm for the scalar recurrence §6, "Center Factors"; app. B. The chunk_states line sums over positions inside each chunk; the listing gives the explicit einsum [listing 1].

6. Mamba-2 block. The block moves all data-dependent projections to the start, so that A, X, B and C come from one projection §7.1. The paper also adds one extra normalisation layer before the output projection, following NormFormer §7.1. The B and C projections use one head shared across all X heads §7.1, fig. 6. Mamba's original selective SSM is described as a multi-input SSM with head dimension 1 §7.2. Head patterns are defined as in attention: MHA, multi-query, multi-key and multi-value analogues §7.2. The authors find that the multi-value pattern, the one Mamba uses, performs best §9.4.2, tab. 5.

The code gives more detail than the paper for this block. The following sketch follows the reference implementation mamba2_simple.py. Items marked (code) are not described in the paper text. batch is the batch size; B and C are the SSM input and output matrices.

# Mamba-2 block forward  [§7.1] [code: mamba_ssm/modules/mamba2_simple.py#L131]
# Input:  u (batch, T, D); H heads of size P; state size N; one B/C group
# Output: (batch, T, D)

def mamba2_block(u):
    # One projection gives the gate, the SSM inputs and the step sizes  [§7.1]
    z, x, B, C, dt_raw = split(in_proj(u))             # z, x: (batch, T, H*P)
    # Depthwise causal conv, kernel 4, over x, B and C  (code, not in paper)
    x, B, C = silu(causal_conv1d(concat(x, B, C)))     # (code)
    # Per-head step size from the raw value and a learned bias  (code)
    dt = softplus(dt_raw + dt_bias)                    # (batch, T, H)
    # Per-head scalar decay a_t = exp(dt * A) with A < 0  (code)
    log_decay = dt * (-exp(A_log))                     # (batch, T, H), <= 0
    # Scale the input by dt and run chunked SSD  [§6] [§5.1]
    x_heads = reshape(x, (batch, T, H, P)) * dt[..., None]
    y, _ = ssd_chunked(x_heads, log_decay, B, C, chunk_len)
    # Skip connection with a learned per-head weight  (code, not in paper)
    y = y + D[None, None, :, None] * reshape(x, (batch, T, H, P))
    # Gate first, then normalise: the extra norm layer of §7.1  [§7.1]
    gated = reshape(y, (batch, T, H * P)) * silu(z)    # (code)
    return out_proj(rms_norm(gated))                   # (code)

7. Discretisation. The paper omits discretisation in its equations and defers to Mamba §2, remark. The reference code applies a step size dt = softplus(raw + bias), takes the discrete decay as exp(dt · A) with A = -exp(A_log) < 0, and scales the input by dt mamba2_simple.py. This is the usual Mamba convention. It is not in the paper text.

8. Costs. With state size N, head size P = N, and chunk length Q = N, the authors state training FLOPs O(T N²), memory O(T N), and work dominated by N × N matrix products §6, Theorem 6.1; §6.3. The cost table gives N² FLOPs per inference step for SSD, the same as a plain SSM, with state size N §6.3, table. Attention costs T²N total training FLOPs and a cache of size T §6.3.

Hyperparameters

Values are as written in the paper. "not stated" means the paper text I read does not give a value. The code column is separate and says where it comes from.

Results

Author-reported. Setup: Pile validation perplexity and zero-shot accuracy on seven standard tasks, models trained on the Pile for up to 300B tokens §9.2.2, tab. 1. Compared with Pythia (same tokenizer and data) and Mamba-1 (same recipe) §9.2.2.

Table 1: zero-shot results, main text (from tab. 1; markings from source omitted).

Smaller sizes (130M to 370M) and larger baselines (Pythia-6.9B, RWKV4-7.4B, GPT-J-6B, OPT) are in tab. 9 in app. D.3. Headline claims from the text:

Head-pattern ablation. Multi-input (MVA-style) is best in both settings. At 125M and 2.5B tokens: MIS 11.66, MCS 12.62, MES 12.59, MHS 12.06, MSS 12.00. At 360M and 7B tokens: MIS 8.73, MCS 9.33, MES 9.36, MHS 9.01, MSS 9.04 tab. 5.

Kernel ablation (130M, perplexity). None 11.58; Swish 11.66; Exp 11.62; ReLU 11.73; ReLU+normalisation 11.64; cosFormer 11.97; Random Feature Attention 11.57; Performer 12.21 tab. 6. Based and ReBased variants: Swish 11.67; Swish+Taylor 12.19; LayerNorm 11.50; LayerNorm+Square 11.84 (130M); at 380M with N = 256: Swish 8.58; Swish+Taylor 8.71; LayerNorm 8.61; LayerNorm+Square 8.63 tab. 7. The authors conclude that the kernel variants do not beat a plain pointwise activation §9.4.3.

Scaling laws. Models from about 125M to about 1.3B were trained with the Chinchilla token budget. Mamba-2 matches or beats Mamba and a Transformer++ recipe in perplexity and wall-clock time §9.2.1, fig. 9.

Paper vs code

Searched: the paper links https://github.com/state-spaces/mamba as the code and checkpoint location §1. This repository is by the authors' organisation. Cloned at git rev-parse HEAD = e9594ce1c732d97440f0332fdc43170a2294dbfa. The LICENSE file is Apache License 2.0.

The repository is newer than the paper: its tree includes a mamba3.py module that the paper does not describe. The paper's exact code state is not known from the text. The citations below are to this commit.

1. Activation on the convolution output. The paper applies the feature map ψ (default Swish) to the B and C branches, and says it may also be applied to X "for simplicity and symmetry" §7.3. The code applies SiLU to the whole convolution output, which contains x, B and C mamba2_simple.py. Paper: X optional. Code: X always.
2. Causal convolution. The code has a depthwise causal convolution with kernel 4 over x, B and C before the SSD mamba2_simple.py. The paper's architecture section does not describe a convolution §7.1. Treat it as an implementation detail not in the paper.
3. Skip term D. The code adds a learned per-head skip D · x mamba2_simple.py. None of the equations in the paper include it.
4. Default state and head size. Mamba2Simple defaults to d_state = 64, headdim = 128 mamba2_simple.py. Mamba2 defaults to d_state = 128, headdim = 64 mamba2.py. The ablations in the paper use 64 and 64 §9.4.2, tab. 5. The paper does not state which values were used for the language-model runs.
5. Chunk length. The reference layer defaults to chunk_size = 256 mamba2_simple.py. The paper's listing defaults block_len = 64 [listing 1], and the repository's minimal test uses 64 ssd_minimal.py. The paper gives no value for Q in its experiments.
6. Segment sum. The paper's listing uses a direct cumulative-sum difference with a -inf mask [listing 1]. The repository keeps that form as segsum_unstable ssd_minimal.py and uses a different segsum ssd_minimal.py. Its docstring still says the file is "the same as Listing 1" ssd_minimal.py. The two versions are intended to compute the same values.
7. Discretisation. The paper's listing takes discrete A (log-decays) as input and does not show how dt enters [listing 1]. The repository folds dt into X and into A in its test, ssd_minimal_discrete(x*dt, A*dt, B, C, ...) ssd_minimal.py. The fused path takes dt and dt_bias and applies softplus internally mamba2.py. The paper states none of this.
8. Initialisation and parameter groups. A is initialised uniformly in [1, 16] and stored as A_log, with A = -exp(A_log) mamba2_simple.py. dt is sampled log-uniformly in [0.001, 0.1] with floor 1e-4 mamba2_simple.py. dt_bias, A_log and D are excluded from weight decay mamba2_simple.py. The paper states none of these.
9. Normalisation type. The paper allows LayerNorm, GroupNorm or RMSNorm §7.1. The code uses a gated RMSNorm with eps = 1e-5 and norm_before_gate = False, so the gate is applied before the norm mamba2_simple.py.
10. Groups. The code supports ngroups (G) sharing B and C across head groups mamba2_simple.py. The paper describes grouped-value attention and omits the details §7.2. Ablations do not cover it.

No discrepancies were found in the core SSD algorithm itself: the block-decomposition steps in the paper's listing and in the repository's minimal implementation compute the same four terms [listing 1; code: mamba_ssm/modules/ssd_minimal.py#L53-L74].

Pitfalls

• Scalar decay per head is the restriction that makes SSD work. SSD does not cover Mamba's per-channel diagonal A. Mamba-2 uses one A per head mamba2_simple.py. Results from this paper do not show that a per-channel SSM can use the same algorithm.
• Discretisation is missing from the paper. Implement it from the code (Paper vs code, items 7–8), or from Mamba-1, which the paper says it follows §2, remark. The paper's listing takes already-discretised decays as input and does not show how the step size scales the input [listing 1].
• Chunk length Q is not given. The paper's listing and the reference layer use different defaults (64 vs 256). Memory and speed both depend on Q. Choose it and state it.
• The chunk-level decay matrix grows as (T/Q)². The reference listing builds a (K+1) × (K+1) decay matrix for the chunk recurrence ssd_minimal.py. The paper says the scan cost is negligible §6.3. This holds for moderate T/Q. For very long sequences with small Q, use a scan instead.
• The diagonal-block mask costs memory. exp(segsum(...)) materialises a Q × Q mask per chunk and head [listing 1]. The memory is O(T · Q · H), not the O(T N) in the theorem. The theorem assumes Q = N §6, Theorem 6.1.
• Mask before exp. Upper-triangular cumulative-sum differences are positive and can overflow if they are exponentiated before masking. The listing masks with -inf first [listing 1]. Keep this order.
• Dtype must match. The listing and the repository both assert that X, A, B and C share one dtype [listing 1; code: mamba_ssm/modules/ssd_minimal.py#L44]. Cast the decays to the input dtype before calling the layer.
• Inference cost is stated two ways. Theorem 6.1 gives O(TN) inference FLOPs §6. The cost table gives O(N²) per step §6.3. The paper does not reconcile these. A per-step recurrence with state P × N costs about N² FLOPs when P = N, so the per-step figure is the one to use.
• Some benchmark claims lack detail. The speed comparisons do not give the GPU model or the precision §9.3. The associative-recall results are qualitative: figure only, no numbers in the text fig. 8.
• Ablation differences are small and have no variance estimate. Some kernel-table gaps are about 0.1 perplexity (for example ReLU 11.73 vs Swish 11.66). The paper does not report seeds or repeats tab. 6.
• Two Swish values differ. The 130M Swish row is 11.66 in tab. 6 and 11.67 in tab. 7. The paper does not explain the difference.
• Hybrid results are single comparisons. The attention-layer gains are small (about 0.1 perplexity) and the paper does not state repeats tab. 3.
• The arXiv version may differ from later releases. The repository has moved on (mamba3.py). Pin the commit in code_commit and do not assume later behaviour matches this card.

Lineage

• Builds on:
  • Mamba, the selective SSM with a hardware-aware scan, arXiv:2312.00752 §1, §2, §7.2. Mamba-2 keeps the selective form but restricts A to one scalar per head.
  • The structured SSM family S4 and its diagonal variants §2.
  • Linear attention (Katharopoulos et al., ICML 2020), whose recurrent form the paper extends §1, §4.
  • Multi-query and grouped-query attention, which define the multi-contract and grouped head patterns §7.2.
  • NormFormer, which motivates the extra normalisation layer §7.1.
• Compared against:
  • Mamba-1's selective scan, in speed and in model quality §9.3, §9.2.
  • FlashAttention-2 §9.3.
  • Transformer++, Pythia, RWKV-4 and GPT-J, as zero-shot baselines §9.2.2, tab. 9.
  • Based, ReBased, cosFormer, Random Feature Attention and Performer, as kernel variants §9.4.3.
• Superseded by: not stated in this wiki's content.