Grid-World Representations in Transformers Reflect Predictive Geometry

• URL: http://arxiv.org/abs/2603.16689v1
• Date: Tue, 17 Mar 2026 15:47:33 GMT
• Title: Grid-World Representations in Transformers Reflect Predictive Geometry
• Abstract summary: Next-token predictors often appear to develop internal representations of the latent world and its rules.<n>We train decoder-only transformers on vectors sampled from the exact distribution of random walks.<n>Across models and layers, the learned representations align strongly with the ground-truth predictive vectors and are often low-dimensional.<n>Although demonstrated in a simplified toy system, the analysis suggests that prefixes representations supporting optimal prediction may provide a useful lens for studying how neural networks internalize grammatical and other structural constraints.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Orthogonal JEPA: Factorized Predictive States for Latent World Models [95.20158869896393]
  We introduce method, a latent world-modeling framework based on predictive factorization.<n>We show that method can allocate redundant capacity to dominant signals while providing weak or conflicting gradients to less dominant predictive structure.<n>Experiments on controlled vision, single-cell transcriptomics, longitudinal health records, continuous control, and molecular dynamics evaluate representation quality, forecasting, planning, and long-horizon stability.
  arXiv  Detail & Related papers  (2026-08-20T13:59:57Z)
• Predicting When Random Low-Dimensional Reparameterizations Train Neural Networks [3.980892205755342]
  We present an orientation-resolved quadratic master formula that predicts the random-slice residual from both the curvature spectrum and the reference-to-solution displacement profile.<n>We then introduce Random Mapping Networks (RaMaN), which instantiate the predicted latent dimension using structured Hadamard or seed-regenerated Gaussian maps.<n>RaMaN closely tracks measured transition locations and outperforms orientation-agnostic approximations when displacement direction matters.
  arXiv  Detail & Related papers  (2026-08-12T21:09:54Z)
• Geometry-Aware Dataset Condensation for Diffusion Model Training [103.45641113998839]
  We propose to reformulate real subset selection as a geometry-aware distribution alignment problem.<n>By incorporating one-sided partial optimal transport, our method selectively aligns a compact subset with the full data distribution.<n>Experiments across diffusion variants, subset sizes, image resolutions, and training rounds show that our method achieves superior fidelity and distributional coverage.
  arXiv  Detail & Related papers  (2026-06-04T08:53:58Z)
• Generalized Linear Mode Connectivity for Transformers [87.32299363530996]
  A striking phenomenon is linear mode connectivity (LMC), where independently trained models can be connected by low- or zero-loss paths.<n>Prior work has predominantly focused on neuron re-ordering through permutations, but such approaches are limited in scope.<n>We introduce a unified framework that captures four symmetry classes: permutations, semi-permutations, transformations, and general invertible maps.<n>This generalization enables, for the first time, the discovery of low- and zero-barrier linear paths between independently trained Vision Transformers and GPT-2 models.
  arXiv  Detail & Related papers  (2025-06-28T01:46:36Z)
• Learning low-dimensional representations of ensemble forecast fields using autoencoder-based methods [0.08192907805418585]
  Large-scale numerical simulations often produce high-dimensional gridded data that is challenging to process for downstream applications.<n>We present two alternative frameworks, which yield low-dimensional representations of ensemble forecasts while respecting their probabilistic character.
  arXiv  Detail & Related papers  (2025-02-06T10:16:47Z)
• Lines of Thought in Large Language Models [3.281128493853064]
  Large Language Models achieve next-token prediction by transporting a vectorized piece of text across an accompanying embedding space.<n>We aim to characterize the statistical properties of ensembles of these 'lines of thought'<n>We find it remarkable that the vast complexity of such large models can be reduced to a much simpler form, and we reflect on implications.
  arXiv  Detail & Related papers  (2024-10-02T13:31:06Z)
• Bayesian Circular Regression with von Mises Quasi-Processes [57.88921637944379]
  In this work we explore a family of expressive and interpretable distributions over circle-valued random functions.<n>For posterior inference, we introduce a new Stratonovich-like augmentation that lends itself to fast Gibbs sampling.<n>We present experiments applying this model to the prediction of wind directions and the percentage of the running gait cycle as a function of joint angles.
  arXiv  Detail & Related papers  (2024-06-19T01:57:21Z)
• Curve Your Attention: Mixed-Curvature Transformers for Graph Representation Learning [77.1421343649344]
  We propose a generalization of Transformers towards operating entirely on the product of constant curvature spaces. We also provide a kernelized approach to non-Euclidean attention, which enables our model to run in time and memory cost linear to the number of nodes and edges.
  arXiv  Detail & Related papers  (2023-09-08T02:44:37Z)
• Semisupervised regression in latent structure networks on unknown manifolds [7.5722195869569]
  We consider random dot product graphs, in which an edge is formed between two nodes with probability given by the inner product of their respective latent positions. We propose a manifold learning and graph embedding technique to predict the response variable on out-of-sample nodes.
  arXiv  Detail & Related papers  (2023-05-04T00:41:04Z)
• VTAE: Variational Transformer Autoencoder with Manifolds Learning [144.0546653941249]
  Deep generative models have demonstrated successful applications in learning non-linear data distributions through a number of latent variables. The nonlinearity of the generator implies that the latent space shows an unsatisfactory projection of the data space, which results in poor representation learning. We show that geodesics and accurate computation can substantially improve the performance of deep generative models.
  arXiv  Detail & Related papers  (2023-04-03T13:13:19Z)
• Neural Jacobian Fields: Learning Intrinsic Mappings of Arbitrary Meshes [38.157373733083894]
  This paper introduces a framework designed to accurately predict piecewise linear mappings of arbitrary meshes via a neural network. The framework is based on reducing the neural aspect to a prediction of a matrix for a single point, conditioned on a global shape descriptor. By operating in the intrinsic gradient domain of each individual mesh, it allows the framework to predict highly-accurate mappings.
  arXiv  Detail & Related papers  (2022-05-05T19:51:13Z)
• Joint Estimation of Image Representations and their Lie Invariants [57.3768308075675]
  Images encode both the state of the world and its content. The automatic extraction of this information is challenging because of the high-dimensionality and entangled encoding inherent to the image representation. This article introduces two theoretical approaches aimed at the resolution of these challenges.
  arXiv  Detail & Related papers  (2020-12-05T00:07:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.