Structured State Space Duality

Selectivity solved the modelling problem and created an engineering one. Making the state space parameters depend on the input destroys the global convolution that made S4 trainable in parallel, and the replacement, a hardware-aware associative scan, runs elementwise. Elementwise work on a modern accelerator is a waste of silicon: an H100 does roughly an order of magnitude more arithmetic per second through its tensor cores than through its general-purpose units. Mamba trained fine and left most of the machine idle. The structured state space duality is the reformulation that got the matrix units back.

The recurrence as one matrix

Unroll the selective recurrence \(h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t\), \(y_t = C_t^\top h_t\), and the map from the input sequence to the output sequence is a single lower-triangular matrix \(M\) whose entries are

When \(\bar{A}_k\) is a scalar times the identity, \(a_k I\), the product collapses to \(\prod a_k\) and

That is a causal attention matrix. The term \(C_i^\top B_j\) is a query-key inner product with no softmax, and the product of scalars is a learned, input-dependent decay mask applied on top. Dao and Gu named the resulting class 1-semiseparable matrices and showed that a selective SSM with a scalar-times-identity state matrix and a masked linear attention layer are the same sequence transformation written two ways (Dao and Gu, 2024, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, ICML, arXiv:2405.21060).

The duality is the useful part. The same transformation has a linear form, \(O(T)\) in sequence length, which materialises the state and steps through time, and a quadratic form, \(O(T^2)\), which materialises \(M\) and does one big matrix multiply. Neither is universally better. The linear form wins asymptotically; the quadratic form wins on hardware, because it is a matmul.

The chunked algorithm takes both

Split the sequence into chunks of length \(Q\), typically 64 or 256. Within a chunk, compute the quadratic form on the \(Q \times Q\) block: one masked matmul, straight onto tensor cores. Across chunks, carry a single state vector using the linear form: one recurrence step per chunk rather than per token. The block decomposition of \(M\) makes this exact, not an approximation, because the off-diagonal blocks factor as (chunk output projection) times (accumulated decay) times (chunk input projection), which is a low-rank product the algorithm never has to expand.

Total cost is \(O(TQ)\) for the diagonal blocks plus \(O(TN)\) for the inter-chunk recurrence, still linear in \(T\), with the constant paid almost entirely in matmul. This is why Mamba-2's core layer runs 2 to 8 times faster than Mamba's while remaining competitive with transformers on language modelling, and why it can afford a state dimension over an order of magnitude larger. Chunk size \(Q\) is now a tuning knob trading the quadratic term against recurrence overhead, in the same way FlashAttention tiles are tuned.

What the duality gives up

The equivalence holds for the scalar-times-identity case. A general diagonal \(\bar{A}\), where each of the \(N\) state dimensions decays at its own input-dependent rate, is strictly more expressive and does not reduce to a single attention mask. Mamba-2 accepts that restriction deliberately; Mamba-1 did not. So SSD is not a proof that selective SSMs are linear attention. It is a proof that a well-chosen subclass is, and that the subclass is fast enough to be worth the expressivity it forfeits.

When it breaks

Chunk size interacts with decay. With aggressive decay inside a chunk, the inter-chunk term contributes almost nothing and a larger \(Q\) is nearly free. With slow decay, the low-rank inter-chunk path carries real signal and numerical care is needed on the accumulated products, which can underflow across long chunks in low precision. Production kernels keep the decay accumulation and the state in fp32 for this reason.

The attention view invites bad intuitions. Reading \(M_{ij}\) as attention suggests every position can attend to every earlier position with a learned weight, but the weight is forced to factor through a rank-\(N\) bottleneck and a monotone decay product. The matrix has \(T^2\) entries and roughly \(TN\) degrees of freedom. That bottleneck is the capacity limit discussed in State Capacity and Associative Recall, now visible as a structural constraint on the attention matrix rather than as a property of a recurrence.

Numerical results differ between the two forms. The linear and quadratic paths are mathematically identical and floating-point different, so a model trained with the chunked kernel and served with a pure recurrence will not produce bit-identical logits. That matters for reproducibility testing and for any evaluation harness comparing training-time and inference-time outputs.