Skip to main content

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

Derive Mamba-2 from scalar state updates to a causal mixing matrix, chunked SSD, complete block structure and runnable CPU references.

Keep the recurrent state, use more matrix multiplication

Mamba ends with a hardware question. Its input-dependent recurrence compresses earlier tokens into a fixed-size state, and affine composition lets training compute prefixes in parallel. But its selective scan mainly uses elementwise arithmetic. GPUs can execute matrix multiplication much more efficiently through specialized hardware. How can a recurrent sequence model use that capability?

Mamba-2 answers by studying the same sequence operator in recurrent and matrix forms. A restricted state transition produces a structured causal matrix. Breaking that matrix into chunks exposes local matrix products, while compact states carry information between chunks. This structured state space duality, or SSD, supplies both an algorithm and a useful connection to attention. Paper v1, §§5–6

Only a short prerequisite is needed: an SSM retains part of a previous state, writes the current input, then reads the updated state. Input-dependent coefficients provide selection. Mamba-2 keeps that idea while changing the transition structure and the surrounding block. Its paper title refers to a precise connection between particular operators; it does not identify arbitrary softmax attention with Mamba-2.

A head has a width axis and a state axis

Consider one head with input vector x of width P and state matrix S of shape [P, N]. N is the number of state coordinates for each input channel. The current B and C vectors each have N coordinates. The update writes an outer product into the state, and the readout contracts its N axis. The explorer connects that memory to the matrix of input-to-output paths.

For this explorer and the chunk figure below, use one head with P=2, N=3 and six positions. The initial state is a provided nonzero 2-by-3 array, so the outputs include incoming memory as well as the six new inputs. The coefficients are authored, not learned from text. This small example is enough to expose how every earlier input reaches a later output.

One recurrence, one causal matrix

Six positions, one head, width P = 2 and state size N = 3. Select a causal cell to expand an earlier input’s contribution.

Causal mixing: input j → output i

i / j012345012345

Teal + / violet −; intensity = |M| / 0.348. Blank cells are future inputs. Every causal cell is a keyboard button.

Read the selected input through memory

x[0,0] = 1.000
inject Δ[0] once

B at source 0 · three state coordinates

n00.700
n1-0.200
n20.300
empty interval · product = 1
Retention product 1.000

C at output 0 · read the three coordinates

n00.500
n10.300
n2-0.200
Δ 0.200 × C·B 0.230 × retention = M[0,0] 0.046
Selected contribution = 0.046

Actual memory h[0] · P = 2 × N = 3

0
0.314
-0.127
0.321
1
-0.230
0.364
0.063

Teal + / violet −. Intensity = |value| / 0.364 (largest magnitude in this array).

All input contributions 0.046

+ Initial-state contribution 0.009

= output 0.055

The nonzero incoming memory contributes separately; it is not another column of M.

Inspect exact values and equations
M[0,0] = Δ[0] × (C[0] · B[0]) × ∏ retention[k], k = 1…0
0.200 × 0.230 × 1.000 = 0.046
Contribution at p = 0: M × x[0,0] = 0.046

When i = j, the retention product is empty and equals 1. Continuous a = −0.7; discrete retention[k] = exp(a × Δ[k]). Input injection includes Δ[j] exactly once.

Coefficient factors

coordinateB[j,n]C[i,n]product
n = 00.7000.5000.350
n = 1-0.2000.300-0.060
n = 20.300-0.200-0.060

Incoming state h[−1] · [P = 2, N = 3]

coordinaten = 0n = 1n = 2
p = 00.200-0.1000.300
p = 1-0.2000.4000.100

Updated state h[0] · [P = 2, N = 3]

coordinaten = 0n = 1n = 2
p = 00.314-0.1270.321
p = 1-0.2300.3640.063
initial-state output = (∏ retention[k], k = 0…0) × Σ C[0,n] h[−1,0,n]
Σ M[i,j] x[j,p] 0.046 + initial-state output 0.009
= dense 0.055 = recurrent 0.055

Difference = 1.39e-17. Selected input x[0,0] = 1.000; coefficients may be signed.

Step 1 / 6

Output 0, source 0, width coordinate 0: compare matrix mixing plus initial memory with recurrence.

Calculated authored coefficients, not attention probabilities or pretrained outputs. Dense materialization explains six positions; production SSD avoids a full sequence matrix. The incoming state is nonzero and contributes separately.

Select a causal cell to inspect its B/C factors, intervening retentions, signed coefficient and contribution. Select an output and width coordinate to compare the dense expression with the recurrent result. The matrix operates along the six-position sequence axis; the same matrix acts on each of the head’s P channels. The figure also shows the state array so the matrix view stays connected to the running computation.

