Predictive Statistics Shape Emergent World Representations of Grid Walkers
Authors: Sasha Brenner, Thomas R. Knösche, Nico Scherf
Organizations: Max Planck Institute for Human Cognitive and Brain Sciences, Leipzig, Germany · Leipzig University, Germany · ScaDS.AI, Dresden/Leipzig, Germany
Next-token predictors often appear to develop internal representations of the latent world and its rules. The probabilistic nature of these models suggests a deep connection between the structure of the world and the geometry of probability distributions. In order to understand this link more precisely, we use a minimal stochastic process as a controlled setting: constrained random walks on a two-dimensional lattice that must reach a fixed endpoint after a predetermined number of steps. Optimal prediction of this process solely depends on a sufficient vector determined by the walker's position relative to the target and the remaining time horizon; in other words, the probability distributions are parametrized by the world's grid geometry. We train decoder-only transformers and recurrent networks on prefixes sampled from the exact distribution of these walks and compare their hidden activations to sufficient statistics of prediction, by measuring alignment and linear readability across layers. We find that the transformer's computation factors into two stages: the first attention block extracts the sufficient statistic from the input, and later layers transform it into the next-step predictive geometry. Across constraint variants the post-attention representation is universal: a shared world-state of the lattice that can be read directly as a world model, traced to the predictive geometry of the data. Later layers then specialize it to each variant's next-step distribution. Recurrent networks reach the same Bayes-optimal loss but do not isolate this world-state as a separate stage, showing that the world-model geometry also depends on architecture. Although demonstrated in a toy system, the results suggest that the geometry of the predictive distribution is a useful lens on how neural networks internalize the structure of their data.
Figures & tables
Figure 1: The three walker variants and the prediction task. Each column shows one walker (left to right: walker L (looper), walker D (double-endpoint), walker O (obstacled)) and reads top to bottom: a length-7 prefix on the lattice, one color per move type, final position as a black dot; the emitted token sequence, with a “?” for the unknown next move; and the next-step distribution Pr(c8∣s1:7) produced by the per-walker model fθL,fθD,fθO . Endpoints are red stars, obstacles dark squares, the start a white circle.
Figure 2: Activation geometry across targets and layers. Left: Last-LayerNorm activations (right column of each pair) closely resemble the ground-truth next-step targets (left column) for walker L (top) and walker D (bottom). For both walkers the target is the effective field qt of Equation 2 ; for walker L qt is effectively 2D since quv=0 , whereas for walker D it has all three components nonzero. Activations are Procrustes-aligned to the ground-truth target coordinates. Right: 3D PCA views of all valid walk-prefix activations across post-attention residual, post-MLP residual, and output logits for walkers L, D, and O-L2. Each point is a distinct path on the lattice, colored by path length with a shared colorbar.
Figure 3: Linear probe alignment and cross-variant similarity ( dhead=64 ). (A,B) All-walk in-sample affine probes on the residual stream, fit in opposite directions at each hook; gaps mark second-layer hooks, absent in one-layer models. (A) Decoding direction ( activations→q ): joint R2 of an affine probe predicting the 3D effective field from the activation. (B) Reverse direction ( [xt/t,yt/t,logt]→activations ): joint R2 of an affine probe predicting the activation from the 3D geometric features. High R2 here means the activation lies in the affine span of the features. (C) Held-out state-class R2 of a linear probe onto the 3D geometric features, comparing post-attention (hatched) and post-MLP (solid) bars across walker variants and architectures. Diamonds show the maximum off-diagonal feature-axis leakage. (D,E) Pairwise linear CKA between state-mean activations at the post-attention (D) and post-MLP (E) hooks, computed on the common support intersection of states valid across the selected runs. Each entry averages the pairwise CKA over all distinct run pairs across the two groups; diagonal entries average over distinct same-group run pairs and thus measure within-variant reproducibility. Before the MLP, representations are nearly identical across variants (CKA ≥0.72 ); after the MLP, cross-variant similarity drops sharply, with the structurally most distinct Walker O-1L diverging the most.
Figure 4: First-layer mechanism in a representative trained Walker L transformer ( L=1,H=1,dhead=128 ). (A) The single attention head is nearly uniform over the causal prefix. (B) The token component of the value-write vectors forms a two-dimensional compass code for the four lattice moves. (C) The residual after removing the best affine xt/t,yt/t component from the empirical attention output (orange dashed) is almost exactly the time average of the position component of the value-write vectors (black). Panel C is plotted in the PCA basis of the seven empirical residual centroids.
Figure 5: Layerwise transformation of synthetic feature grids. We construct grids spanning the actual datapoints at the post-attention residual stream of Walker L (top row) and Walker D (bottom row), with points colored by the fitted logt coordinate (low to high, see colorbar). Each subsequent column shows the image of these grids at LayerNorm 2, post-MLP, the final LayerNorm, and the output logits. The annotated PCA EV 3 is the fraction of variance captured by the top three principal components of the activations at each stage.
Appendix figures & tables13 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Held-out state-class R2 of a linear probe onto the 3D geometric features [xt/t,yt/t,logt] , by selected walker/architecture row and head dimension (columns). Left: post-attention residual stream. Right: post-MLP residual stream. Table 1 reports the complete architecture sweep.
dhead=32
dhead=64
dhead=128
walker
L
H
attn
MLP
attn
MLP
attn
MLP
L
1
1
0.962
0.850
0.982
0.793
0.967
0.775
L
1
2
0.985
0.733
0.969
0.755
0.937
0.707
L
2
1
0.955
0.856
0.988
0.825
0.952
0.811
L
2
2
0.988
0.791
0.951
0.785
0.934
0.785
D
1
1
0.969
0.785
0.988
0.695
0.981
0.750
Appendix
Table 1: Complete architecture-resolved held-out state-class R2 for linear recovery of the geometric features [xt/t,yt/t,logt] across the 36 trained transformer models. “attn” denotes the post-attention residual stream and “MLP” denotes the post-MLP residual stream. The consistently high post-attention values indicate that the attention block makes this sufficient statistic linearly available across architectures. The lower post-MLP values show that the MLP has already applied a variant-specific nonlinear transformation, so the activations are no longer as faithfully linearly obtainable from the coordinates. Figure 6 gives a compact visual summary of this diagnostic, while this table makes the layer count L and number of attention heads H explicit for every architecture in the sweep.
Figure 7: All-walk in-sample R2 for affine probes between residual-stream activations and target features, across dhead∈{32,64,128} . At each width we show the four main-text selected runs (Walker L L1H1, Walker D L1H1, Walker O L1H2, and Walker O L2H2), not every network in the architecture sweep. Decode rows fit the geometric target zt=[xt/t,yt/t,logt] or effective field qt from raw per-position activations; encode rows fit the state-mean activation from the corresponding state-mean target. Gaps mark second-layer hooks, absent in one-layer runs.
Figure 8: Full pairwise linear CKA matrix between state-mean activations of every individual run in our sweep, at the post-attention residual stream (left) and the post-MLP residual stream (right).
Figure 9: Variance outside the best two-dimensional PCA subspace for the top representation (left: transformer final LayerNorm activations or RNN top hidden states) and raw output logits (right). Each point is one final checkpoint. Color and marker shape denote architecture family; vertical bars show the within-family min–max range for each walker, and horizontal ticks show the median.
Figure 10: RNN capacity scan: predictive loss. Final validation loss is plotted as excess over the Bayes-optimal value, as a function of hidden size, recurrent cell type, depth, and walker variant. Loss is a model-level quantity, so two-layer models appear once rather than separately for each recurrent layer. All RNN families approach the optimum as width increases; Walker O remains the hardest variant at small widths.
Figure 11: RNN capacity scan: linear accessibility of geometric and predictive targets. Each panel shows held-out CV R2 as a function of hidden size for LSTM and GRU recurrent networks trained on the three walker variants. The target zt=[xt/t,yt/t,logt] is the explicit geometric coordinate system used in the main text, while qt is the effective field. Decode rows fit the target from the raw per-position activation A via OLS; encode rows fit the state-mean activation from the target via OLS. For two-layer RNNs, h0 and h1 denote the first and second recurrent hidden layers; for one-layer RNNs, h0 is the only hidden layer. The logits series probes the output-layer logits as a sanity check: qt is a fixed linear transformation of the centered logits and saturates to R2=1 at this hook for Walkers L and D.
Figure 12: Hidden-state geometry across architectures, depths, and walkers at H=128 . (A) LSTM. (B) GRU. Within each panel, rows are L=1 top hidden ( ≡ layer 0 ), L=2 layer 0 , and L=2 layer 1 (top); columns are Walker L, Walker D, and Walker O. Each cell shows the 3D PCA of the state-mean recurrent activation, coloured by context position t , with axes drawn on equal scale in data units so a near-flat manifold renders as flat. The geometry is set primarily by the task: Walker L collapses to a near-2D dome, Walker D forms a 3D bowl, and Walker O is capacity-sensitive. Neither the LSTM/GRU choice nor the depth substantially alters the manifold’s intrinsic shape.
Figure 13: All-vs-all state-mean lCKA, transformer residual streams vs. RNN hidden layers, per walker. Rows and columns of each heatmap are: every transformer residual-stream hook of the canonical dhead=64 run (the first few entries, separated by the white guide lines), followed by every per-layer hidden state of every RNN in the sweep, grouped by (cell type, depth, hidden size, layer index). Three blocks are visible per walker: the transformer’s own residual stream is internally dissimilar; the RNN family is internally very similar; the off-diagonal block between the two architectures is moderate ( ≈0.5 – 0.7 for Walkers L/D, lower for Walker O). LayerNorm hooks are excluded.
L
H
dhead
walker L
walker D
walker O
1
1
32
1.000011
1.000006
1.000138
1
1
64
1.000010
1.000004
1.000044
1
1
128
1.000012
1.000005
1.000033
1
2
32
1.000010
1.000004
1.000042
1
2
64
1.000012
1.000005
1.000035
1
2
128
1.000011
1.000005
1.000038
Appendix
Table 2: Final validation loss for every trained model, reported as the ratio CEmodel/H(ct+1∣st) averaged over context positions. The denominator is the exact per-position conditional entropy of the ground-truth process, so a value of 1.000 corresponds to a Bayes-optimal predictor. Each column is one of the three walker variants (Section 3.2 ); each row is one architecture from the appendix sweep ( L∈{1,2} layers, H∈{1,2} heads, dhead∈{32,64,128} ).
Figure 14: Validation loss vs. tokens seen for the full sweep. Rows are head dimensions ( dhead∈{32,64,128} ); columns are layer/head configurations. Each panel shows three curves, one per walker variant. The vertical axis is the same dimensionless ratio as Table 2 : CEmodel/H(ct+1∣st) , averaged over context positions; the dashed line marks the Bayes-optimal value 1.0 . The horizontal axis is the cumulative number of training tokens (log scale).
cell
Lrnn
Hhid
walker L
walker D
walker O
LSTM
1
16
1.000008
1.000021
1.002547
LSTM
1
32
1.000005
1.000007
1.000892
LSTM
1
64
1.000004
1.000006
1.000376
LSTM
1
128
1.000005
1.000006
1.000190
LSTM
1
256
1.000006
1.000006
1.000130
LSTM
2
16
1.000006
1.000010
1.000625
Appendix
Table 3: Final validation loss for every trained RNN, reported as the normalized ratio CEmodel/H(ct+1∣st) . Rows enumerate recurrent cell type, recurrent depth, and hidden size; columns are the three walker variants.
Figure 15: RNN validation loss vs. tokens seen for the full recurrent sweep. Rows are walker variants; columns are recurrent cell and depth. Each panel overlays hidden sizes H∈{16,32,64,128,256} . The vertical axis is the same normalized validation-loss ratio used in Tables 2 and 3 ; the dashed line marks the Bayes-optimal value 1.0 .
Prior-data fitted networks (PFNs) have recently emerged as a powerful approach for Bayesian prediction tasks, approximating the posterior predictive distribution (PPD) through in-context learning. Despite their strong empirical performance and ability to go beyond point predictions, theoretical understandings of the algorithmic capability of transformers to learn distributions in context are still lacking. Focusing on Gaussian process regression problems, we show by construction that transformers can implement a gradient descent algorithm targeting the posterior predictive mean and variance, followed by nonlinear mappings that yield binned probabilities of PPD. We study the error bounds of the approximated PPD in terms of attention depth and bin resolution. Based on these results, we further demonstrate the key role of normalization and the choice of attention depth in enabling the extrapolation abilities of transformers beyond the pretraining sample size range. We conduct simulations that corroborate our findings, providing insight into the expressivity of PFNs targeting PPDs and how architectural choices may influence generalization capabilities.
Gyeonghun Kang, Changwoo J. Lee, Xiang Cheng
Department of Statistical Science, Duke University, Durham, NC, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA
World models learn to predict future states of an environment, enabling planning and mental simulation. Current approaches default to Transformer-based predictors operating in learned latent spaces. This comes at a cost: O(N^2) computation and no explicit spatial inductive bias. This paper asks a foundational question: is self-attention necessary for predictive world modeling, or can alternative computational substrates achieve comparable or superior results? I introduce FluidWorld, a proof-of-concept world model whose predictive dynamics are governed by partial differential equations (PDEs) of reaction-diffusion type. Instead of using a separate neural network predictor, the PDE integration itself produces the future state prediction. In a strictly parameter-matched three-way ablation on unconditional UCF-101 video prediction (64x64, ~800K parameters, identical encoder, decoder, losses, and data), FluidWorld is compared against both a Transformer baseline (self-attention) and a ConvLSTM baseline (convolutional recurrence). While all three models converge to comparable single-step prediction loss, FluidWorld achieves 2x lower reconstruction error, produces representations with 10-15% higher spatial structure preservation and 18-25% more effective dimensionality, and critically maintains coherent multi-step rollouts where both baselines degrade rapidly. All experiments were conducted on a single consumer-grade PC (Intel Core i5, NVIDIA RTX 4070 Ti), without any large-scale compute. These results establish that PDE-based dynamics, which natively provide O(N) spatial complexity, adaptive computation, and global spatial coherence through diffusion, are a viable and parameter-efficient alternative to both attention and convolutional recurrence for world modeling.
We present a framework for building and interpreting minimal transformer models. By constraining a transformer's embedding dimension and head size to 2, we enable full two-dimensional visualization of its internal representations. Embeddings, query/key/value transforms, attention outputs, residual streams, and decision boundaries can all be seen directly. Our central claim is that the learned geometry implies an algorithm; the arrangement of points and boundaries in R^2 can be read as a step-by-step procedure. We train a transformer on a simple task where it must produce the most recently observed even number whenever the '+' operator appears in a sequence of digits. Once trained, we visually walk through every step of the transformer's computation. We show how the model embeds the tokens and their respective positions in the sequence, transforms them via the Q, K, and V matrices, uses the dot product between the Q and K representations to form the attention matrix, and uses the attention matrix to select values that move the representation of each input token to the region of the domain of the output layer that will correctly predict the next token. We introduce a suite of interpretability visualizations that make the algorithmic interpretation of this procedure explicit. Our framework offers a pedagogical and experimental testbed to explore how transformers use informational geometry to implement next-token prediction.
Raneem Mahajne, Toviah Moldwin
Edmond and Lily Safra Center for Brain Sciences, The Hebrew University of Jerusalem