Transformers and Neural State-Space Architectures › Long-Context Memory and Efficient Recurrence

Michael BrenndoerferJune 26, 20267 min read
World Models HandbookAll chapters
Part 5 of 12World-Model Architectures
Chapter 26Transformers and Neural State-Space Architectures › Long-Context Memory and Efficient Recurrence
Reading section 4 of 10 · Long-Context Memory and Efficient Recurrence · 7 minSection 4 of 10 · 7 min

Book contents

Section 4 of 10 · 7 min read

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:

h˙(t)=A h(t)+B x(t),y(t)=C h(t),\dot{h}(t) = A\, h(t) + B\, x(t), \qquad y(t) = C\, h(t),

where:

  • h(t)∈Rnh(t)\in\mathbb R^n: the continuous-time state
  • x(t)∈Rmx(t)\in\mathbb R^m: the input signal
  • y(t)∈Rpy(t)\in\mathbb R^p: the output
  • A∈Rn×nA\in\mathbb R^{n\times n}: the state-transition generator
  • B∈Rn×mB\in\mathbb R^{n\times m}: the input map into the state derivative
  • C∈Rp×nC\in\mathbb R^{p\times n}: 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 Δ\Delta via the zero-order-hold rule gives the recurrent form used at inference:

ht=Aˉ ht−1+Bˉ xt,yt=C ht,Aˉ=eAΔ,Bˉ=∫0ΔeAuB du.h_t = \bar{A}\, h_{t-1} + \bar{B}\, x_t, \qquad y_t = C\, h_t, \qquad \bar{A} = e^{A\Delta}, \qquad \bar B=\int_0^\Delta e^{Au}B\,du.

where:

  • hth_t: the discrete-time latent state after processing token tt, with the initial state h0h_0 set to zero
  • Aˉ=eAΔ\bar{A} = e^{A\Delta}: the discretized transition matrix, obtained by exact matrix exponentiation of AA over the step size Δ\Delta
  • Bˉ\bar{B}: the discretized input matrix, derived from AA, BB, and Δ\Delta under the zero-order-hold assumption that xx is constant over the step
  • xtx_t: the input token at timestep tt
  • CC: the same output matrix as in continuous time

Here xtx_t is held constant over ((t−1)Δ,tΔ]((t-1)\Delta,t\Delta] and hth_t is the state at the interval's end. Variation of constants gives the integral for Bˉ\bar B: it accumulates the input's contribution while the continuous dynamics evolve. If AA is invertible, Bˉ=A−1(eAΔ−I)B\bar B=A^{-1}(e^{A\Delta}-I)B; the integral also covers singular AA. For a fixed stable mode aa, retention per token is eaΔe^{a\Delta}. A smaller positive step moves this multiplier toward one and, under zero-order hold, also shrinks the input write approximately to BΔB\Delta. A larger step forgets more per token. Learning Δ\Delta therefore changes token-scale dynamics, not the physical-time constant of a fixed AA in isolation.

Because Aˉ=eAΔ\bar{A} = e^{A\Delta}, a fixed positive step Δ\Delta maps each eigenvalue aia_i with negative real part to a discrete multiplier of magnitude ∣eaiΔ∣=eRe⁡(ai)Δ<1|e^{a_i\Delta}|=e^{\operatorname{Re}(a_i)\Delta}<1. 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 AA. 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 [0,t][0,t] uniformly and has time-dependent 1/t1/t 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 h0=0h_0=0, unrolling gives ht=∑k=0t−1AˉkBˉxt−kh_t=\sum_{k=0}^{t-1}\bar A^k\bar Bx_{t-k} and yt=∑k=0t−1Kkxt−ky_t=\sum_{k=0}^{t-1}K_kx_{t-k} with Kk=CAˉkBˉK_k=C\bar A^k\bar B. 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 O(Tlog⁡T)O(T\log T) work for a full sequence, while cached recurrence has O(1)O(1) 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 BtB_t, CtC_t, and the discretization step Δt\Delta_t depend on the current input while retaining a learned input-independent continuous-time AA. Its Appendix B.1 explicitly credits S5's scan and distinguishes Mamba's selection mechanism and hardware-aware implementation. For real ai<0a_i<0 and Δt,i>0\Delta_{t,i}>0, the effective decay Aˉt,i=exp⁡(Δt,iai)\bar A_{t,i}=\exp(\Delta_{t,i}a_i) 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.

The memory-computation exchange

At fixed dimensions, dense full-sequence attention has O(T2)O(T^2) pairwise work, an uncompressed cache has O(T)O(T) storage, and a cached new query has O(T)O(T) attention work. A fixed-width recurrence has O(T)O(T) sequence work and constant state storage; a work-efficient scan preserves linear work, while FFT convolution uses O(Tlog⁡T)O(T\log T). 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 Rdk\mathbb R^{d_k} and values be columns in Rdv\mathbb R^{d_v}. With S0=0S_0=0, the numerator accumulator is St=St−1+ktvt⊤∈Rdk×dvS_t=S_{t-1}+k_tv_t^\top\in\mathbb R^{d_k\times d_v}, and its column readout is St⊤qt=∑j≤t(qt⊤kj)vjS_t^\top q_t=\sum_{j\le t}(q_t^\top k_j)v_j. Normalized linear attention also keeps ct=ct−1+ktc_t=c_{t-1}+k_t and divides by qt⊤ctq_t^\top c_t 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.