The state update and readout are:

S_t=a_t S_{t-1}+\Delta_t x_t B_t^{\top},\qquad y_t=S_t C_t

The important restriction is that a is a scalar per head and position. Every P-by-N state coordinate within a head shares that retention. In the pinned block, continuous A is a learned negative scalar per head, delta is positive and token-dependent, and their product gives the log decay:

A_h=-\exp(A_{\log,h}),\qquad \ell_{t,h}=\Delta_{t,h} A_h,\qquad a_{t,h}=\exp(\ell_{t,h})

This differs from Mamba’s learned [E, N] diagonal dynamics, which can assign a different continuous rate to each channel/state coordinate. Mamba-2 gives up that variation within a head to obtain a simpler factorization. Multiple heads still provide different rates and different sequence transformations. With H heads, expanded width E equals H × P and persistent SSM state has shape [b, H, P, N]. Paper §5.1 and §7.2

Unroll the recurrence into a causal matrix

Input j first writes delta_j × x_j × B_j into the state. To contribute to output i, that write must survive the transitions at j+1 through i, then be read by C_i. Consequently the mixing coefficient is:

M_{ij}=\begin{cases}\Delta_j\,(C_i^{\top}B_j)\displaystyle\prod_{k=j+1}^{i}a_k,&j\leq i\\0,&j>i\end{cases}

The diagonal is included: at i=j the product is empty and equals one, so a token can affect its own output immediately. The current retention a_j applies to the state before that token’s write; multiplying it into that token’s own contribution would be an indexing error. Incoming memory has a separate path that does pass through every retention from zero to i:

y_i=\sum_{j=0}^{i}M_{ij}x_j+\left(\prod_{k=0}^{i}a_k\right)S_{-1}C_i

Written compactly, the zero-state transformation is M = L ⊙ (C Bᵀ), with delta absorbed into B or the value input exactly once. L contains the causal products of scalar retentions. This is a structured masked kernel-attention form: C plays the query role, B the key role and x the value role. There is no softmax, no requirement that coefficients be positive, and no row-sum normalization. The matrix is a linear map once its coefficients are fixed; the full neural block remains input-dependent. Paper Equation 16 and §5.3

The useful structure is semiseparability: off-diagonal interactions can factor through the N-dimensional state rather than requiring an independent coefficient for every token pair. The full causal matrix need not have low rank. Instead, strictly off-diagonal blocks have rank at most N for this single-head operator. That local rank structure is what makes a block decomposition useful. Paper §3.1

Materializing a T-by-T matrix would still take quadratic space and work. The dense view explains the operator and checks a short example. The next step is to exploit its structure without building the full matrix.

Compute within chunks, then pass what crosses their boundaries

Partition the sequence into chunks of length Q. The diagonal matrix blocks describe interactions entirely within a chunk. Those small [Q, Q] blocks can use matrix multiplication and can be evaluated independently. The missing off-diagonal interactions describe information arriving from earlier chunks, which the recurrent state can summarize.

Original paper · Figure 5. Tri Dao and Albert Gu connect diagonal input-to-output blocks with the state factors between chunks. Reproduced unchanged under CC BY 4.0. The explorer below applies these steps to our separate six-position example.

Local matrices, carried state

Reuse the six-position example. Work through four stages while inspecting a chunk’s actual P × N boundary arrays.

Each chunk first computes its local causal outputs from a zero state. These local matrix products are independent across chunks.

Diagonal blocks work locally

L
L
L
↗
↗
L
↗
↗
L
L
↗
↗
↗
↗
L
↗
↗
↗
↗
L
L

Each 2 × 2 diagonal block computes local outputs (L). Amber regions are recovered by carried boundary state rather than materialized as dense blocks.

Boundary memory connects chunks

Chunk 0 · t0…1Zero-state local outputs
t00.046
t10.161
carry P × N boundary from t1
Chunk 1 · t2…3Zero-state local outputs
t20.158
t30.087
carry P × N boundary from t3
Chunk 2 · t4…5Zero-state local outputs
t40.057
t50.144

Decay incoming memory, then add this chunk’s zero-state summary

Incoming boundary before t0 · p0

n00.200
n1-0.100
n20.300

× Chunk retention 0.657Product of retention factors at t0…1

Multiply incoming memory by retention

→ Decayed incoming = retention × incoming

n00.131
n1-0.066
n20.197

+ Local final state from zero incoming memory

n00.086
n1-0.190
n20.125

= True outgoing boundary

