Long sequences: state-space models, linear attention, and subquadratic architectures

Why this subject matters now

For five years the answer to "how do I model a sequence" was a single word: attention. The transformer's self-attention is exact, fully parallel over the sequence during training, and it has no recurrent bottleneck, which is exactly why it displaced the LSTM. Its cost, however, is quadratic in sequence length, and that cost stopped being an abstraction the moment context windows grew past a few thousand tokens. A practitioner today is expected to know not just the transformer block but the entire subquadratic frontier that grew up to work around it: state-space models, linear attention, and the gated and selective variants that sit between them. Five years ago this material was a niche; it is now the architecture of shipped long-context models.

The quadratic cost is worth seeing as a number rather than a symbol. The attention benchmark in this repository, run on an NVIDIA H100 80GB HBM3, times a single-head attention at increasing sequence length. At 512 tokens a naive PyTorch attention (materialize the scores, softmax, multiply) takes 0.371 ms; the fused FlashAttention kernel takes 0.05 ms, a 7.4× speedup. By 8,192 tokens the naive kernel takes 98.9 ms against FlashAttention's 3.58 ms, a 27.7× gap, and the peak memory of the naive path has climbed to 35 GB while the fused path holds at 0.57 GB. At 16,384 tokens the naive kernel does not run at all: the score matrix alone is 68.72 GB and does not fit, so the benchmark records an out-of-memory failure. FlashAttention still runs there in 13.7 ms, but note what fusion did and did not fix. It removed the \(O(N^2)\) memory by never materializing the score matrix, and it moved the computation to the fast on-chip SRAM, but it did not change the \(O(N^2)\) arithmetic. FlashAttention is derived in full at building a language model from scratch and attention and the transformer block; this page begins where that one ends, with the observation that no kernel can fix quadratic compute and that the fix has to be a different algorithm.

The different algorithm is old. A linear time-invariant system, the workhorse of control theory and signal processing since the 1960s, maps an input sequence to an output in linear time and constant memory through a hidden state. The contribution of the last four years was to make such systems trainable at scale, numerically stable over tens of thousands of steps, and competitive in quality with attention. That began with HiPPO and S4 from Gu, Goel, Ré and collaborators, matured into the diagonal simplifications S4D and DSS and the scan-based S5, and arrived, with Gu and Dao's Mamba, at a selective recurrence that beat transformers of the same size on language modeling while running in linear time. In parallel a separate community reached similar architectures from the attention side by replacing the softmax with a kernel feature map, producing linear transformers, Performer, RetNet, RWKV, and gated linear attention. The two lines have now largely converged: a selective state-space model and a gated linear-attention layer are, up to parameterization, the same recurrence. Knowing that convergence, and knowing exactly where these models still lose to attention, is the price of admission to a conversation about long-context architecture today.

The binding constraint, quantified

Fix a single attention head with head dimension \(d\) and sequence length \(N\). Self-attention forms scores \(S = QK\T \in \R^{N\times N}\), applies a row softmax, and multiplies by \(V\): \(\text{out} = \softmax(QK\T/\sqrt{d})\,V\). The two matrix multiplies each cost \(2N^2 d\) floating-point operations, so the arithmetic is \(\Theta(N^2 d)\) and, decisively, the score matrix is \(\Theta(N^2)\) in memory. That memory term is what breaks first. In float32 a single \(N\times N\) matrix at \(N = 16{,}384\) is \(16384^2 \times 4\) bytes \(= 1.07\) GB per matrix, and the benchmark's 68.72 GB figure reflects the several score-sized intermediates a naive forward and backward pass keeps alive at once. The lesson is not that attention is slow but that its memory and compute both grow with the square of the thing users most want to grow: context length.

Every architecture on this page attacks the \(N^2\) directly. A recurrent formulation carries a fixed-size state \(h\in\R^{n}\) forward one step at a time, costing \(O(n)\) per step and \(O(Nn)\) total, with \(O(n)\) memory that does not depend on \(N\). The catch, historically, was that a serial recurrence cannot be parallelized across the sequence the way attention can, so training was slow and gradients vanished. The state-space and linear-attention families remove both objections: their recurrences are linear, which makes them expressible either as a global convolution computable by FFT or as an associative scan computable in \(O(\log N)\) depth, so training regains parallelism while inference keeps the constant-memory recurrent form.

The continuous state-space model

A single-input single-output linear time-invariant state-space model is the pair of equations

where \(u(t)\in\R\) is the scalar input signal, \(x(t)\in\R^{n}\) is the hidden state, \(y(t)\in\R\) is the output, and \(A\in\R^{n\times n}\), \(B\in\R^{n\times1}\), \(C\in\R^{1\times n}\), \(D\in\R\) are the parameters. The state \(x\) is a running summary of the input's history; \(A\) governs how that summary decays and rotates, \(B\) writes the current input into it, and \(C\) reads a scalar out. The feedthrough \(D\) is a skip connection from input to output and is usually kept as a learnable scalar or dropped. In a deep network one such system is instantiated per channel of the feature vector, so a layer with model width \(H\) holds \(H\) independent SSMs, each with its own small state.

The solution of the state equation is the variation-of-constants formula \( x(t) = e^{A(t-t_0)}x(t_0) + \int_{t_0}^{t} e^{A(t-\tau)}B\,u(\tau)\,d\tau \), which is worth keeping in mind because the discretization below is nothing more than this formula evaluated over one time step under an assumption about \(u\). The matrix exponential \(e^{At}\) is the continuous analogue of raising a transition matrix to a power, and the eigenvalues of \(A\) set the timescales: an eigenvalue \(\lambda\) contributes a mode \(e^{\lambda t}\), decaying if \(\operatorname{Re}\lambda < 0\) and oscillating at frequency \(\operatorname{Im}\lambda\). Choosing \(A\) is therefore choosing what range of timescales the state can remember, which is precisely what HiPPO will do in a principled way.

Discretization by zero-order hold, derived

