Transformers and Neural State-Space Architectures › Code Implementation

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.

Apart from the zero start marker, each token carries a scalar observation and its aligned action; the target is the next observation. Optimal conditional prediction can use history without recovering the hidden state exactly, and fresh process and sensor noise prevent perfect forecasts. Training and test trajectories use the same simulator and independent control/noise laws. Recursive evaluation changes observation feedback after its prefix; its start position and prefix-length distribution also differ from the pooled all-position teacher-forced score. We therefore compare conditioning regimes without claiming to isolate every causal source of their difference.

Next, the two architectures. The transformer uses a single-head causal attention block, learned positional embeddings, a feed-forward sublayer, and a linear head. The SSM uses a real diagonal recurrence with input-independent $B$ and $C$. We clamp each sigmoid decay to $[\epsilon,1-\epsilon]$, with $\epsilon=10^{-4}$, so sigmoid saturation cannot produce a unit multiplier in this float32 implementation. With bounded finite writes, the uniformly contractive recurrence has a bounded state in exact arithmetic. Finite-precision overflow is still possible at extreme scales; the bound is not a guarantee for other dtypes, unbounded inputs, nonfinite parameters, or the entire training process. A learned skip connection passes the current embedding to the output without routing it through recurrent memory.

A few implementation details connect directly to the chapter's formalism. In the transformer, the mask is built with torch.triu above the diagonal, so position $i$ is forbidden from attending to positions $j > i$, the causal constraint rendered as a boolean matrix. The forward method returns the attention weights as well as the predictions, which we will use later for the diagnostic plot; average_attn_weights=True averages over heads into a single weight matrix; the model as defined uses a single head (nhead=1), so this averaging is a no-op and the returned matrix is that head's attention directly. In the SSM, the loop over time steps is the explicit recurrence $h_t = \bar{A} h_{t-1} + \bar{B} x_t$, written out step by step for clarity rather than hidden inside a parallel scan; on this short sequence the loop is harmless, but on longer sequences, parallel scan or convolution can avoid this Python-level sequential loop. For fixed state size, a work-efficient scan has linear total work and logarithmic parallel depth, while FFT convolution takes $O(T\log T)$ work. Both models are small (roughly 3,300 and 1,200 parameters respectively), but their unequal capacities remain a confound. Both use full-batch Adam and the same learning rate and epoch count; matching that recipe does not guarantee equally effective optimization.

Both models learn a useful one-step predictor in this run; the printed test errors are measured rather than assumed. The observation-noise variance is $0.1^2=0.01$, but that is not the full achievable one-step error: process noise and uncertainty about the hidden state also matter. We do not implement a Kalman or fitted linear baseline here, so these two scores do not establish how close either model is to optimal. The environment has coupled, oscillatory dynamics, while our SSM has only real positive decay channels; it is not the exact model class for this system. To examine recursive prediction, we warm-start each model with the start token $(0,u_0)$ and the first genuine observation token $(y_1,u_1)$, preserving their training positions. The first forecast targets $y_2$. Each subsequent token pairs the predicted observation at time $t$ with the known control $u_t$, and the prediction is scored against the raw held-out observation at time $t+1$. After the warm start, no new real observation enters the prefix.

Two details matter for interpreting this procedure. The warm start consists of the zero start marker and the first genuine observation, each paired with its aligned control; later observation components are predictions. Known controls remain aligned with the timestep of each token. Thus the experiment changes observation feedback after the prefix without shifting the initial token positions or pairing observations with the wrong actions.

We also inspect how the trained transformer distributes attention across temporal lags. To avoid mixing query positions with different available context lengths, every example contributes only its final query, which can access all 32 tokens. We compare the mean weights with uniform attention over that same context. This is a descriptive memory diagnostic, not proof that the model performs Bayesian filtering or uses the most informative observations.

At the final query $i=L-1$, lag $k$ selects key $j=L-1-k$. Reversing that query's key axis places self-attention at lag zero, and averaging across held-out sequences keeps query length fixed. The complete $L$-lag profile sums to one up to numerical error. The plot displays only its first eight entries, not a renormalized eight-token distribution. Its dashed reference is $1/L$, uniform attention over all available tokens.

Pooling all query positions would introduce a baseline effect: even uniform causal attention has weights $1/(i+1)$ at query $i$, so averaging its eligible pairs at each lag produces a declining curve. The matched-query construction avoids that particular confound. We can now ask whether this final-query profile favors recent tokens relative to uniform attention, while remembering that learned values, residual paths, and other components also determine predictions.

The environment's coupled oscillatory transition is not exactly representable by this real positive-decay SSM. The model can still fit useful finite-horizon predictions, but an observed advantage would not prove an exact model-class match. Nor does this experiment isolate architecture alone: the models have different parameter counts and parameterizations, and identical optimizer settings need not suit them equally. It is one controlled data-generating task and one training seed, not evidence of a consistent ranking across environments.

The two rollout curves below plot recursive open-loop MSE against horizon, with each model’s all-position teacher-forced MSE as a dashed reference.

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.