n00.217
n1-0.256
n20.322
t0, p00.046+0.009=0.055
t1, p00.161+0.112=0.273

Teal local output + amber contribution from incoming memory = final output. Boundary tiles are the selected width coordinate of actual P × N states.

Inspect exact values and equations

Incoming boundary state before t = 0 · [P, N]

coordinaten = 0n = 1n = 2
p = 00.200-0.1000.300
p = 1-0.2000.4000.100

Local outputs with zero incoming state · [chunk length, P]

coordinatep = 0p = 1
t = 00.046-0.018
t = 10.161-0.180

Full sequence final-state agreement · selected p

coordinatechunkedrecurrentdifference
n = 00.5260.526-1.11e-16
n = 10.1400.1400
n = 2-0.063-0.0630
Step 1 / 4

Local output: positions 0 through 1, selected width coordinate 0.

Calculated authored coefficients with nonzero initial memory. This teaching algorithm passes boundaries through a linear chronological recurrence; the authors’ minimal listing instead uses a quadratic boundary-matrix convenience. This browser illustration measures no kernel speed.

The chunked algorithm follows four connected operations:

  1. Compute each chunk’s local outputs assuming zero incoming state.
  2. Compute each chunk’s local final state under that same zero-state assumption.
  3. Pass the true incoming state through the chunk and add its local final state to obtain the true outgoing boundary state.
  4. Read the incoming state at every position inside the chunk, with the appropriate retention, and add those carried outputs to the local outputs.

For a chunk beginning at s and ending at e, both endpoints included, these last two operations are:

S_e=\left(\prod_{k=s}^{e}a_k\right)S_{s-1}+S_e^{local}
y_i=y_i^{local}+\left(\prod_{k=s}^{i}a_k\right)S_{s-1}C_i,\qquad s\leq i\leq e

Choose chunks of two or three positions, then inspect the four stages. The same six inputs and nonzero incoming memory from the duality figure are reused. A chunk’s zero-state local output is an intermediate result; adding its carried contribution restores the true recurrent output. The final state must agree too, because the next chunk or decode call will consume it. Checking only the visible outputs would miss a broken continuation state.

The teaching implementation passes boundary states in a chronological recurrence, with work linear in the number of chunks. Local block products dominate useful parallel work. At fixed head width, state size and chunk size, total work grows linearly with T; its constants still depend on those dimensions. The paper’s formal cost statement chooses P=N, whereas this toy example deliberately uses P=2 and N=3 to keep the axes distinct. Paper Theorem 6.1 and Figure 5

There is a practical distinction in the authors’ compact Listing 1: it constructs a dense matrix over chunk boundaries for convenience. That particular boundary calculation is quadratic in the number of chunks. It should not be used as the justification for asymptotically linear boundary passing. The downloadable teaching reference uses the explicit linear recurrence, while production implementations organize the stages into specialized kernels. Pinned minimal implementation

Numerical stability matters even in the small matrices. Products are formed as exponentials of segment sums of log decays. Subtracting two large cumulative sums can erase a small intervening decay through cancellation. The reference builds the needed local segment sums directly. These details preserve the intended operator; they do not turn the Python example into a latency benchmark.

The matrix and recurrence now agree at both token and chunk boundaries. To obtain a complete model block, we still need the projections, local convolution, direct skip, gate and normalization around SSD.

The Mamba-2 block changes more than the scan

Mamba derives B, C and delta after projecting and convolving x. The pinned Mamba2Simple block instead produces z, x, raw B, raw C and raw delta in a single input projection. It applies the causal depthwise convolution and SiLU to the combined x/B/C branches. Delta bypasses that convolution and becomes positive through softplus with a learned bias. Paper §7.1, pinned Mamba2Simple

Compare the complete mixers

Inspect matching operations in two recorded CPU blocks, with the changed projection, transition and normalization paths made explicit.

Mamba

Outside input · residual bypass ↘

RMS prenormmodel width 4
Project x / z8 content + 8 gate = 16
Conv(x) → SiLU8 channels · causal width 3
Produce Δ / B / CΔ per channel · learned A[8,3]
Selective scanstate [8,3]
D × x
+ direct skipD per expanded channel
z → SiLU
× gatesigned SiLU(z)
Output projectionexpanded 8 → model 4
+ outside inputresidual output [1,6,4]

Mamba-2

Outside input · residual bypass ↘

