Transformers and Neural State-Space Architectures › Long-Context Memory and Efficient Recurrence
World Models HandbookAll chapters Transformers and Neural State-Space Architectures › Long-Context Memory and Efficient RecurrenceOne section · Sections in this chapter4 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
This section develops fixed linear state-space layers, then distinguishes their input-selective relatives and other subquadratic sequence models. Stability and efficient evaluation depend on the particular recurrence and implementation.
For a long trajectory, full attention's pairwise arithmetic and retained KV storage can become expensive. An hour at 30 Hz contains 108,000 steps. At fixed batch size, width, depth, and precision, an append-only KV cache grows linearly with retained tokens. Truncation discards old entries; compression reduces their representation and can lose information, but need not create a globally fixed-size state. A recurrent layer instead maintains a state whose width is independent of sequence length. Whether that state preserves enough information is a task question, not a consequence of efficient implementation alone.
The canonical continuous-time form is the linear state-space system we met during system identification in Part II, now with learned parameters:
where:
- : the continuous-time state
- : the input signal
- : the output
- : the state-transition generator
- : the input map into the state derivative
- : the state readout
The lineage here is worth appreciating. Classical system identification, from subspace methods to prediction-error methods, spent decades estimating exactly these matrices from input-output data for engineering models of known dimension. The neural state-space program takes the same mathematical object and embeds it as a layer inside a deep network, with the matrices as parameters trained by gradient descent end to end. What was once a standalone modeling methodology becomes a differentiable component, and the analysis tools, stability theory in particular, come along for free.
Discretizing with step via the zero-order-hold rule gives the recurrent form used at inference:
where:
- : the discrete-time latent state after processing token , with the initial state set to zero
- : the discretized transition matrix, obtained by exact matrix exponentiation of over the step size
- : the discretized input matrix, derived from , , and under the zero-order-hold assumption that is constant over the step
- : the input token at timestep
- : the same output matrix as in continuous time
Here is held constant over and is the state at the interval's end. Variation of constants gives the integral for : it accumulates the input's contribution while the continuous dynamics evolve. If is invertible, ; the integral also covers singular . For a fixed stable mode , retention per token is . A smaller positive step moves this multiplier toward one and, under zero-order hold, also shrinks the input write approximately to . A larger step forgets more per token. Learning therefore changes token-scale dynamics, not the physical-time constant of a fixed in isolation.
Because , a fixed positive step maps each eigenvalue with negative real part to a discrete multiplier of magnitude . For a fixed diagonal recurrence, the initial-state contribution therefore decays. With bounded inputs and fixed finite input and output matrices, its forced response is bounded as well. This is a property of that recurrent layer, not a guarantee that every activation or gradient in the surrounding network remains bounded. Three design choices help explain the modern SSM lineage:
- Diagonal . S4 (Gu et al., 2021) used a diagonal-plus-low-rank representation; S4D, S5, and Mamba use diagonal state transitions. The scalar channels can be real or complex. Real negative continuous-time eigenvalues yield nonoscillatory decays, whereas complex conjugate channels can encode damped oscillations. A diagonal representation is not intrinsically less expressive than a diagonalizable full matrix when complex coordinates and suitable input/output maps are allowed. The real, positive-decay example below is more restrictive. Negative real parts, a positive discretization step, and bounded forcing give the fixed-layer stability conditions stated above.
- HiPPO-style initialization. HiPPO (High-order Polynomial Projection Operators) derives online dynamics for optimal polynomial projection under a specified time-dependent weighting measure. LegS weights the elapsed interval uniformly and has time-dependent dynamics; LagT uses exponentially recency-weighted history. These are properties of the specified projection operators. S4 uses HiPPO-related structure to motivate its parameterization and initialization, but an arbitrary fixed or subsequently learned SSM does not inherit exact projection optimality.
- Parallel training via convolution. For fixed coefficients and , unrolling gives and with . This matrix-valued impulse response measures each past input's linear influence. Stable modes have decaying envelopes, but their mixture may initially grow or oscillate; it need not be a monotone low-pass curve. At fixed channel dimensions, zero-padded FFT convolution costs work for a full sequence, while cached recurrence has work per new token with respect to history length. This duality applies to the fixed linear layer. A nonlinear RNN lacks one global LTI kernel, though local Jacobians and trajectory perturbations remain inspectable.
A fixed-size state must compress its history, and input dependence changes what the model retains. Classical structured SSM layers use input-independent coefficients. S5 replaces S4's bank of single-input single-output systems with a multi-input multi-output formulation and evaluates its recurrence by parallel scan; it should not be grouped with Mamba as introducing learned input-dependent selectivity. See the S5 paper, introduction and Section 2.2. Mamba makes , , and the discretization step depend on the current input while retaining a learned input-independent continuous-time . Its Appendix B.1 explicitly credits S5's scan and distinguishes Mamba's selection mechanism and hardware-aware implementation. For real and , the effective decay lies strictly between zero and one in exact arithmetic. Complex eigenvalues instead give complex multipliers whose magnitudes are below one. Selection changes both retention and the input-dependent write; a large discretization step can suppress old state while admitting new information. Because these coefficients depend on the sequence, one fixed convolution kernel no longer describes the layer, so Mamba uses a hardware-aware scan. Conditional on a sequence of coefficients, the update is affine in the previous state; the full input-to-output mapping is not a linear time-invariant system. Pointwise decay below one alone does not ensure bounded response to arbitrary forcing: a uniform contraction bound and bounded writes are sufficient, whereas decays approaching one need separate analysis. Finite-precision exponentials and softplus also require numerical care.
At fixed dimensions, dense full-sequence attention has pairwise work, an uncompressed cache has storage, and a cached new query has attention work. A fixed-width recurrence has sequence work and constant state storage; a work-efficient scan preserves linear work, while FFT convolution uses . These costs concern specified implementations, not guarantees of recall. If a task admits an adequate filtering state, recurrence can preserve what matters. Addressable memory can make isolated-event retrieval easier, but a suitable recurrent or symbolic state can also retain such events. This is a design tradeoff to test, not an exclusive division of capability.
Linear transformers make the connection explicit. Let feature-mapped keys and queries be columns in and values be columns in . With , the numerator accumulator is , and its column readout is . Normalized linear attention also keeps and divides by when that denominator is positive. These are statistics sufficient to compute this attention rule, not necessarily the environment's state. RWKV uses recurrent, temporally weighted key-value memory; Transformer-XL reuses preceding-segment representations; Block-Recurrent Transformers update persistent state vectors with attention and gates. Each changes what is retained and how it is read. Tasks with smoothly accumulating evidence and tasks with consequential object relocations suggest different tests, but neither task label establishes an architecture winner.