State Space Duality

By 2024 three lines of work had independently arrived at a linear recurrence with input-dependent decay: selective SSMs from control theory, RWKV from RNNs, gated linear attention from transformers. The resemblance was noted informally. Dao and Gu made it an identity.

Unrolling into a matrix

Take a diagonal selective SSM and substitute the recurrence into itself.

That is a linear map from the input sequence to the output sequence, so it is a matrix. Read off the entries:

$C_t^\top B_s$ is a query–key inner product with $C$ as the query and $B$ as the key. The product of $A$s is lower-triangular by construction — it is empty for $s > t$ — so it is a causal mask whose entries decay rather than being ones and zeros. The whole model is attention without a softmax, under a semiseparable mask.

Two algorithms for one operator

The last row is where the speed comes from. Split the sequence into chunks: within a chunk use the quadratic form, which is a dense matmul and saturates the hardware; between chunks pass a single summarised state, which is the recurrent form. Diagonal blocks quadratic, off-diagonal blocks low-rank.

This is why Mamba-2 can afford a state eight times larger than Mamba-1 and still train faster. Mamba-1’s selective scan was a bespoke CUDA kernel doing elementwise work; the chunked form is torch.matmul on tensor cores, and the hardware gap between those two is roughly an order of magnitude.

import torch
from torch import Tensor


def ssd(x: Tensor, a: Tensor, b: Tensor, c: Tensor, chunk: int = 64) -> Tensor:
    """x: [B, T, D], a: [B, T] decays, b/c: [B, T, N]. Chunked SSD."""
    bsz, t, d = x.shape
    n = b.size(-1)

    xs, as_, bs, cs = (z.reshape(bsz, t // chunk, chunk, -1) for z in (x, a[..., None], b, c))
    cumdecay = as_.cumsum(2)                                  # log-space within chunk

    y, state = [], x.new_zeros(bsz, n, d)
    for i in range(t // chunk):
        # Within the chunk: quadratic form, a plain masked matmul.
        scores = (cs[:, i] @ bs[:, i].transpose(-1, -2))      # [B, C, C]
        mask = (cumdecay[:, i] - cumdecay[:, i].transpose(-1, -2)).exp().tril()
        local = (scores * mask) @ xs[:, i]

        # Across chunks: one recurrent step, carrying a summarised state.
        y.append(local + cs[:, i] @ state)
        decayed = (cumdecay[:, i, -1:] - cumdecay[:, i]).exp()
        state = state * cumdecay[:, i, -1:].exp().unsqueeze(-1) + \
                bs[:, i].transpose(-1, -2) @ (xs[:, i] * decayed)

    return torch.cat(y, dim=1).reshape(bsz, t, d)

Everything in that loop that looks like attention is attention, and everything that looks like a recurrence is one. They are separated by the chunk boundary and nothing else.

The theory transfers freely in both directions now. Multi-head and grouped-query structure carry over to SSMs. Tensor and sequence parallelism developed for transformers apply. And the attention side gains a constant-memory decoding path for any model whose mask is semiseparable.

What the duality does not do is settle whether the softmax matters. Every model on this side of it has a mask that decays with distance, and that is a strictly weaker object than a content-dependent score over all pairs — a decaying mask cannot single out one distant token, however much the model would like to. The persistent recall gap between linear and softmax models is that limitation, and no reformulation removes it. Which is why the frontier models are still hybrids.