RMS prenormmodel width 4
Project z / x / B / C / Δ8 + 8 + 3 + 3 + 4 = 26
Conv(x,B,C)SiLU · 14 channels
Δ
4 heads × width 2Δ bypasses convolution
SSD recurrencescalar A/head · state [4,2,3]
D × x
+ direct skipD per head
z → SiLU
× gatethen gated RMSNorm
Output projectionexpanded 8 → model 4
+ outside inputresidual output [1,6,4]

Mamba-2 head 0 · actual state after t5

P = 2 rows × N = 3 state coordinates

0
-0.028
-0.003
-0.022
1
-0.010
-0.010
-0.009

Teal + / violet −. Intensity = |value| / 0.0283 (largest magnitude in this array).

Recorded sample at the highlighted operation

outside input · full shape [1 × 6 × 4]

flat 0-1.278
flat 1-1.167
flat 2-1.056
flat 3-0.944

outside prenorm · full shape [1 × 6 × 4]

flat 0-1.143
flat 1-1.043
flat 2-0.944
flat 3-0.845

projection · full shape [1 × 6 × 26]

flat 0-0.895
flat 1-0.591
flat 20.467
flat 30.929

Only these recorded flattened entries are drawn. The full shape is not reconstructed from four samples.

Mamba projects x/z (8 + 8). Mamba-2 projects z/x/B/C/delta in parallel (8 + 8 + 3 + 3 + 4 = 26); B/C and delta no longer depend on the convolved content projection.

Inspect exact values and equations
Mamba: x/z → conv(x) → Δ,B,C → scan + D*x → gate → out
Mamba-2: z/x/B/C/Δ → conv(x,B,C) → heads → SSD + D*x → gate → RMSNorm → out

Mixer input projection after outside RMS prenorm

outside input · [1 × 6 × 4]

coordinate+0+1+2+3
flat 0…-1.278-1.167-1.056-0.944
flat 4…-0.833-0.722-0.611-0.500

outside prenorm · [1 × 6 × 4]

coordinate+0+1+2+3
flat 0…-1.143-1.043-0.944-0.845
flat 4…-1.229-1.065-0.901-0.737

projection · [1 × 6 × 26]

coordinate+0+1+2+3
flat 0…-0.895-0.5910.4670.929
flat 4…0.207-0.779-0.7710.220

Recorded CPU float64 trace with untrained, authored sinusoidal weights. Shapes describe full tensors; tables show their first eight flattened values at most.

Mamba-2 head 0, after t = 5 · state [P = 2, N = 3]

coordinaten = 0n = 1n = 2
p = 0-0.028-0.003-0.022
p = 1-0.010-0.010-0.009

Mamba-2 persistent buffers: convolution [14,3] and SSM [4,2,3]. Mamba persistent buffers: convolution [8,3] and SSM [8,3]. Here the first dimension of each convolution buffer counts its actual convolved branches.

Step 1 / 8

Projection: showing the Mamba-2 recorded stage.

Both traces use deterministic, authored sinusoidal weights and identical outside inputs, not pretrained activations. This architecture comparison does not rearrange Mamba weights into Mamba-2 weights. Both streams start from zero buffers.

This comparison uses recorded untrained CPU traces for both tiny mixers. Mamba-2 has model width 4, expanded width 8, four heads of width 2, state size 3 and one B/C group shared across heads. Its projection therefore emits 2E + 2N + H = 26 features, while its local convolution operates on E + 2N = 14 channels. The z and delta branches are outside that convolution. Select a stage and model to inspect the actual tensor and its shape.

After SSD, the selected implementation adds D_h × x per head, flattens the heads and multiplies by SiLU(z). It then applies RMS normalization and the output projection. The order is important: norm_before_gate=False means gate first, normalize second. The pinned gated-normalization reference confirms that it multiplies x by SiLU(z) before calculating the mean square. This internal normalization is also separate from the enclosing block’s prenorm and residual. Pinned gated RMS normalization

Parallel projection placement removes a dependency between large projections and makes the block easier to partition across devices. Sharing B/C across heads reduces repeated coefficient features. These architecture choices accompany the SSD operator; duality alone does not prescribe every projection or normalization. Nor can arbitrary Mamba weights simply be rearranged into Mamba-2 weights: the transition constraints and coefficient sources differ.

Streaming needs both kinds of history here too. For G B/C groups, the convolution state is [b, E+2GN, K], and the SSM state is [b, H, P, N]. The one-group reference stores [b,14,3] convolution history and [b,4,2,3] SSM state. Reusing the Mamba-1 convolution-cache formula without accounting for B/C would undercount this block’s storage.

With the full data path established, the executable reference can test equivalence beyond the small matrix shown in the browser.

Original paper, equations and runnable CPU code