Sequence data is discrete: a token at each of \(N\) positions. To run a continuous SSM on it, discretize with a step size \(\Delta > 0\), sampling the input at times \(t_k = k\Delta\). The question is what recurrence in the samples \(u_k = u(t_k)\) reproduces the continuous system. Zero-order hold answers it by assuming the input is piecewise constant between samples: \(u(\tau) = u_k\) for \(\tau \in [t_k, t_{k+1})\). This is exactly the assumption a digital-to-analog converter makes, and it is exact for held inputs rather than a first-order approximation.

Apply the variation-of-constants formula over one step, from \(t_k\) to \(t_{k+1} = t_k + \Delta\), with \(u\) held at \(u_k\):

where the substitution \(s = t_{k+1}-\tau\) turned the convolution integral into a plain integral of the matrix exponential over \([0,\Delta]\), and \(u_k\) came out of the integral because it is constant. The remaining integral has a closed form when \(A\) is invertible: since \(\frac{d}{ds}\big(A^{-1}e^{As}\big) = e^{As}\),

Writing \(x_k\) for \(x(t_k)\), the discrete system is the recurrence

The two forms of \(\bar{B}\) are equal because \(A^{-1}\) and \(e^{A\Delta}-I\) are both power series in \(A\) and therefore commute. Two facts about this map matter later. First, \(\Delta\) is a genuine parameter, not just a bookkeeping constant: it controls how much of one continuous time unit each discrete step advances, and in the selective models it will be made input-dependent, which is the entire trick. Second, a common and cheaper choice keeps \(\bar{A} = e^{\Delta A}\) but approximates \(\bar{B}\approx\Delta B\), the first term of the expansion \(A^{-1}(e^{\Delta A}-I) = \Delta I + \tfrac{\Delta^2}{2}A + \cdots\); Mamba uses this simplified \(\bar{B} = \Delta B\), which is accurate for small \(\Delta\) and removes an inverse.

Two views of the same object: recurrence and convolution

Unroll the discrete recurrence from a zero initial state. Substituting repeatedly, \(x_k = \sum_{j=0}^{k} \bar A^{\,k-j}\bar B\,u_j\), and reading out,

This is a causal convolution of the input with the fixed kernel \(\bar K = (\bar K_0, \bar K_1, \dots)\), where \(\bar K_i = \bar C \bar A^{i}\bar B\) is a scalar (an inner product with the \(i\)-step-powered transition). The two views are mathematically identical: the recurrence computes \(y\) left to right in \(O(Nn)\) time and \(O(n)\) memory, while the convolution computes the same \(y\) as one length-\(N\) filter, computable by FFT in \(O(N\log N)\) time. The recurrent view is what you want at inference, where tokens arrive one at a time and only the current state need be kept; the convolutional view is what you want at training, where the whole sequence is present and a single parallel FFT beats \(N\) serial steps. A state-space layer therefore trains as a convolution and generates as a recurrence, from the same parameters. The one hard part is the kernel: computing \(\bar K_i = \bar C\bar A^{i}\bar B\) for \(i = 0,\dots,N-1\) naively means \(N\) matrix powers, \(O(Nn^2)\), and this is the bottleneck S4 exists to remove.

The equivalence is exact and can be checked at machine precision. For a diagonal transition \(\bar A = \diag(\bar a_1,\dots,\bar a_n)\), the kernel is a sum of geometric series, \(\bar K_i = \sum_{k=1}^{n} \bar C_k\, \bar a_k^{\,i}\,\bar B_k\), and a short program confirms the recurrent and convolutional outputs agree to \(4\times10^{-16}\). The implementation section runs exactly that check in both PyTorch and JAX.

The three computational forms a linear SSM can take, and when each is used, are worth keeping in one picture. The same parameters \((\bar A, \bar B, \bar C)\) drive all three.

   parameters:  A, B, C, Delta   --ZOH-->   Abar = exp(dt A),  Bbar = (Abar-I)A^-1 B

   RECURRENT                 CONVOLUTIONAL                  ASSOCIATIVE SCAN
   x_k = Abar x_{k-1}+Bbar u_k    K_i = C Abar^i Bbar            e_k = (Abar_k, Bbar_k u_k)
   y_k = C x_k                    y   = K * u  (causal)          x_k = e_k o ... o e_0
   ------------------------       ------------------------       ------------------------
   O(N n) time, O(n) state        O(N log N) via FFT             O(N) work, O(log N) depth
   serial, 1 token at a time      needs time-INVARIANT A         survives time-VARYING A
   used at INFERENCE              used at TRAINING (S4/S4D)      used at TRAINING (S5/Mamba)
   

The convolution requires a single fixed kernel and therefore a time-invariant system, which is why it is the S4 training path; the scan needs only associativity of the affine operator and therefore survives the input-dependent transitions of a selective model, which is why it is the Mamba training path. The recurrence is the common inference form for all of them.

HiPPO: choosing \(A\) so the state remembers

A random or poorly chosen \(A\) forgets too fast or blows up; the modes \(e^{\lambda t}\) either decay before the state accumulates anything useful or diverge. HiPPO, from Gu, Dao, Ermon, Rudra, and Ré (2020), constructs \(A\) so the state is the optimal online summary of the entire input history. The framing is function approximation. At each time \(t\), the input seen so far is a function \(u_{\le t}\) on \([0,t]\); project it onto the first \(n\) orthogonal polynomials with respect to a measure on that interval, and let the state \(x(t)\in\R^{n}\) be the coefficients of that projection. HiPPO derives the ordinary differential equation the coefficients must satisfy for the projection to stay optimal as \(t\) advances, and that ODE is exactly \(\dot x = A x + B u\) for specific \(A, B\). The state is then, by construction, a compressed reconstruction of the whole past.

The choice of measure sets what "remember" means. The scaled Legendre variant, HiPPO-LegS, uses a uniform measure over the growing window \([0,t]\), which weights all of history equally and rescales as time passes, giving timescale-invariant memory. Its matrix is lower-triangular plus diagonal,

