Transformers and Neural State-Space Architectures

Imagine a household robot operating in a kitchen. Ten minutes ago, it saw a coffee mug placed behind the microwave. The mug is now out of view. Remembering its location could let the robot fetch it directly; forgetting might force a search. In many partially observable tasks, useful evidence arrives at separated times. A short episode may fit in memory in full, but retaining and consulting an arbitrarily long raw history becomes costly under a finite storage and compute budget.

A recurrent state-space model of the kind we built in the previous chapter carries a fixed-size joint state $(h_t,z_t)$. With finite numerical precision, that state has only finitely many representable values. A task may nevertheless admit a finite-dimensional sufficient statistic, such as a correctly specified Kalman filter's mean and covariance in ideal arithmetic; its numerical implementation is an approximation. The difficulty arises when the retained state merges histories that require different predictions or actions. Contractive linear recurrences attenuate old input contributions; gated and other recurrences can retain or transform them differently. Once a relevant distinction is absent from the state and any external memory, later computation cannot recover it from that state alone.

This chapter studies two architectural families with different memory costs. Transformers retain token representations and use attention to compute content-dependent weighted readouts within the available context. Retention makes access possible; it does not guarantee that the representation preserves a detail or that training learns to retrieve it. Neural state-space layers retain a fixed-size recurrent state, using structured transitions that can be evaluated efficiently. Fixed-coefficient and input-selective variants differ in both their mathematics and implementation. Hybrid sequence architectures combine these mechanisms. We will examine named examples, then compare two small predictors on a controlled task rather than infer a universal ranking.

Throughout, keep our standing distinction in mind: a world model predicts environment-relevant quantities that an agent may use for planning or policy learning. Its performance must be checked under the trajectories where it will be used. Four quality axes help organize that evaluation:

• Predictive fidelity: accuracy of forecasts at specified horizons and under specified conditioning.
• Representation quality: retention of information relevant to the prediction or decision task, not necessarily reconstruction of every physical state variable.
• Uncertainty quality: calibration and informativeness of predictive distributions, not simply a large or small variance.
• Decision usefulness: performance of plans or policies when evaluated in the environment.

These axes are related but not interchangeable. The experiment below measures point-prediction errors and an attention diagnostic. It does not test distant event recall, calibrated uncertainty, or a controller.

This section develops self-attention as the core mechanism of transformer world models: how a trajectory becomes tokens, how attention computes a content-based weighted average over the past, and why that mechanism is so well suited to partially observable decision problems.

A transformer can retain multiple token representations instead of depending only on a single recurrent history summary. For each query it computes a new fixed-width readout from that memory. Tokenization, embeddings, value projections, and the readout can all discard information; this is not lossless storage. One possible tokenization, used below, embeds an observation-action pair. For observation width $d_o$, action width $d_a$, and embedding width $d$, define $o_t\in\mathbb R^{d_o}$, $a_t\in\mathbb R^{d_a}$, and $z_t,p_t\in\mathbb R^d$:

where:

• $o_t$: the observation received at timestep $t$
• $a_t$: the action taken at timestep $t$
• $W_{\text{in}}\in\mathbb R^{d\times(d_o+d_a)}$: a learned embedding of the stacked observation-action pair
• $p_t$: an additive positional encoding for timestep $t$

The projection embeds the pair; it need not reduce dimension or information. Our code expands two inputs to sixteen features. Without positional signals or an order-dependent mask, self-attention is permutation-equivariant: permuting input rows permutes output rows, and the key-value collection has no intrinsic order. A causal mask changes that statement because different positions have different accessible prefixes; repeated content can then produce different outputs even without explicit positions. Positional encodings provide additional temporal information. The original Transformer uses sinusoidal or learned additive encodings. Rotary encodings instead rotate query and key coordinates inside attention, so they are not literally the additive $p_t$ in this equation. Transformer blocks combine attention, feed-forward layers, residual connections, and normalization. Residual paths allow incremental updates and can improve gradient propagation, but they do not guarantee stable gradients.

The heart of the mechanism is scaled dot-product attention. Each token emits three vectors through learned linear maps:

where:

• $Z\in\mathbb R^{T\times d}$: the matrix with row $t$ equal to $z_t^\top$
• $W_Q,W_K\in\mathbb R^{d\times d_k}$ and $W_V\in\mathbb R^{d\times d_v}$: learned projections, giving $Q,K\in\mathbb R^{T\times d_k}$ and $V\in\mathbb R^{T\times d_v}$
• $q_i=Q_{i,:}^\top$, $k_j=K_{j,:}^\top$, and $v_j=V_{j,:}^\top$: column vectors for the query, key, and value; the attention output below is also a column

It is worth pausing on why the same tokens produce three different roles. The query encodes what the current position is looking for; the key encodes what each position advertises about itself, the label under which it can be found; and the value encodes the information each position contributes once selected. Separating these roles, rather than using the raw embedding for all three, is what lets the network learn asymmetric relationships: a token predicting a collision can learn to search for tokens advertising recent velocities, while advertising, under its own key, something unrelated such as its object identity. The attention output for token $i$ is a weighted average over all value vectors (in the causal models trained below, the sum runs only over $j \le i$, but we first present the unmasked form):

where:

• $j$: the index ranging over all tokens in the sequence
• $j'$: a dummy index used only inside the normalizing constant, so that the sum in the denominator is independent of the sum in the numerator
• $q_i^\top k_j$: the dot product between the query of token $i$ and the key of token $j$, measuring how relevant token $j$ is to token $i$
• $\sqrt{d_k}$: the square root of the key dimension $d_k$, by which the scores are scaled to keep the softmax gradients well behaved when $d_k$ is large
• $\sum_{j'} \exp\!\big(q_i^\top k_{j'} / \sqrt{d_k}\big)$: the normalizing constant that turns the scores into convex weights summing to one
• $v_j$: the value vector of token $j$, the content that gets retrieved

For finite logits and a nonempty allowed set, the softmax weights are nonnegative and sum to one. Thus the per-head output before its output projection is a convex combination of value vectors. The unmasked expression includes current and future positions; a causal mask restricts it to $j\le i$. This convexity does not describe the entire block after residual, projection, and feed-forward operations. If query and key components are independent, zero-mean, unit-variance variables, independent across coordinates, then $\operatorname{Var}(q^\top k)=d_k$. Dividing by $\sqrt{d_k}$ gives unit variance under those assumptions. It can reduce softmax saturation; learned projections can still produce saturated scores.

Two properties explain why this matters for world models:

• Content-based addressing. Scores depend on token content and any positional mechanism, rather than only on a fixed decay schedule. An old retained token can receive high weight if the learned query matches it. A Kalman filter instead weights the current innovation using predicted covariance, measurement covariance, and the observation model; its fixed-size state can still be sufficient for its specified task. Attention recomputes compatibility with retained learned keys, not necessarily raw evidence.
• Short gradient paths. Full attention provides a token-to-token path within a layer whose length does not grow with temporal separation. A recurrent path may contain many state-update Jacobians. Attention therefore shortens temporal dependency paths, but saturated scores, depth, and normalization can still cause optimization problems.

The price depends on which operation is measured. At fixed widths and layer count, dense full-sequence attention performs $O(T^2)$ pairwise arithmetic. Materializing its score matrix uses $O(T^2)$ memory, but this is not an intrinsic memory lower bound: FlashAttention computes exact attention with memory-efficient tiled evaluation. An uncompressed fixed-width KV cache uses $O(T)$ storage. With that cache, one new query performs $O(T)$ attention work; a fixed-width recurrent update instead has cost independent of elapsed history. These are asymptotic statements at fixed batch size, precision, and architecture, not wall-clock speed predictions.

Multi-head attention repeats the lookup with separate projections, then concatenates and linearly mixes the head outputs. Heads may learn different patterns, but assigning one head to each physical state factor requires evidence. The predictive representation is produced jointly by embeddings, attention, residual paths, and feed-forward layers; heads alone do not define a Bayesian belief. Inspecting their weights is one diagnostic, alongside state probes, interventions, ablations, and predictive tests. Attention is not Explanation gives examples where attention weights do not uniquely explain predictions.

This section explains how transformers are trained as predictive models of trajectories: the causal mask, the autoregressive factorization of the likelihood, and what a transformer trained this way computes when an agent uses it inside a decision loop.