The six correspondence rows connect heads, recurrence, dense mixing, chunks, the whole mixer and streaming buffers to actual synchronized Python source. They preserve the distinction between the paper’s operator and the pinned implementation’s block details. Original Figure 5 shows why diagonal chunk blocks and low-rank off-diagonal factors describe the same transformation.

This self-contained CPU reference adapts official v2.2.4 (commit 95d8aba8), retaining attribution. It compares recurrent, dense and chunked SSD on the same inputs, including nonzero incoming memory and gradients. It does not load a checkpoint or benchmark a fused kernel.

Original-paper images: Tri Dao and Albert Gu, Transformers are SSMs (arXiv v1), under CC BY 4.0. PDF excerpts are cropped; Figure 5 is unchanged. Select any source image to enlarge it.

Download Python reference · Pinned requirements

python -m pip install -r requirements-mamba.txt
python mamba2.py --self-check
python mamba2.py --smoke --backward

1. Give each head a width axis and a state axis

Original paper

§7.2 and Equations 17–20 · head structure

A head applies the same structured sequence transformation to P channels. Its memory is a P-by-N array; sharing the transition across N is the restriction that enables the SSD factorization.

Math and tensor shapes

E=HP,\quad X\in\mathbb{R}^{b\times T\times H\times P}
S\in\mathbb{R}^{b\times H\times P\times N}

Reshape expanded channels into heads without changing their values. A is a learned scalar per head, delta varies per token/head, and the selected one-group B/C vectors are shared across heads. Width P and state size N describe different axes and need not be equal.

Tiny mixer: E=8; H=4; P=2; N=3; one shared B/C group

Pinned official code

Python implementation

1. Give each head a width axis and a state axis
python
def split_heads(x, head_width):
    """Split expanded width D into H heads of P channels: [B,T,D] -> [B,T,H,P]."""
    if (
        not isinstance(head_width, int)
        or head_width <= 0
        or x.shape[-1] % head_width
    ):
        raise ValueError("head_width must divide expanded width")
    return x.reshape(*x.shape[:-1], x.shape[-1] // head_width, head_width)

2. Update each head with scalar retention and an outer-product write

Original paper

§5.1 · scalar-identity state spaces

With one scalar transition per head, every coordinate of the head’s state has the same retention. Delta scaling is part of both discretization and input injection.

Math and tensor shapes

\ell_t=\Delta_t A,\quad a_t=\exp(\ell_t)
S_t=a_t S_{t-1}+\Delta_t x_t B_t^{\top},\quad y_t=S_t C_t

This core excludes D*x, the gate and normalization. The Python API receives continuous A and delta separately and forms the discrete factors internally. In upstream ssd_minimal_discrete, the argument named A is instead the already-scaled log decay, and X is already delta-scaled; applying either scale twice changes the operator.

Per head: x [P]; B/C [N]; S [P,N]; y [P]

Pinned official code

Python implementation

2. Update each head with scalar retention and an outer-product write
python
def recurrent_ssd(x, delta, a, b, c, initial_state=None):
    """Return y [B,T,H,P], states [B,T,H,P,N] from continuous scalar A [H].

    delta is [B,T,H]; B/C are [B,T,H,N]. Scale the input by delta once.
    The pre-sequence state is [B,H,P,N]; retention at index zero acts on it.
    """
    _validate_ssd(x, delta, a, b, c, initial_state)
    state = (
        x.new_zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1])
        if initial_state is None
        else initial_state
    )
    outputs, states = [], []
    for index in range(x.shape[1]):
        retention = torch.exp(delta[:, index] * a)
        write = (
            delta[:, index, :, None, None]
            * x[:, index, :, :, None]
            * b[:, index, :, None, :]
        )
        state = retention[..., None, None] * state + write
        outputs.append(torch.einsum("bhpn,bhn->bhp", state, c[:, index]))
        states.append(state)
    return torch.stack(outputs, 1), torch.stack(states, 1)

3. Expand the same recurrence into causal coefficients

Original paper

Equation 16 · structured masked attention form

Unrolling the recurrence produces a lower-triangular matrix. Its pairwise B/C terms multiply a structured retention mask, with no softmax.

Math and tensor shapes

M_{ij}=\mathbf{1}_{j\leq i}\,\Delta_j(C_i^{\top}B_j)\exp\!\left(\sum_{k=j+1}^{i}\ell_k\right)
y_i=\sum_{j=0}^{i}M_{ij}x_j+\exp\!\left(\sum_{k=0}^{i}\ell_k\right)S_{-1}C_i