Two properties of LegS are what make it valuable, and both are proven in the HiPPO paper rather than reproduced here. First, the reconstruction error of the length-\(n\) state decreases as \(n\) grows and does not depend on where in the (arbitrarily long) history a feature sits, which is the formal sense in which it gives long-range memory: an event far in the past is represented as accurately as a recent one. Second, the discretized LegS recurrence has gradients whose magnitude is bounded independently of the sequence length, so it does not suffer the vanishing-gradient pathology that killed plain RNNs on long sequences. The intuition is that the state is not trying to carry information forward by repeated multiplication, which is what vanishes; it is holding an explicit polynomial reconstruction whose coefficients are updated by a well-conditioned linear map. S4 and every model after it initialize \(A\) at this matrix (or a diagonal approximation of it) precisely to inherit these two properties, then let training adjust from there.

S4: making the convolution kernel cheap

S4, from Gu, Goel, and Ré (2021), is the model that made a HiPPO-initialized SSM trainable at scale. The obstacle it removes is the kernel cost. Computing \(\bar K_i = \bar C\bar A^{i}\bar B\) for all \(i < N\) by repeated multiplication is \(O(N n^2)\) and numerically fragile, and the LegS \(A\) is not normal, so it cannot simply be diagonalized in a stable basis. S4's structural observation is that the LegS matrix is normal plus low-rank: it can be written \(A = V\Lambda V^{*} - PQ^{\T}\) with \(\Lambda\) diagonal (the normal part) and \(P,Q\) of small rank (here rank one). Conjugating by \(V\) turns this into a diagonal plus low-rank (DPLR) system, \(A = \Lambda - PQ^{\T}\), which is the form S4 actually parameterizes and trains.

With a DPLR \(A\), the kernel is computed not directly but through its generating function, the truncated \(z\)-transform \(\hat K(z) = \sum_{i=0}^{N-1} \bar K_i z^{i}\). Evaluating \(\hat K\) at the \(N\)-th roots of unity and taking an inverse FFT recovers the kernel \(\bar K\); this replaces matrix powers with an evaluation of a rational function. The generating function of a DPLR SSM has a closed rational form, and after applying the Woodbury identity to fold in the low-rank correction, each evaluation reduces to sums of the shape

which is a Cauchy matrix-vector product: the matrix \(1/(\omega_j - \lambda_k)\) is a Cauchy matrix. What this sum computes is the SSM's transfer function sampled at the FFT frequencies; what makes it fast is that Cauchy matrix-vector products have stable near-linear algorithms, so the whole kernel construction runs in \(\tilde O(n + N)\) rather than \(O(Nn^2)\). The full derivation, including the Woodbury reduction and the stability argument, is long and lives in the S4 paper's appendices; the load-bearing idea to carry away is that DPLR structure turns "power a matrix \(N\) times" into "evaluate a rational function at \(N\) points," and rational-function evaluation is a Cauchy kernel. S4 with this machinery was the first model to solve the hardest Long Range Arena task, Path-X at length 16,384, which no transformer variant had done, and that result is what put state-space models on the map.

S4D and DSS: dropping the low-rank term

The DPLR bookkeeping is intricate, and two papers in 2022 showed most of it is unnecessary. DSS, from Gupta, Gu, and Berant, parameterizes a purely diagonal state matrix and computes the kernel through a softmax-normalized combination of the diagonal modes, and it matched S4 on Long Range Arena. S4D, from Gu, Gupta, Goel, and Ré, gave the clean analysis: keep only the diagonal part of the DPLR system, so \(A = \Lambda = \diag(\lambda_1,\dots,\lambda_n)\), and the kernel becomes an explicit Vandermonde computation,

which needs no Cauchy machinery at all: it is a sum of \(n\) geometric sequences, computable as a Vandermonde matrix-vector product in \(O(nN)\) or faster. The remaining question was initialization, since the long-range memory came from the specific structure of the HiPPO matrix, and throwing away the low-rank part throws some of it away. S4D answered it with principled diagonal initializations, S4D-Lin placing the eigenvalues on a line in the left half-plane and S4D-Inv approximating the diagonal part of the HiPPO-LegS spectrum, and showed these recover almost all of S4's quality. The practical consequence is large: essentially every state-space model since is diagonal, because a diagonal SSM is a per-dimension scalar recurrence, trivially parallelizable and free of the Cauchy kernel's complexity. S5, from Smith, Warrington, and Linderman (2022), took the next step by using a single multi-input multi-output diagonal SSM computed with a parallel associative scan instead of a convolution, which is the bridge to how Mamba is computed.

The associative scan and parallel prefix

The diagonal recurrence \(x_k = \bar A\, x_{k-1} + \bar b_k\), with \(\bar b_k = \bar B u_k\), looks serial, but a linear recurrence is a prefix computation under an associative operator, so it parallelizes in logarithmic depth. Represent each step as the affine map \(e_k = (\bar A_k,\, \bar b_k)\) acting by \(x \mapsto \bar A_k x + \bar b_k\). Composing two steps, first \(e_1\) then \(e_2\), gives \(x \mapsto \bar A_2(\bar A_1 x + \bar b_1) + \bar b_2 = (\bar A_2 \bar A_1)x + (\bar A_2 \bar b_1 + \bar b_2)\), so the composition operator is

This operator is associative, which is what a parallel scan requires, but it is not commutative: the transition matrices multiply, and matrix (or, for time-varying scalars, the ordered) product depends on order. Getting the operand order right, later step on the left acting after the earlier step, is the single most common bug in a hand-written scan, and Problem 3 shows the wrong order produces a completely different and silently plausible sequence. Given associativity, the prefix \(x_k = e_k \bullet e_{k-1} \bullet \cdots \bullet e_0\) applied to \(x_{-1}=0\) is computed by a Blelloch scan in \(O(N)\) total work and \(O(\log N)\) depth, versus the \(O(N\log N)\) work of a Hillis-Steele scan; the work-depth analysis of both, and the measured cost of the tradeoff on this machine, is derived at parallel computing and the memory hierarchy. The point for this page is that the scan gives a state-space layer the same sequence-parallel training as the FFT convolution, and unlike the FFT it survives when the recurrence stops being time-invariant, which is exactly what happens in Mamba.