An action-conditioned predictor must not read its future targets. During training, a trajectory with $T$ observed positions provides $T-1$ adjacent next-observation targets. A causal mask excludes attention from position $i$ to $j>i$, usually by assigning those scores negative infinity. Provided an allowed finite score remains in each row, the masked weights are zero. With aligned targets and no other leakage paths, all positions can be supervised in one layerwise-parallel forward pass. This reuses computation across prediction tasks without removing attention's quadratic pairwise arithmetic.

Formally, the model defines an autoregressive factorization of trajectory likelihood, echoing the probabilistic framing of Part II:

where:

• $o_{1:T}$ and $a_{1:T}$: the full sequences of observations and actions from timestep $1$ to $T$
• $p(o_t \mid o_{<t}, a_{<t})$: the probability mass or density of observation $o_t$ given all earlier observations and actions, the world-model factor
• $p(a_t \mid o_{\le t}, a_{<t})$: the probability mass or density of action $a_t$ given the observations seen up to and including time $t$ and all earlier actions, the policy factor

The chain rule uses the ordering $o_1,a_1,o_2,a_2,\ldots,o_T,a_T$. For $t=1$, the observation factor is the marginal $p(o_1)$. Training both observation and action factors would model environment observations and the data-collecting policy, respectively. However, the paired token $(o_t,a_t)$ used below already contains $a_t$: an action head reading that same token would leak its target despite a causal mask. Joint supervision needs separately ordered observation/action tokens or an action-excluded, appropriately shifted input path. The observation-only objective below does not train action probabilities. A behavior-policy estimate can support off-policy weighting only with reliable propensities and suitable support; learning such a head is not itself a validity guarantee.

A next-observation world model maximizes the expected log-likelihood of observations under the data distribution, estimated as the empirical average over $N$ training trajectories, with teacher forcing:

where:

• $\mathcal{L}(\theta)$: the empirical trajectory-average log-likelihood, maximized during training; minimizing $-\mathcal{L}$ is equivalent
• $\theta$: the model parameters
• $N$: the number of trajectories; superscript $(n)$ identifies observations and actions from trajectory $n$
• $T$: the common number of observed positions in this equation, giving $T-1$ adjacent prediction targets
• $p_\theta(o_{t+1}^{(n)}\mid o_{\le t}^{(n)},a_{\le t}^{(n)})$: the next-observation distribution conditioned on the recorded prefix and current action

This is a conditional likelihood, not a variational ELBO. The later simulator uses $T$ for its number of transitions and allocates $T+1$ observation slots, including a zero start marker; its target count therefore differs from this equation's convention.

For continuous observations, $p_\theta$ is a density, not the probability of one exact value. The code below minimizes mean squared error. For a scalar Gaussian observation model with the network output as its mean and a fixed positive variance, negative log-likelihood differs from squared error only by a positive constant factor and an additive constant. With the same targets and weighting, the objectives therefore have the same minimizers; this interpretation does not mean that the code learns or calibrates a predictive variance.

Teacher forcing conditions on recorded observations; a recursive open-loop forecast under fixed candidate actions conditions on its own later predictions. Closed-loop control is different: it uses newly observed environment feedback to choose or revise actions. Identical one-step errors need not produce identical rollout errors. To see why, consider the local error recursion $e_{t+1}\approx J_t e_t+\delta_t$, where $J_t$ propagates an existing error and $\delta_t$ is a new disturbance or model mismatch. In the idealized scalar unit-gain additive case with zero initial error, independent zero-mean disturbances with common finite variance $\sigma^2$ give root-mean-square error $\sigma\sqrt H$; a constant bias $b$ instead gives signed error $Hb$. Contractive propagation can bound either effect; unstable propagation can amplify it. Feeding predictions back can introduce dependent errors, but neither correlation nor monotonic growth is inevitable.

There is a second, quieter gap in the same objective: it is an average over the data-collecting distribution, so it weights trajectories exactly as the logging policy produced them. A planning agent, by contrast, will steer its rollouts toward states it finds interesting or valuable, which may be rare or absent in the training data. Predictive fidelity on the data distribution therefore bounds neither fidelity nor uncertainty quality under the agent's own exploration distribution, a fragility we return to in the limitations section.

