The Recurrence-Convolution Duality
Why a linear state space layer can be run as a parallel convolution during training and as a constant-memory recurrence at inference, and what the word "linear" is buying.
RNNs train slowly because they are sequential and infer cheaply because their state is fixed size. Transformers train fast because attention is parallel and infer expensively because the KV cache grows with every token. For decades those looked like the two ends of an unavoidable tradeoff. Linear state space models sit on both ends at once, and the reason is a piece of algebra rather than an architectural trick.
The layer
A continuous-time linear state space model maps a scalar input signal to a scalar output through a hidden state:
Discretised with step size \(\Delta\), this becomes a recurrence with matrices \(\bar{A}\) and \(\bar{B}\) derived from \(A\), \(B\) and \(\Delta\):
Run this forward one token at a time and you have an RNN with a fixed-size state, constant per-token cost, and no cache that grows.
Unrolling into a convolution
Because the recurrence is linear and its parameters do not depend on the input, it can be unrolled in closed form. Starting from \(h_{-1} = 0\):
That is a convolution of the input with a kernel \(\bar{K} = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \dots)\) whose length is the sequence length. Once you have \(\bar{K}\), the whole sequence can be computed in one shot with an FFT-based convolution in \(O(L \log L)\) time, fully parallel across positions.
So the same parameters give two computational forms of the identical function. Train with the convolution, because training sees the whole sequence at once and wants parallelism. Infer with the recurrence, because generation is sequential anyway and constant memory is exactly what you want. Nothing is approximated in the switch; the two forms are algebraically equal.
What linearity is paying for
The duality depends entirely on \(\bar{A}\), \(\bar{B}\) and \(C\) being independent of the input. Insert any input-dependent nonlinearity inside the recurrence and \(\bar{A}^j\) no longer factors out, the closed form disappears, and you are back to a sequential RNN. This is the price: the sequence mixing must be linear. Nonlinearity in the architecture comes from what sits between layers, not from inside the scan.
The constraint is also the reason the model can be analysed. A linear time-invariant system has a transfer function, poles, and a frequency response, so the kernel's behaviour over long horizons is a property of the eigenvalues of \(\bar{A}\) rather than an emergent mystery. Eigenvalues inside the unit circle mean memory decays; how fast is set by how close they sit to the boundary.
When it breaks
Computing the kernel naively is the hard part. Forming \(\bar{K}\) requires powers of \(\bar{A}\) up to the sequence length. For a general state matrix of dimension \(N\) over length \(L\), that is \(O(N^2 L)\) work and \(O(NL)\) memory, which is worse than the attention it was meant to replace. The entire technical content of S4 (Gu et al., 2022, Efficiently Modeling Long Sequences with Structured State Spaces, arXiv:2111.00396) is a structured parameterisation of \(A\) that makes this computable, and later work simplified it to a diagonal form that is far easier to implement with almost no quality loss.
Numerical conditioning of \(\bar{A}^j\). Repeated powers either vanish or explode unless the eigenvalues are controlled tightly. Stable implementations parameterise eigenvalues in a form that keeps them inside the unit disc by construction, typically through a negative real part in continuous time, rather than clipping after the fact.
Discretisation is a real choice. Zero-order hold and the bilinear transform give different \(\bar{A}\) from the same \(A\), with different stability properties. The step size \(\Delta\) is a learned per-channel parameter, and it sets the timescale each channel attends to; getting its initialisation range wrong produces a model that either forgets immediately or never forgets, and both look like a capacity problem rather than a scaling problem.
The convolution form has its own memory cost. FFT convolution over a long sequence materialises the full kernel and intermediate spectra, so training memory scales with \(L\) even though inference does not. On very long sequences this becomes the binding constraint, which is part of why chunked and scan-based training formulations replaced pure FFT convolution in later systems.
12 flashcards for this concept
Click a card to reveal the answer.