Mamba: selective, input-dependent state spaces

Every model so far is linear time-invariant: \(\bar A, \bar B, \bar C, \Delta\) are fixed across the sequence, which is what let the recurrence become a fixed convolution. That time-invariance is also a weakness. An LTI SSM processes every token with the same filter, so it cannot decide, based on content, to remember this token and forget that one; its dynamics are chosen at initialization and frozen per position. Mamba, from Gu and Dao (2023), makes the system selective by letting the parameters depend on the input at each step:

with \(A\) kept as a fixed diagonal (per channel) and discretized on the fly, \(\bar A_t = e^{\Delta_t A}\) and \(\bar B_t = \Delta_t B_t\). The recurrence \(x_t = \bar A_t x_{t-1} + \bar B_t u_t\), \(y_t = C_t x_t\) is now linear time-varying. This is called the S6 layer, a selective S4. The selectivity is what closes the quality gap with attention: an input-dependent \(\Delta_t\) acts as a content-based gate, a large \(\Delta_t\) advancing the dynamics a long way and effectively resetting the state to focus on the current token, a small \(\Delta_t\) barely moving it and thus holding the past. Input-dependent \(B_t\) and \(C_t\) let the model choose what to write and read as a function of content, which is exactly the ability an LTI filter lacks.

Selectivity has a cost: it breaks the convolutional view. A global convolution kernel exists only for a time-invariant system, because the kernel \(\bar K_i = \bar C \bar A^{i}\bar B\) assumes the same \(\bar A\) at every step; once \(\bar A_t\) varies with \(t\), there is no single filter and the FFT training path is gone. The recurrence, however, is still linear in the state, so it is still an affine scan, only now with a per-step transition \(\bar A_t\). Mamba therefore computes with the associative scan of the previous section rather than a convolution. The remaining problem is memory. The selective state is expanded: for a channel of input, the state has \(n\) dimensions (the SSM state size, typically 16), so materializing the full sequence of states is \(n\) times larger than the input, and writing it to and from high-bandwidth memory would dominate. Mamba's hardware-aware selective scan fuses the discretization, the scan, and the output projection into one kernel that keeps the expanded state in on-chip SRAM and never writes it to HBM, recomputing what it needs in the backward pass, the same recomputation strategy FlashAttention uses. That fusion is why Mamba is fast in practice despite the scan; it is an engineering match to the memory hierarchy, not a change to the algorithm's asymptotics.

Mamba-2 and state-space duality

Mamba-2, from Dao and Gu (2024), restricts the transition further, to a scalar times identity per head, \(\bar A_t = a_t I\) with \(a_t\in\R\), and in exchange exposes a duality with attention that unlocks the tensor cores. With a scalar transition, unroll the recurrence and read out: the output is \(y_t = \sum_{s\le t}\big(\prod_{r=s+1}^{t} a_r\big) (C_t^{\T} B_s)\, u_s\). Collect the coefficients into an \(N\times N\) lower-triangular matrix \(M\) with \(M_{ts} = \big(\prod_{r=s+1}^{t} a_r\big) C_t^{\T}B_s\) for \(s\le t\); then \(y = M u\). The scalar products form a rank-structured (semiseparable) matrix, and the whole layer is the single masked matrix multiply \(y = M u\), which is exactly the shape of masked attention with scores \(C_t^{\T}B_s\) and a decay mask \(\prod a_r\) in place of the softmax. This is state-space duality (SSD): a scalar-transition selective SSM and a linear-attention layer with a decay mask are the same computation viewed two ways. The payoff is the SSD algorithm, which blocks the sequence into chunks, computes the diagonal blocks as small dense attention matmuls (tensor-core friendly) and the off-diagonal contributions by passing a chunk-level state forward, combining the quadratic-within-chunk efficiency of attention with the linear-across-chunks efficiency of the recurrence. Mamba-2 is both simpler and several times faster to train than Mamba while matching or beating it, and the SSD viewpoint is the cleanest bridge between the two halves of this page.

Linear attention as a recurrent network

The other route to a linear-time sequence model starts from attention and removes the softmax. Softmax attention computes, for query \(i\),

and the quadratic cost lives entirely in the fact that the score \(\exp(q_i^{\T}k_j)\) couples every query to every key, so it cannot be factored across \(i\) and \(j\). Katharopoulos, Vyas, Pappas, and Fleuret (2020) observed that if the exponential is replaced by a factorable similarity \(\phi(q_i)^{\T}\phi(k_j)\) for some feature map \(\phi:\R^{d}\to\R^{m}\), the sums reorganize. Substitute and use associativity of the sums:

where \(S_i = \sum_{j\le i}\phi(k_j)v_j^{\T}\in\R^{m\times d}\) and \(z_i = \sum_{j\le i}\phi(k_j)\in\R^{m}\). The pair \((S_i, z_i)\) is a state that satisfies the recurrence \(S_i = S_{i-1} + \phi(k_i)v_i^{\T}\), \(z_i = z_{i-1} + \phi(k_i)\). This is exactly an RNN: a fixed-size matrix state updated by a rank-one addition per token, read out by a query. The cost is \(O(Nmd)\) time and \(O(md)\) memory, linear in \(N\), and inference carries only the \(m\times d\) matrix \(S\) forward with no growing KV cache. The move that made this possible was moving the parentheses: \((\phi(q)^{\T}\phi(k))v = \phi(q)^{\T}(\phi(k)v^{\T})\), an associativity trick that trades the \(N\times N\) score matrix for a \(d\times d\) state. Problem 2 quantifies exactly when that trade wins.