A causal transformer produces a learned predictive representation, not automatically a Bayesian belief or a statistically sufficient state. Let $\mathcal H_t=(o_{\le t},a_{<t})$ denote available history and $b_t=p(s_t\mid\mathcal H_t)$ the Bayesian posterior in a specified POMDP. Under its transition and observation assumptions, this belief supports prediction for supplied future actions. A learned feature $r_t=g_\theta(\mathcal H_t,a_t)$ always determines its own prediction head, but statistical sufficiency would additionally require $p(o_{t+1}\mid\mathcal H_t,a_t)=p(o_{t+1}\mid r_t)$ for the true data-generating distribution. Finite context, lossy representations, limited capacity, or unsuccessful optimization can break that equality. Retained token representations provide addressable memory from which the network computes features; they are neither the full raw history nor the posterior by definition.

This view is what makes sequence-modeling world models work at all:

• Trajectory Transformer (Janner et al., 2021) discretizes trajectory components including observations, actions, and rewards and fits an autoregressive model. A separate beam-search procedure uses that model for trajectory optimization; action proposals and reward predictions are not the whole planning algorithm.
• Decision Transformer (Chen et al., 2021) produces actions from desired return-to-go and state/action history. It is a return-conditioned policy, not necessarily an environment-transition predictor or an online search planner.
• VideoGPT uses autoregressive prediction of learned video tokens, while Gato applies sequence modeling to multimodal policy outputs. Shared sequence machinery is useful, but modality-specific representations, data, and objectives remain important.

A planner knows each candidate open-loop action sequence it proposes; a feedback policy instead chooses later actions in response to later observations. An action-conditioned model can compute a forecast for either kind of query, but observational conditioning is not automatically a valid intervention model. Appropriate action support and causal assumptions are needed, for example, actions randomized independently of hidden disturbances, or a history sufficient to remove action-outcome confounding. If a logging policy uses hidden information unavailable to the model, its action may reveal that information rather than cause the observed outcome. Our simulator uses independent random controls and a known transition mechanism. System identification likewise may use deliberately designed excitation inputs, not only inputs selected by someone else.

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:

• $h(t)\in\mathbb R^n$: the continuous-time state
• $x(t)\in\mathbb R^m$: the input signal
• $y(t)\in\mathbb R^p$: the output
• $A\in\mathbb R^{n\times n}$: the state-transition generator
• $B\in\mathbb R^{n\times m}$: the input map into the state derivative
• $C\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:

where:

• $h_t$: the discrete-time latent state after processing token $t$, with the initial state $h_0$ set to zero
• $\bar{A} = e^{A\Delta}$: the discretized transition matrix, obtained by exact matrix exponentiation of $A$ over the step size $\Delta$
• $\bar{B}$: the discretized input matrix, derived from $A$, $B$, and $\Delta$ under the zero-order-hold assumption that $x$ is constant over the step
• $x_t$: the input token at timestep $t$
• $C$: the same output matrix as in continuous time

Here $x_t$ is held constant over $((t-1)\Delta,t\Delta]$ and $h_t$ is the state at the interval's end. Variation of constants gives the integral for $\bar B$: it accumulates the input's contribution while the continuous dynamics evolve. If $A$ is invertible, $\bar B=A^{-1}(e^{A\Delta}-I)B$; the integral also covers singular $A$. For a fixed stable mode $a$, retention per token is $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\Delta$. A larger step forgets more per token. Learning $\Delta$ therefore changes token-scale dynamics, not the physical-time constant of a fixed $A$ in isolation.

Because $\bar{A} = e^{A\Delta}$, a fixed positive step $\Delta$ maps each eigenvalue $a_i$ with negative real part to a discrete multiplier of magnitude $|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 $A$. 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]$ uniformly and has time-dependent $1/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 $h_0=0$, unrolling gives $h_t=\sum_{k=0}^{t-1}\bar A^k\bar Bx_{t-k}$ and $y_t=\sum_{k=0}^{t-1}K_kx_{t-k}$ with $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(T\log T)$ work for a full sequence, while cached recurrence has $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 $B_t$, $C_t$, and the discretization step $\Delta_t$ depend on the current input while retaining a learned input-independent continuous-time $A$. Its Appendix B.1 explicitly credits S5's scan and distinguishes Mamba's selection mechanism and hardware-aware implementation. For real $a_i<0$ and $\Delta_{t,i}>0$, the effective decay $\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.

