Search papers, labs, and topics across Lattice.
The authors trained decoder-only transformers to predict constrained random walks on a 2D lattice, where optimal prediction depends on the walker's position relative to the target and remaining time. They found that the hidden activations of the trained transformers strongly align with the analytically derived, ground-truth predictive vectors representing the world's geometry. This demonstrates a direct link between world-model-like representations and the predictive geometry of the data.
Transformers trained on a simple grid-world learn hidden representations that directly reflect the underlying predictive geometry, offering a glimpse into how neural networks internalize structural constraints.
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 geometry. We train decoder-only transformers on prefixes sampled from the exact distribution of these walks and compare their hidden activations to the analytically derived sufficient vectors. Across models and layers, the learned representations align strongly with the ground-truth predictive vectors and are often low-dimensional. This provides a concrete example in which world-model-like representations can be directly traced back to the predictive geometry of the data itself. Although demonstrated in a simplified toy system, the analysis suggests that geometric representations supporting optimal prediction may provide a useful lens for studying how neural networks internalize grammatical and other structural constraints.