Why not simply pick \(\phi\) so that \(\phi(q)^{\T}\phi(k) = \exp(q^{\T}k)\) and recover softmax exactly? Because no finite feature map does. The Gaussian and exponential kernels have infinite-dimensional reproducing-kernel Hilbert spaces; \(\exp(q^{\T}k)\) is an inner product only in an infinite-dimensional feature space, so any finite \(m\) is an approximation. Two responses define the sub-field. Katharopoulos and colleagues abandon softmax and use a cheap deterministic map, \(\phi(x) = \operatorname{elu}(x)+1\), which is positive but is a different (non-softmax) attention. Performer, from Choromanski and colleagues at DeepMind (2020), keeps softmax as the target and approximates it with random features (FAVOR+): a randomized \(\phi\) of dimension \(m\) whose inner product is an unbiased estimator of \(\exp(q^{\T}k)\), with error that shrinks as \(m\) grows. The choice is the usual one: an exact-but-different objective, or an unbiased-but-noisy approximation to the original.

RetNet and RWKV: decay and retention

Plain linear attention with an unnormalized cumulative sum tends to let the state \(S_i\) grow without bound over a long sequence, and it weights the distant past as heavily as the present. Adding a decay fixes both and connects the two halves of this page, because a decay is a scalar transition \(\bar A\). RetNet, from Sun and colleagues at Microsoft Research (2023), replaces the running sum with a decayed one, \(S_i = \gamma S_{i-1} + k_i v_i^{\T}\) for a per-head constant \(\gamma\in(0,1)\), which is precisely a diagonal SSM with scalar transition \(\gamma\) and identity feature map. Retention comes in three equivalent computational forms, exactly the recurrence/convolution/chunk trichotomy above: a parallel form for training that materializes a decay-masked score matrix \(D_{ij}=\gamma^{\,i-j}\) for \(i\ge j\), a recurrent form for \(O(1)\)-state inference, and a chunkwise form that is quadratic within a chunk and recurrent across chunks, the same block decomposition Mamba-2's SSD uses.

RWKV, from Peng and a large open collaboration (2023), reaches a similar place from the RNN side. It is an attention-free recurrent network with a per-channel time-decay \(w\) and a bonus \(u\) that gives the current token extra weight, plus a token-shift that mixes each position with its predecessor. Its "WKV" operator is a decayed weighted sum over the past, computable both as a parallelizable form for training and as a strict \(O(1)\)-state recurrence for generation, which is what lets RWKV serve as a drop-in transformer replacement at constant inference memory. The early versions (through RWKV-4) use a scalar per-channel state; RWKV-5 and RWKV-6 (Eagle and Finch) move to a matrix-valued state, which makes them, structurally, gated linear attention. The convergence is the recurring theme: retention, RWKV's WKV, and a scalar-transition selective SSM are the same recurrence wearing different notation.

Gated linear attention and the state-size tradeoff

A fixed scalar decay \(\gamma\) is to RetNet what a fixed \(A\) was to S4: it cannot choose, per token, what to keep. Gated linear attention (GLA), from Yang, Wang, Shen, Panda, and Kim (2023), makes the decay input-dependent, exactly the selectivity move Mamba made. The state update becomes \(S_t = \diag(\alpha_t)\,S_{t-1} + k_t v_t^{\T}\) with a data-dependent gate \(\alpha_t\in(0,1)^{m}\) computed from the input, a per-dimension forget gate on the memory. This is the linear-attention twin of the selective SSM: an input-dependent diagonal transition. GLA's contribution beyond the formulation is a hardware-efficient chunkwise algorithm that keeps the data-dependent gates while still using matmuls on tensor cores, which is what the flash-linear-attention library implements for GLA, RetNet, Mamba-2's SSD, and their relatives under one roof. Whether one calls the diagonal, input-dependent transition a "selective SSM" or a "gated linear attention" is now a matter of which community's notation is in front of you.

All of these models share one structural limit, and it is the right note to end the theory on. Their state is a fixed \(m\times d\) matrix, independent of sequence length. That is the source of their efficiency and the source of their one real weakness: a fixed matrix can hold only a bounded number of distinct key-value associations before they collide. To recall an arbitrary earlier token exactly, the model must have stored it, and a rank-\(m\) memory can store at most \(m\) linearly independent keys' values without interference, as Problem 4 makes precise. Softmax attention has no such limit because its "state," the KV cache, grows with the sequence; that is why it pays \(O(N)\) memory and why it recalls perfectly. The entire design space of this page is a single tradeoff: how large a fixed state to carry, trading recall capacity against memory and compute. Hybrids, next, refuse to choose.

Where pure models break, and the hybrids

The diagnostic task is associative recall: present a stream of key-value pairs, then a query key, and ask for its value. Its multi-query version, MQAR, studied by Arora and colleagues in the Zoology work (2023), is the clean stress test, and it exposes a consistent ordering. Softmax attention solves it at any length because it can attend directly to the matching key. Pure linear-attention and pure SSM layers solve it only up to a capacity set by their state size, and Zoology showed MQAR accuracy scaling with model state dimension: to match attention, a linear model needs a state large enough to hold the associations, which erodes the efficiency advantage. Selectivity helps, which is much of why Mamba beat earlier SSMs on language, but it does not remove the ceiling; the hardest exact-copy and in-context retrieval tasks still favor attention. This is not a tuning problem, it is the capacity argument of Problem 4.

The production answer is to interleave. Jamba, from Lieber and colleagues at AI21 (2024), builds a large model mostly out of Mamba layers with a full-attention layer inserted every few blocks (and mixture-of-experts on top), so that the bulk of the sequence mixing is subquadratic while a small number of attention layers restore exact recall and long-range retrieval. The empirical finding across several such studies is that a surprisingly small fraction of attention layers, often under one in eight, recovers most of the recall quality while keeping most of the throughput and memory advantage of the state-space backbone. Related hybrids reach the same conclusion from the other side: the Based architecture from the Zoology group pairs a short sliding-window softmax attention, which handles exact local recall, with a linear attention using a Taylor-series feature map, which handles the long tail cheaply. The consensus of the frontier is that neither pure attention nor a pure fixed-state recurrence is the right answer at scale; a mixture that spends quadratic cost only where recall demands it is.

Worked problems

Implementation

