Grid-World Representations in Transformers Reflect Predictive Geometry

TL;DR

This paper demonstrates that transformer models trained on simple stochastic processes develop internal representations aligned with the underlying predictive geometry of the data, revealing how neural networks internalize structural constraints.

Contribution

It provides a controlled analysis showing that transformer representations reflect the predictive geometry of the world, linking internal states to optimal prediction in a minimal setting.

Findings

Transformers' hidden states align with analytically derived predictive vectors.

Representations are often low-dimensional and correspond to the world's geometry.

The study offers a framework for understanding neural internalization of structural constraints.

Abstract

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.