The diagonal’s empty retention product is one. Nonzero incoming memory contributes separately from M*X. Stable local segment sums avoid subtracting two large cumulative sums. Materializing this dense matrix is quadratic and serves only as an inspectable numerical reference.

For each head: M [T,T] acts on X [T,P]

Pinned official code

Python implementation

3. Expand the same recurrence into causal coefficients
python
def segment_sum(log_decay):
    """Stable sums over j+1..i without subtracting long cumulative sums.

    Input shape [...,T]; output [...,T,T]. The inclusive diagonal has log 0;
    the upper triangle has -inf, so its exponential is exactly zero.
    """
    length = log_decay.shape[-1]
    expanded = log_decay[..., :, None].expand(*log_decay.shape, length)
    strict_lower = torch.ones(
        length, length, dtype=torch.bool, device=log_decay.device
    ).tril(-1)
    lower = strict_lower | torch.eye(
        length, dtype=torch.bool, device=log_decay.device
    )
    sums = expanded.masked_fill(~strict_lower, 0).cumsum(-2)
    return sums.masked_fill(~lower, -torch.inf)


def dense_ssd(x, delta, a, b, c, initial_state=None):
    """Materialize a quadratic teaching matrix; return y, final, M, initial_y.

    M[i,j] = delta[j] * (C[i] dot B[j]) * exp(sum(delta[k]*A, k=j+1..i)).
    This is signed causal mixing, not softmax-normalized attention.
    """
    _validate_ssd(x, delta, a, b, c, initial_state)
    log_decay = (delta * a).transpose(1, 2)
    decay = segment_sum(log_decay).exp()
    mixing = torch.einsum("bthn,bshn,bhts,bsh->bhts", c, b, decay, delta)
    state = (
        x.new_zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1])
        if initial_state is None
        else initial_state
    )
    initial_outputs = torch.einsum(
        "bhpn,bthn,bht->bthp", state, c, log_decay.cumsum(-1).exp()
    )
    output = torch.einsum("bhts,bshp->bthp", mixing, x) + initial_outputs
    final = torch.einsum(
        "bht,bthn,bthp,bth->bhpn", decay[:, :, -1], b, x, delta
    )
    final = final + log_decay.sum(-1).exp()[..., None, None] * state
    return output, final, mixing, initial_outputs

4. Compute local outputs, then carry the missing state

Original paper

Figure 5 and §6 · block decomposition

Original Figure 5 by Tri Dao and Albert Gu, arXiv v1, reproduced unchanged under CC BY 4.0. Diagonal blocks handle within-chunk interactions; state factors carry earlier chunks into later outputs.

Math and tensor shapes

S_e=\left(\prod_{k=s}^{e}a_k\right)S_{s-1}+S_e^{local}
y_i=y_i^{local}+\left(\prod_{k=s}^{i}a_k\right)S_{s-1}C_i

Compute zero-state local outputs and local final states, pass true boundary states chronologically, then add their output contributions. This authored reference uses a linear recurrence over chunks. Upstream Listing 1 uses a dense inter-chunk matrix for brevity; that convenience is quadratic in the number of chunks and is not the asymptotically linear boundary algorithm.

Chunk size Q: local matrix [Q,Q]; boundary state [b,H,P,N]

Pinned official code

Python implementation

4. Compute local outputs, then carry the missing state
python
def chunked_ssd(x, delta, a, b, c, chunk_size, initial_state=None):
    """Dense local blocks plus linear boundary passing; return y, final, chunks."""
    _validate_ssd(x, delta, a, b, c, initial_state)
    if (
        isinstance(chunk_size, bool)
        or not isinstance(chunk_size, int)
        or chunk_size <= 0
        or x.shape[1] % chunk_size
    ):
        raise ValueError(
            "chunk_size must be a positive integer dividing sequence length"
        )
    incoming = (
        x.new_zeros(x.shape[0], x.shape[2], x.shape[3], b.shape[-1])
        if initial_state is None
        else initial_state
    )
    outputs, chunks = [], []
    for start in range(0, x.shape[1], chunk_size):
        end = start + chunk_size
        local, local_final, _, _ = dense_ssd(
            x[:, start:end],
            delta[:, start:end],
            a,
            b[:, start:end],
            c[:, start:end],
        )
        log_decay = delta[:, start:end] * a
        carried = torch.einsum(
            "bhpn,bthn,bth->bthp",
            incoming,
            c[:, start:end],
            log_decay.cumsum(1).exp(),
        )
        final = log_decay.sum(1).exp()[..., None, None] * incoming + local_final
        output = local + carried
        chunks.append(
            {
                "start": start,
                "end": end,
                "localOutputs": local,
                "localFinal": local_final,
                "initialState": incoming,
                "finalState": final,
                "carriedOutputs": carried,
                "outputs": output,
            }
        )
        outputs.append(output)
        incoming = final
    return torch.cat(outputs, 1), incoming, chunks