The first block implements a minimal diagonal state-space layer and checks the recurrent and convolutional forms against each other, exactly Problem 5. The PyTorch tab computes the recurrence as a serial scan and the kernel by the Vandermonde sum; the JAX tab computes the same recurrence with jax.lax.associative_scan, using the non-commutative affine operator of Problem 3, and also builds the convolution, so the two tabs cover both the serial and the parallel-scan routes. All three outputs agree to machine precision.

import torch

# Minimal diagonal (real) SSM. Complex A is the usual choice in practice;
# real diagonal keeps the check readable. Shapes in comments.
torch.manual_seed(0)
n, L = 4, 16                          # state size, sequence length
A   = -torch.rand(n) - 0.1            # (n,)  stable: negative real parts
B   = torch.randn(n)                  # (n,)  input -> state
C   = torch.randn(n)                  # (n,)  state -> output
dt  = 0.5                             # step size (Delta)

Abar = torch.exp(dt * A)             # (n,)  = exp(Delta A), ZOH transition
Bbar = (Abar - 1.0) / A * B          # (n,)  = (Abar - I) A^-1 B

u = torch.randn(L)                    # (L,)  input signal

# --- recurrent form: serial scan, O(L n) time, O(n) state ---
def ssm_recurrent(u, Abar, Bbar, C):
    x = torch.zeros_like(Abar)        # (n,)
    ys = []
    for k in range(u.shape[0]):
        x = Abar * x + Bbar * u[k]    # x_k = Abar x_{k-1} + Bbar u_k
        ys.append((C * x).sum())      # y_k = C x_k
    return torch.stack(ys)            # (L,)

# --- convolutional form: kernel K_i = sum_r C_r Abar_r^i Bbar_r, then convolve ---
def ssm_conv(u, Abar, Bbar, C):
    L = u.shape[0]
    powers = Abar.unsqueeze(0) ** torch.arange(L).unsqueeze(1)  # (L, n)
    K = (powers * (C * Bbar)).sum(dim=1)                        # (L,) kernel
    y = torch.zeros(L)
    for k in range(L):                                          # causal conv
        y[k] = (K[:k + 1].flip(0) * u[:k + 1]).sum()
    return y

y_rec  = ssm_recurrent(u, Abar, Bbar, C)
y_conv = ssm_conv(u, Abar, Bbar, C)
print("max |recurrent - conv| =", (y_rec - y_conv).abs().max().item())
assert torch.allclose(y_rec, y_conv, atol=1e-5)

import jax, jax.numpy as jnp
from jax import lax

key = jax.random.PRNGKey(0)
n, L = 4, 16
kA, kB, kC, kU = jax.random.split(key, 4)
A    = -jax.random.uniform(kA, (n,)) - 0.1     # (n,) stable
B    = jax.random.normal(kB, (n,))             # (n,)
C    = jax.random.normal(kC, (n,))             # (n,)
dt   = 0.5
Abar = jnp.exp(dt * A)                         # (n,)
Bbar = (Abar - 1.0) / A * B                    # (n,)
u    = jax.random.normal(kU, (L,))             # (L,)

# --- parallel associative scan of the affine recurrence ---
# element k is (Abar, Bbar * u_k); compose so the LATER step acts last:
#   combine(left, right) = (aR*aL, aR*bL + bR)   -- NON-commutative
def combine(left, right):
    aL, bL = left
    aR, bR = right
    return aL * aR, aR * bL + bR                # order matters

def ssm_scan(u, Abar, Bbar, C):
    a = jnp.broadcast_to(Abar, (u.shape[0], n)) # (L, n) per-step transitions
    b = Bbar * u[:, None]                        # (L, n) per-step affine terms
    _, x = lax.associative_scan(combine, (a, b)) # x[k] = state after step k
    return (x * C).sum(axis=1)                   # (L,) outputs

# --- convolutional form for cross-check ---
def ssm_conv(u, Abar, Bbar, C):
    powers = Abar[None, :] ** jnp.arange(L)[:, None]   # (L, n)
    K = (powers * (C * Bbar)).sum(axis=1)              # (L,) kernel
    return jnp.array([jnp.dot(K[:k + 1][::-1], u[:k + 1]) for k in range(L)])

y_scan = ssm_scan(u, Abar, Bbar, C)
y_conv = ssm_conv(u, Abar, Bbar, C)
print("max |scan - conv| =", float(jnp.abs(y_scan - y_conv).max()))
assert jnp.allclose(y_scan, y_conv, atol=1e-5)

The second block is a causal linear-attention layer in the running-state form of the theory section, the RNN that a kernelized attention becomes. Both tabs use the \(\operatorname{elu}(x)+1\) feature map, carry the matrix state \(S\) and the normalizer \(z\), and assert equality against the quadratic reference so the associativity trick is verified rather than asserted. This is the computation Problem 2 counts.

import torch
import torch.nn.functional as F

torch.manual_seed(0)
N, d = 32, 8                          # sequence length, head dim
Q = torch.randn(N, d); K = torch.randn(N, d); V = torch.randn(N, d)

def feat(x):                          # feature map phi: positive, deterministic
    return F.elu(x) + 1.0             # (N, d)

# --- linear attention, running-state recurrence: O(N d^2) ---
def linear_attention(Q, K, V):
    Qf, Kf = feat(Q), feat(K)         # (N, d), (N, d)
    S = torch.zeros(d, d)             # (d, d) matrix state  = sum phi(k) v^T
    z = torch.zeros(d)                # (d,)   normalizer     = sum phi(k)
    out = torch.zeros(N, d)
    for i in range(N):
        S = S + torch.outer(Kf[i], V[i])   # rank-one update
        z = z + Kf[i]
        num = Qf[i] @ S               # (d,)
        den = Qf[i] @ z + 1e-6        # scalar
        out[i] = num / den
    return out

# --- quadratic reference: O(N^2 d), for the equality check only ---
def linear_attention_quadratic(Q, K, V):
    Qf, Kf = feat(Q), feat(K)
    scores = Qf @ Kf.t()              # (N, N)
    mask = torch.tril(torch.ones(N, N))
    scores = scores * mask
    num = scores @ V                  # (N, d)
    den = scores.sum(dim=1, keepdim=True) + 1e-6
    return num / den