Linear transformers make the connection explicit. Let feature-mapped keys and queries be columns in $\mathbb R^{d_k}$ and values be columns in $\mathbb R^{d_v}$. With $S_0=0$, the numerator accumulator is $S_t=S_{t-1}+k_tv_t^\top\in\mathbb R^{d_k\times d_v}$, and its column readout is $S_t^\top q_t=\sum_{j\le t}(q_t^\top k_j)v_j$. Normalized linear attention also keeps $c_t=c_{t-1}+k_t$ and divides by $q_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.

This section surveys architectures that combine attention and state-space recurrence, the design patterns they follow, and why mixing the two mechanisms is a hypothesis worth testing for long-horizon world models.

Can recurrence reduce memory cost while attention preserves useful access to selected context representations? Several 2024 hybrid language-model reports explore this design hypothesis. Static layer interleaving does not dynamically identify valuable events: it applies its assigned layers to every token. Savings come from using fewer attention layers, restricting their span, or sharing parameters. Their value for a world-model task still requires task-specific evaluation.

One pattern is layer interleaving. Jamba (AI21, 2024) uses an attention-to-Mamba ratio of 1:7 in its reported architecture. Zamba (Zyphra, 2024) repeatedly invokes a shared attention-plus-MLP module among Mamba blocks, saving parameters relative to separate copies. Griffin and Hawk (De et al., 2024) use gated linear recurrence: Griffin retains local attention, whereas Hawk is recurrent-only. Their reported quality and inference-cost comparisons concern specified language benchmarks, model scales, contexts, and hardware, not all world models. Griffin's held-out-book experiments also test length extrapolation beyond training context. A length-independent recurrence is a plausible contributing bias, but that alone neither proves the mechanism nor guarantees extrapolation.

For world models specifically, hybridization predates the recent wave of SSM-based hybrids and takes a second form: transformer dynamics inside a stochastic latent world model.

• TransDreamer (Chen et al., 2022) introduces the Transformer State-Space Model (TSSM). It replaces recurrent dynamics with transformer dynamics and modifies posterior conditioning to permit parallel training. Its 2D and 3D Hidden Order Discovery tasks involve discovering and remembering object orders through interaction. Reported gains over Dreamer are empirical, not a proof that finite recurrent state cannot solve these tasks or that memory is the only changing factor.
• TWM (Robine et al., 2023), not TSSM, uses Transformer-XL dynamics over latent states, actions, and rewards with a separate observation model. It reports results under Atari 100k's 100,000-interaction budget. Short attention paths are a possible benefit, but the reported comparison does not isolate them as the cause of sample efficiency.
• A possible chunked hybrid runs local recurrence within windows and attention over their summaries. This should not be confused with the named Block-Recurrent Transformer, whose state update uses attention and gates. For one full summary-attention pass with $M\approx T/W$ fixed-width summaries and linear within-window work, total work is $O(T+M^2)=O(T+(T/W)^2)$. Fixed $W$ leaves a quadratic term; bounded $M$, or growth no faster than order $\sqrt T$, gives linear-order total work. Recomputing summary attention at every token would change this accounting.

A third design axis changes attention's retained memory. Sliding windows restrict direct access to recent tokens; cache compression can select, merge, or quantize stored representations. Global tokens provide shared access points, but they become a recurrent cross-segment summary only if an explicit update-and-carry mechanism is implemented with the required causal restrictions. Chunked prefill is different: it schedules processing in chunks and need not reduce retained KV state or change the memory representation. These choices trade cost against available information, with consequences revisited in Part VI.

For a new task, ask what history distinctions predictions or decisions require. One useful test places a decision-critical observation at varying temporal distances while controlling its content, distractors, and task difficulty. Plot error against distance alongside aggregate prediction loss. A fixed-decay system may lose usable event information as distance grows; an attention model may fail when the relevant representation leaves its context. These are hypotheses, not guaranteed curve shapes: learned retrieval can fail inside a window, and selective recurrence can retain events through distractors. A distance-controlled experiment reveals failures that an average can hide, but our smooth-dynamics experiment below does not perform this rare-event test.

Before the full experiment, let us build direct intuition for attention with a tiny hand-computed lookup. Suppose a trajectory contains three tokens, each describing an object with a one-hot identity feature and an associated position value. A query token asks, in effect, "where is the red object?" In a real transformer the query, keys, and values are learned projections, but the hand-designed example uses the same weighted-lookup operation. Training optimizes the projections for an objective; useful semantic alignment must be learned rather than assumed.

