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
Teal + / violet −; intensity = |M| / 0.348. Blank cells are future inputs. Every causal cell is a keyboard button.
Read the selected input through memory
B at source 0 · three state coordinates
C at output 0 · read the three coordinates
Selected contribution = 0.046
Actual memory h[0] · P = 2 × N = 3
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
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
| coordinate | B[j,n] | C[i,n] | product |
|---|---|---|---|
| n = 0 | 0.700 | 0.500 | 0.350 |
| n = 1 | -0.200 | 0.300 | -0.060 |
| n = 2 | 0.300 | -0.200 | -0.060 |
Incoming state h[−1] · [P = 2, N = 3]
| coordinate | n = 0 | n = 1 | n = 2 |
|---|---|---|---|
| p = 0 | 0.200 | -0.100 | 0.300 |
| p = 1 | -0.200 | 0.400 | 0.100 |
Updated state h[0] · [P = 2, N = 3]
| coordinate | n = 0 | n = 1 | n = 2 |
|---|---|---|---|
| p = 0 | 0.314 | -0.127 | 0.321 |
| p = 1 | -0.230 | 0.364 | 0.063 |
Σ 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.
Output 0, source 0, width coordinate 0: compare matrix mixing plus initial memory with recurrence.
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:
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:
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:
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:
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
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
Decay incoming memory, then add this chunk’s zero-state summary
Incoming boundary before t0 · p0
× Chunk retention 0.657Product of retention factors at t0…1
→ Decayed incoming = retention × incoming
+ Local final state from zero incoming memory
= True outgoing boundary
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]
| coordinate | n = 0 | n = 1 | n = 2 |
|---|---|---|---|
| p = 0 | 0.200 | -0.100 | 0.300 |
| p = 1 | -0.200 | 0.400 | 0.100 |
Local outputs with zero incoming state · [chunk length, P]
| coordinate | p = 0 | p = 1 |
|---|---|---|
| t = 0 | 0.046 | -0.018 |
| t = 1 | 0.161 | -0.180 |
Full sequence final-state agreement · selected p
| coordinate | chunked | recurrent | difference |
|---|---|---|---|
| n = 0 | 0.526 | 0.526 | -1.11e-16 |
| n = 1 | 0.140 | 0.140 | 0 |
| n = 2 | -0.063 | -0.063 | 0 |
Local output: positions 0 through 1, selected width coordinate 0.
The chunked algorithm follows four connected operations:
- Compute each chunk’s local outputs assuming zero incoming state.
- Compute each chunk’s local final state under that same zero-state assumption.
- Pass the true incoming state through the chunk and add its local final state to obtain the true outgoing boundary state.
- 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:
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 ↘
Mamba-2
Outside input · residual bypass ↘
Mamba-2 head 0 · actual state after t5
P = 2 rows × N = 3 state coordinates
Teal + / violet −. Intensity = |value| / 0.0283 (largest magnitude in this array).
Recorded sample at the highlighted operation
outside input · full shape [1 × 6 × 4]
outside prenorm · full shape [1 × 6 × 4]
projection · full shape [1 × 6 × 26]
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-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.591 | 0.467 | 0.929 |
| flat 4… | 0.207 | -0.779 | -0.771 | 0.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]
| coordinate | n = 0 | n = 1 | n = 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.
Projection: showing the Mamba-2 recorded stage.
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
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
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
Python implementation
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
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
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]
Python implementation
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
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
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]
Python implementation
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_outputs4. Compute local outputs, then carry the missing state
Original paper
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
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]
Python implementation
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, chunks5. Assemble the selected Mamba-2 mixer
Original paper
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
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]
Python implementation
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 output6. Continue the same mixer one token at a time
Original paper
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
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]
Python implementation
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).
| Model | Parameters | Pile validation perplexity |
|---|---|---|
| Pythia | 2.8B | 6.73 |
| Mamba | 2.8B | 6.22 |
| Mamba-2 | 2.7B | 6.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.
Related reading
- Mamba: build the selective state update, associative scan and persistent decode buffers before restricting the transition.
- Vision Mamba: see bidirectional patch sequences and the visual readout built around selective mixers.
- VMamba: see four spatial traversals, aligned directional outputs and a hierarchical feature map.
- Linear attention approximations: connect attention feature maps to recurrent summaries.
- Tensor cores: understand the matrix-multiplication hardware that motivates the SSD algorithm.
- Original Mamba-2 paper, arXiv v1, ICML proceedings and official pinned source: the theoretical framework, experiments and implementation.