o1 = linear_attention(Q, K, V)
o2 = linear_attention_quadratic(Q, K, V)
print("max |recurrent - quadratic| =", (o1 - o2).abs().max().item())
assert torch.allclose(o1, o2, atol=1e-5)

import jax, jax.numpy as jnp
from jax import lax, nn

key = jax.random.PRNGKey(0)
N, d = 32, 8
kq, kk, kv = jax.random.split(key, 3)
Q = jax.random.normal(kq, (N, d))
K = jax.random.normal(kk, (N, d))
V = jax.random.normal(kv, (N, d))

def feat(x):
    return nn.elu(x) + 1.0            # (N, d) positive feature map

# --- running-state recurrence via lax.scan ---
def linear_attention(Q, K, V):
    Qf, Kf = feat(Q), feat(K)
    def step(carry, inp):
        S, z = carry                 # S: (d, d), z: (d,)
        qf, kf, v = inp
        S = S + jnp.outer(kf, v)     # rank-one memory update
        z = z + kf
        out = (qf @ S) / (qf @ z + 1e-6)
        return (S, z), out
    init = (jnp.zeros((d, d)), jnp.zeros((d,)))
    _, out = lax.scan(step, init, (Qf, Kf, V))
    return out                        # (N, d)

# --- quadratic reference ---
def linear_attention_quadratic(Q, K, V):
    Qf, Kf = feat(Q), feat(K)
    scores = (Qf @ Kf.T) * jnp.tril(jnp.ones((N, N)))
    return (scores @ V) / (scores.sum(axis=1, keepdims=True) + 1e-6)

o1 = linear_attention(Q, K, V)
o2 = linear_attention_quadratic(Q, K, V)
print("max |recurrent - quadratic| =", float(jnp.abs(o1 - o2).max()))
assert jnp.allclose(o1, o2, atol=1e-5)

How it is done in practice

The distance between these derivations and a deployed model is mostly numerics and memory scheduling. Diagonal SSMs are parameterized with complex eigenvalues in practice, not the real diagonals above, because complex modes \(e^{\lambda t}\) with nonzero imaginary part represent oscillations and cover the LegS spectrum, and the eigenvalues are stored through a parameterization that keeps their real part negative (S4D uses \(-\exp(\cdot)\) for the real part) so the recurrence stays stable across tens of thousands of steps. The step size \(\Delta\) is initialized log-uniformly over a range of timescales, typically \(10^{-3}\) to \(10^{-1}\), so different channels cover different memory lengths; getting that range right matters as much as any other hyperparameter. Mamba's selective scan is a fused CUDA kernel for the same memory-hierarchy reason FlashAttention is: the expanded state, of size (batch, length, channels, state), is \(n\) times the input and must never touch HBM, so the kernel keeps it in SRAM and recomputes it in the backward pass. On this kind of hardware that fusion is the difference between competitive and unusable, the same lesson the fusion benchmark in this repository records for attention, where fusing the projection and attention stages gave a 4.26× speedup.

The efficiency case is strongest at inference, not training. A transformer generating token \(t\) must attend over a KV cache of size \(O(t)\), so per-token latency and memory grow with context, and the KV cache is often the dominant memory cost of serving a long-context transformer. A state-space or linear-attention model carries a fixed state and generates each token in \(O(1)\) time and memory regardless of how long the context is, which is decisive for long-context, high-throughput, or on-device serving. The measured attention costs on this H100 quantify what is being avoided: the naive score matrix that reached 35 GB at 8,192 tokens and 68.72 GB at 16,384 is exactly the object a recurrence never forms. FlashAttention removes that memory during training by tiling, but it cannot remove the growing KV cache at inference, which is the state-space family's structural advantage. The counterweight is training maturity: the transformer's kernels, parallelism, and quantization recipes are years ahead, so at the scales where those matter most the pure state-space model has to overcome a large engineering head start, which is another reason the shipped systems are hybrids.

The current research frontier

The active questions cluster in four places. First is the unification: Dao and Gu's state-space duality showed selective SSMs and linear attention are one family, and a line of work now studies the whole class as linear recurrences with structured, possibly input-dependent transitions, with gated linear attention (Yang and colleagues at MIT and MIT-IBM), the DeltaNet and its parallel training (Yang, Kautz, and colleagues), and Mamba-2 all sitting inside it; the flash-linear-attention library is the shared implementation substrate. Second is expressivity: several groups, including work associated with Merrill and Sabharwal and with the Zoology group at Stanford and Together, have characterized what fixed-state recurrences provably cannot do that attention can, formalizing the recall ceiling of Problem 4 and motivating delta-rule and higher-rank state updates that enlarge effective capacity without a full KV cache. Third is the hybrid recipe: Jamba from AI21, the Zamba models from Zyphra, NVIDIA's Hymba, and IBM's Bamba are all production hybrids interleaving Mamba-style and attention layers, and the open question is the optimal ratio and placement rather than whether to mix at all. Fourth is going beyond sequence models, with state-space backbones for vision (Vision Mamba, VMamba), audio, genomics, and reinforcement learning, where the linear-time long-context property is even more valuable than in text. The competing-approaches picture is genuinely multi-institution: the SSM line runs through work at Stanford, Carnegie Mellon, and Cornell, the linear-attention line through DeepMind, Microsoft Research, MIT, and a large open RWKV community, and the hybrids through several industry labs, with no single group owning the frontier.

Open source to read