5. Assemble the selected Mamba-2 mixer

Original paper

Figure 6 and §7.1 · parallel projections and extra normalization

Mamba-2 changes the surrounding block as well as the inner operator. The pinned Mamba2Simple source produces z, x, B, C and raw delta together, then convolves the x/B/C branches.

Math and tensor shapes

[z,v,B^{raw},C^{raw},d]=W_{in}u,\quad\Delta=\operatorname{softplus}(d+b_{\Delta})
r=\operatorname{SSD}(x,\Delta,A,B,C)+D\odot x
o=W_{out}\operatorname{RMSNorm}\!\left(\operatorname{flatten}(r)\odot\operatorname{SiLU}(z)\right)

With E=8, H=4, N=3 and one B/C group, the projection has 2E+2N+H=26 features and the depthwise convolution has E+2N=14 channels. D is per head. norm_before_gate=False means multiply by SiLU(z), then RMS-normalize over expanded channels, then project. The outside prenorm and residual remain separate.

Tiny projection width 26; convolution width 14; output [b,T,4]

Pinned official code

Python implementation

5. Assemble the selected Mamba-2 mixer
python
def forward(self, u, return_trace=False):
    """Follow pinned source order: parallel projection, conv, SSD, skip, norm."""
    projected = self.in_proj(u)  # Order [z, x, B, C, dt].
    z, xbc, dt = projected.split(
        (self.channels, self.conv_channels, self.heads), -1
    )
    delta = F.softplus(dt + self.dt_bias)
    padded = F.pad(xbc.transpose(1, 2), (self.conv_width - 1, 0))
    convolved = F.silu(
        F.conv1d(
            padded,
            self.conv.weight,
            self.conv.bias,
            groups=self.conv_channels,
        )
    ).transpose(1, 2)
    x, b, c = self._split_bc(convolved)
    headed = split_heads(x, self.head_width)
    mixed, states = recurrent_ssd(headed, delta, -self.A_log.exp(), b, c)
    skipped = mixed + self.D[None, None, :, None] * headed
    # norm_before_gate=False: multiply by SiLU(z) BEFORE RMS normalization.
    gated = skipped.flatten(-2) * F.silu(z)
    normalized = rms_norm(gated, self.norm_weight)
    output = self.out_proj(normalized)
    if return_trace:
        return output, {
            "projected": projected,
            "z": z,
            "xbc": xbc,
            "delta": delta,
            "A": -self.A_log.exp(),
            "convolved": convolved,
            "headed": headed,
            "B": b,
            "C": c,
            "states": states,
            "mixed": mixed,
            "skipped": skipped,
            "gated": gated,
            "normalized": normalized,
            "output": output,
        }
    return output

6. Continue the same mixer one token at a time

Original paper

Figure 7 (right) and §8.2 · passing recurrent state

A chunk boundary and a decode boundary both require the correct incoming state. The authored step method also retains the local convolution history needed by the complete block.

Math and tensor shapes

\mathrm{conv}: [b,E+2GN,K],\quad\mathrm{ssm}: [b,H,P,N]

Shift the projected x/B/C history, apply the convolution and update the same scalar-transition state. G is the number of B/C groups; this example fixes G=1. Full-sequence and token-step outputs agree within explicit tolerances, including both buffers. The step method is an authored educational addition, not a method copied from Mamba2Simple.

One-group tiny buffers: convolution [b,14,3]; SSM [b,4,2,3]

Pinned official code

Python implementation

6. Continue the same mixer one token at a time
python
def initial_cache(self, batch):
    """Persistent buffers: [B,D+2N,K] and [B,H,P,N], with K stored samples."""
    reference = self.in_proj.weight
    return (
        reference.new_zeros(batch, self.conv_channels, self.conv_width),
        reference.new_zeros(
            batch, self.heads, self.head_width, self.state_size
        ),
    )