With the one-hot scores divided by $\sqrt{3}$, the matching token receives weight about 0.47 and each other token about 0.26. The output is 6.62, not the red key's stored position 9.0. Increasing the matching score's margin would concentrate the weights toward that value. This example illustrates standard scaled dot-product attention and maximum-subtracted softmax; these details are not universal to every attention variant. A retained token can be accessed at a long lag if useful keys, queries, and values are learned. When it is removed, direct access is lost, although another retained representation may preserve its content. In a fixed stable SSM, an input's contribution propagates through powers of the transition matrix. Its amplitude can decay while usable recall either deteriorates gradually or fails abruptly through finite precision, interference, or decoding thresholds. Selective retention can change that behavior; it is not implemented by the fixed-coefficient SSM below.

Now we build a synthetic, partially observable linear-Gaussian system: a two-dimensional hidden state, scalar control, and noisy scalar observation. A single scalar reading does not identify the state, so observation and control history can improve prediction. A correctly specified Kalman filter would use dynamics- and covariance-dependent prediction/correction updates, generally with time-varying gains before steady state. Because the dynamics are coupled and oscillatory and controls are consequential, a plain running average is not an established approximation to that filter. Environment structure, model capacity, and optimization all influence the comparison below.

We train two models on the same data with the same next-observation objective: a tiny causal transformer with one self-attention block, and a tiny diagonal SSM with a stable, input-independent learned recurrence (a time-invariant member of the family, since its decay, input matrix, and output matrix do not depend on the token, unlike the selective SSMs described above). Then we compare one-step accuracy against recursive rollout error under held-out action sequences, distinguishing open-loop prediction quality from the compounding behavior an agent would experience. Keeping the models tiny is a deliberate methodological choice: at roughly 3,300 parameters for the transformer and 1,200 for the SSM, we can inspect the mechanisms directly. Their capacities still differ, so the comparison does not isolate architecture from scale.

First, the system and dataset. We simulate trajectories of

where:

• $s_t \in \mathbb{R}^2$: the latent state, hidden from the learner
• $u_t$: the scalar control input chosen by the experimenter
• $y_t$: the noisy scalar observation
• $A, B, C$: the transition, input, and observation matrices, matching the system used in the simulation code below
• $w_t$: process noise added to the state transition
• $v_{t+1}$: observation noise added to the sensor reading at time $t+1$

The observation matrix $C$ projects the two-dimensional state onto a single scalar with a fixed linear combination of the two components, so the observation is a lossy, noise-corrupted shadow of the state. The eigenvalues of $A$ have magnitude slightly below one, so the latent dynamics are stable and gently contract, and the off-diagonal entries couple the two channels so that neither component evolves independently. Tokens are $(y_t, u_t)$ pairs and the target at position $t$ is $y_{t+1}$, so every model is trained with teacher forcing under the causal objective from the previous section. Position 0 is a start-of-sequence exception: the simulator never emits $y_0$, so the first token carries a placeholder zero observation paired with $u_0$: it is $(0, u_0)$, not a real observation pair, and genuine observation-control pairs begin at position 1.

These rollout curves measure prediction error after genuine observations stop arriving. Growth can reflect both model error and irreducible multi-step uncertainty from process noise. The dashed lines summarize teacher-forced error over all test positions; they are not horizon-matched noise floors, so a crossing is a descriptive comparison rather than a theorem about compounding. Both models condition on their own predicted observations after the warm start. The transformer can still attend to the supplied start token and first observation, but it cannot retrieve later observations that were never supplied. Inspect the printed curves for this run rather than assuming monotonic growth or a stable ordering between architectures.

Separating the quality axes prevents overreading this experiment. Teacher-forced test MSE measures one-step predictive fidelity, and the recursive curves measure a different conditioning regime. Neither deterministic point predictor supplies a calibrated uncertainty distribution, so we have not evaluated uncertainty quality. No controller acts on the forecasts, so we have not measured decision usefulness either. Good performance on either plotted error metric does not settle those remaining questions.

The attention-by-lag plot can suggest hypotheses about retrieval or integration, but it cannot uniquely locate a failure. Query formation, key/value representations, residual paths, and dynamics layers can produce similar visible weight patterns. Test a suspected mechanism by intervening on relevant tokens, probing state information, or ablating components before selecting a remedy. Here the model explicitly returns weights, making inspection convenient; in other implementations requesting them can add memory or prevent use of an optimized attention path.