• state-spaces/s4: the reference S4, S4D, DSS, and HiPPO implementations from the original authors. Open models/s4/s4.py for the DPLR kernel and the Cauchy-kernel path, and the hippo module for the LegS matrix construction.
• state-spaces/mamba: official Mamba and Mamba-2. Read mamba_ssm/modules/mamba_simple.py for the S6 block and ops/selective_scan_interface.py for how the fused CUDA scan is called; the SSD algorithm lives under mamba_ssm/ops/triton.
• johnma2006/mamba-minimal: a single-file, few-hundred-line PyTorch Mamba with no custom kernels. The clearest way to read the selective recurrence; start at the ssm and selective_scan functions and compare them to the equations above.
• sustcsonglin/flash-linear-attention: Triton kernels for GLA, RetNet, Mamba-2/SSD, DeltaNet, and more under a common chunkwise interface. The place to see the hardware-efficient chunk decomposition; open the fla/ops/gla directory first.
• BlinkDL/RWKV-LM: the RWKV reference across versions. Read the WKV kernel and the token-shift to see the decay-plus-bonus recurrence and how the parallel training form maps to the \(O(1)\) inference form.
• huggingface/transformers: production reference implementations of Mamba, Mamba-2, RWKV, RetNet, and Jamba side by side. Compare models/mamba, models/jamba, and models/rwkv to see how the same recurrence is packaged for each model.

Common misconceptions

"State-space models are just RNNs, so they have the same vanishing-gradient and non-parallel problems." They are linear recurrences, and that word is everything. Linearity lets the recurrence be computed as a convolution or an associative scan, so training is sequence-parallel, unlike a nonlinear RNN; and a HiPPO-initialized \(A\) has gradients bounded independently of length, so the vanishing-gradient pathology that defeated LSTMs on long sequences does not occur. The similarity to an RNN is real only at inference, where both carry a fixed state.

"FlashAttention already solved the long-context problem." FlashAttention removed the \(O(N^2)\) memory by never materializing the score matrix, which is why it runs where the naive kernel OOMs, but it left the \(O(N^2)\) arithmetic untouched, and it cannot shrink the KV cache that grows with context at inference. It is a better way to compute quadratic attention, not a subquadratic algorithm.

"Linear attention is softmax attention made linear, so it should match it." No finite feature map reproduces \(\exp(q^{\T}k)\), whose kernel is infinite-dimensional, so linear attention is a different similarity, not an approximation-free linearization. Performer approximates softmax with random features at the cost of variance; the elu-plus-one variant abandons softmax entirely. Either way the fixed-size state caps recall in a way softmax attention's growing cache does not.

"Selectivity is a minor tweak to S4." It changes the computational class. Fixed parameters make the SSM time-invariant and thus a convolution trainable by FFT; input-dependent parameters make it time-varying, which destroys the convolution and forces the associative scan. That is why Mamba needed a new hardware-aware kernel, and it is the reason selectivity, not raw scale, was what closed the quality gap with transformers.

"Mamba-2 is unrelated to attention." Its scalar-transition recurrence is algebraically a masked matrix multiply with a decay mask, which is exactly linear attention; that is the state-space duality, and it is what lets Mamba-2 use tensor cores through the SSD chunk algorithm instead of a pure scan.

"A big enough state-space model makes attention obsolete." A fixed \(m\times d\) state can hold at most \(m\) independent associations, so exact recall of arbitrary earlier tokens is capacity-limited no matter the scale, which is why associative-recall benchmarks separate these models from attention and why the strongest long-context systems are hybrids that keep a few attention layers precisely for exact recall.

Self-check

References

1. Kalman, R. E. "A New Approach to Linear Filtering and Prediction Problems," Journal of Basic Engineering, 1960. doi:10.1115/1.3662552
2. Oppenheim, A. V. and Schafer, R. W. Discrete-Time Signal Processing, 3rd ed., Pearson, 2009.
3. Vaswani, A. et al. "Attention Is All You Need," NeurIPS, 2017. arXiv:1706.03762
4. Dao, T. et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS, 2022. arXiv:2205.14135
5. Gu, A., Dao, T., Ermon, S., Rudra, A., and Ré, C. "HiPPO: Recurrent Memory with Optimal Polynomial Projections," NeurIPS, 2020. arXiv:2008.07669
6. Gu, A., Goel, K., and Ré, C. "Efficiently Modeling Long Sequences with Structured State Spaces" (S4), ICLR, 2022. arXiv:2111.00396
7. Gupta, A., Gu, A., and Berant, J. "Diagonal State Spaces are as Effective as Structured State Spaces" (DSS), NeurIPS, 2022. arXiv:2203.14343
8. Gu, A., Gupta, A., Goel, K., and Ré, C. "On the Parameterization and Initialization of Diagonal State Space Models" (S4D), NeurIPS, 2022. arXiv:2206.11893
9. Smith, J. T. H., Warrington, A., and Linderman, S. W. "Simplified State Space Layers for Sequence Modeling" (S5), ICLR, 2023. arXiv:2208.04933
10. Gu, A. and Dao, T. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces," 2023. arXiv:2312.00752
11. Dao, T. and Gu, A. "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" (Mamba-2/SSD), ICML, 2024. arXiv:2405.21060
12. Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention," ICML, 2020. arXiv:2006.16236
13. Choromanski, K. et al. "Rethinking Attention with Performers," ICLR, 2021. arXiv:2009.14794
14. Peng, B. et al. "RWKV: Reinventing RNNs for the Transformer Era," EMNLP Findings, 2023. arXiv:2305.13048
15. Sun, Y. et al. "Retentive Network: A Successor to Transformer for Large Language Models" (RetNet), 2023. arXiv:2307.08621
16. Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. "Gated Linear Attention Transformers with Hardware-Efficient Training" (GLA), ICML, 2024. arXiv:2312.06635
17. Lieber, O. et al. "Jamba: A Hybrid Transformer-Mamba Language Model," 2024. arXiv:2403.19887
18. Tay, Y. et al. "Long Range Arena: A Benchmark for Efficient Transformers," ICLR, 2021. arXiv:2011.04006
19. Arora, S. et al. "Zoology: Measuring and Improving Recall in Efficient Language Models," 2023. arXiv:2312.04927
20. Arora, S. et al. "Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff" (Based), ICML, 2024. arXiv:2402.18668
21. Blelloch, G. E. "Prefix Sums and Their Applications," Carnegie Mellon technical report CMU-CS-90-190, 1990. pdf