def step(self, token, cache):
    """Stream one token with the same convolution, scalar-A SSD and gate norm."""
    conv_state, state = cache
    z, xbc, dt = self.in_proj(token).split(
        (self.channels, self.conv_channels, self.heads), -1
    )
    delta = F.softplus(dt + self.dt_bias)
    conv_state = torch.cat((conv_state[..., 1:], xbc[..., None]), -1)
    convolved = F.silu(
        (conv_state * self.conv.weight[:, 0]).sum(-1) + self.conv.bias
    )
    x, b, c = self._split_bc(convolved)
    x = split_heads(x, self.head_width)
    state = (
        torch.exp(-delta * self.A_log.exp())[..., None, None] * state
        + delta[..., None, None] * x[..., None] * b[..., None, :]
    )
    mixed = (state * c[..., None, :]).sum(-1) + self.D[None, :, None] * x
    normalized = rms_norm(mixed.flatten(-2) * F.silu(z), self.norm_weight)
    return self.out_proj(normalized), (conv_state, state)

The reference accepts continuous A and delta separately. By contrast, upstream ssd_minimal_discrete accepts the already-scaled log decay delta × A under the argument name A and an already-scaled delta × x under X. The scaling occurs exactly once in each term. Confusing these API conventions can preserve tensor shapes while silently changing the answer.

The self-check compares recurrent, dense and multiple chunked calculations on identical parameters and inputs. It checks outputs, final states, nonzero incoming state, causal independence from future inputs, gradients and full-sequence versus token-step block behavior including convolution history. Float64 checks use absolute and relative tolerances 1e−9; float32 checks use 1e−5. Small floating-point differences can follow from a different contraction order, so exact bitwise equality is not the general claim.

These checks establish implementation correspondence for the examples. Useful language modeling and efficient GPU execution require separate evidence from trained-model evaluations and timed kernels.

Published quality and efficiency measure different things

The paper appears in ICML 2024; that publication information is separate from this write-up’s review date. Official PMLR proceedings

The following are published Pile validation perplexities from Mamba-2 Table 1. Lower is better. These selected rows use the GPT-NeoX tokenizer, the Pile and 300 billion training tokens. The model sizes are shown separately to distinguish Mamba (2.8B) from Mamba-2 (2.7B).

ModelParametersPile validation perplexity
Pythia2.8B6.73
Mamba2.8B6.22
Mamba-22.7B6.09

Source: Mamba-2 v1, Table 1 and §9.2.2. The table also includes downstream zero-shot tasks. Its reported Average includes OpenbookQA, while the original Mamba table uses a different task set, so those two averages should not be compared across papers. The perplexity rows above avoid that mismatch; they still compare separately trained architectures rather than isolating the causal effect of one block change.

Efficiency has another boundary. Figure 10 compares the core SSD algorithm with a fused Mamba scan and FlashAttention-2 under its benchmark settings. Its speed ratios concern those timed operators and state expansions, not an unconditional full-model training or generation ratio. The paper itself notes that a complete Mamba-2 model can remain less training-efficient than a Transformer at short sequence lengths because the allocation of SSD, attention and MLP layers differs. Paper Figure 10 and §9.3

Matrix multiplication is the reason SSD can use tensor cores, but measured speed also depends on head width, state size, chunk length, sequence length, precision, hardware and kernel implementation. None of the browser computations or CPU self-checks measures the published GPU speedup.

Finally, the state is still finite. More efficient state expansion can make larger memories practical, and the paper’s associative-recall experiments investigate the resulting quality. It does not make the state a lossless archive. A structured attention interpretation changes how we compute the same compressed operator; it does not give it the unconstrained token-pair weights of softmax attention.

From sequence operators to image backbones

The progression from Mamba to Mamba-2 is now explicit: input-dependent state updates lead to a restricted per-head transition; that transition exposes a structured matrix; chunk decomposition combines local matrix products with boundary states; a revised block surrounds the operator with parallel projections and gated normalization.

Images introduce a different question: which sequence order should carry spatial information? Vision Mamba scans patch sequences forward and backward, while VMamba aligns four spatial traversal orders inside a hierarchical backbone. They adapt the earlier Mamba-style selective recurrence; they are not simply image checkpoints of the Mamba-2 block described here. Read them next to see how order, direction combination and spatial hierarchy change the use of recurrent memory.

Key takeaways

  • SSD uses scalar retention per head, shared across the head’s width-by-state array, to expose a structured matrix factorization.
  • Recurrent and dense forms describe the same operator, including the contribution from a nonzero incoming state.
  • The attention connection concerns structured masked kernel attention. Its coefficients are signed and unnormalized, with no softmax.
  • Chunking computes local interactions with matrix products and carries earlier information through compact boundary states.
  • The selected Mamba-2 block also changes coefficient projection placement, convolution branches, head sharing and gated normalization.
  • Published perplexity and timed kernel efficiency are distinct evidence; the authored examples demonstrate mechanics and numerical agreement.

If you found this paper review helpful, consider sharing it with others.

Mastodon