Attention and structured recurrence offer different ways to represent history in learned predictors. They did not invent long-horizon partially observable control: classical filtering and belief-space planning, as well as pre-transformer external-memory networks such as Neural Turing Machines, already addressed aspects of it. Modern transformer deployments demonstrate large context capacities (for example, Llama 3.1's reported 128K context), but capacity is not a guarantee of useful world-model memory. Structured SSMs connect learned sequence layers with linear-system analysis; hybrids combine mechanisms without removing their task and resource dependencies.

Several limitations matter when using these architectures.

• Recursive prediction. Our two models train with recorded observations and are evaluated with self-fed predictions. No architecture alone guarantees accurate forecasts under model mismatch, but exact or contractive dynamics need not exhibit growing model error. Attention can read retained real context during imagination while newer tokens are synthetic; whether that helps depends on learned representations and queries. Scheduled sampling and suitable rollout-aware objectives can reduce some conditioning mismatch. Their statistical interpretation must be assessed separately: not every multistep objective sacrifices a valid likelihood.
• Retained information. Attention cannot directly read absent KV entries, although summaries or other representations may preserve related information. An append-only cache grows with tokens; bounded windows need not grow without bound. Fixed-width finite-precision recurrence may merge histories requiring different predictions or actions. A sufficient task state avoids that problem, whereas more training data alone cannot reverse a distinction already absent from the representation.
• Hybrid configuration. Layer ratios, state widths, sharing, and cache policies add design choices. Complexity and stability analysis constrain some choices but do not identify the best task-specific configuration. Kernel and debugging support must be checked for the chosen implementation rather than assumed uniformly mature or immature.

Architecture choice must be evaluated against intended use. Memory probes and closed-loop environment tests complement one-step and recursive forecasts. Point MSE, as used here, scores a point prediction. Negative log-likelihood instead scores a predictive distribution: it is sensitive to assigned probabilities, though one aggregate value does not separately diagnose calibration, representation sufficiency, and decision utility. A planner can exploit poorly modeled regions outside the logged distribution. Likelihood directly supervises observed contexts and targets; extrapolation elsewhere depends on learned parameters and inductive biases, without a general off-support calibration guarantee. These issues motivate Part VI and Part XII.

The chapter connects attention-based predictive memory, structured recurrence, and selected hybrids.

• Attention forms query-dependent weighted readouts from retained token representations. Neither the memory nor the readout is necessarily lossless or statistically sufficient. At fixed dimensions, dense sequence attention has quadratic pairwise arithmetic, uncompressed KV storage is linear in context, and cached per-query attention work is linear in context.
• Causal masking supports parallel next-observation supervision only with aligned, leakage-free inputs and targets. Joint action supervision needs an input path that does not already contain the action being predicted. Observational action conditioning requires additional causal and support assumptions before it is used for intervention forecasts.
• Fixed-coefficient linear SSM layers admit recurrent, scan, or convolutional evaluation. Input-selective affine recurrences generally do not have one fixed convolution kernel. Stability claims require the appropriate contraction, forcing, and numerical conditions. HiPPO's optimality concerns its specified weighted projection dynamics, not arbitrary learned fixed recurrences.
• Hybrids combine attention and recurrence or place transformer dynamics inside a stochastic latent world model. Integration versus event retrieval is a useful testing hypothesis, not an exclusive allocation of capabilities.

In the experiment, both models learn finite one-step errors, and recursive errors vary with horizon without a fixed ordering or monotonic-growth guarantee. Unequal capacities and optimization remain confounders. The matched-query attention plot compares learned weights with uniform attention at equal context length; it does not prove a memory mechanism.

The methodological lesson is to compare predictions under the conditioning and actions where they will be used. One-step loss, recursive error, and attention profiles provide complementary diagnostics, not a complete assessment of state representation, uncertainty, or decision usefulness. Those require additional probes, probabilistic evaluation, and controllers tested in the environment. Our short smooth-dynamics task also leaves the opening's distant-event memory problem untested. The next chapter develops autoregressive token world models, including video and action representations, with these distinctions still in force.

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about transformers, neural state-space architectures, and their role in world models.