[论文解读] Grid-World Representations in Transformers Reflect Predictive Geometry

Sasha Brenner, Thomas R. Knösche+1 more|arXiv (Cornell University)|Mar 17, 2026

Neural Networks and Applications被引用 0

一句话总结

本文在受限的二维网格行走序列上训练解码器仅 transformers，并展示其内部表示与解析推导的预测向量对齐，揭示在一个 toy 系统中通过下一个 token 预测出现的世界模型风格几何。

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.

研究动机与目标

• 在受控设定中调查下一个 token 预测如何诱发内部的世界模型风格表示。
• 推导受限网格行走过程的解析预测向量。
• 评估 transformer 激活与真实预测向量在各层之间的对齐情况。
• 考察学习表示的维度性和几何结构。
• 讨论对神经网络语法与世界模型的解释意义。

提出的方法

• 定义具固定端点与时间步长的二维网格行走过程，以生成解析可处理的下一步概率。
• 在从这些游走的精确分布中采样的前缀上训练解码器仅 transformers，使用下一个 token 的交叉熵损失。
• 计算真实的二维充分预测向量 v_t，并通过加权仿射探针将 transformer 的激活映射到这些向量。
• 在各层使用仿射拟合的 R^2 和线性中心核对齐（lCKA）评估对齐情况。
• 通过对层激活进行 PCA 分析表示的内在维度。
• 使用六个在六种网格行走配置上训练的相同 transformers（六个行走者）以及三个时间 horizons 来探索预测几何。

Figure 1: Transformer layer representations closely resemble the ground-truth predictive vectors of Equation 4 (as computed in Appendix A.3 ). The left column corresponds to the ground-truth principal components (PCs), and the right column shows the PCs of the last layer norm representation, after P

实验结果

研究问题

• RQ1 变换器激活是否与网格行走任务中解析推导的充分预测向量对齐？
• RQ2 预测几何是否在 transformer 表征中编码为低维流形？
• RQ3 层数和时间 horizon 如何影响学习表示与预测向量之间的对齐？
• RQ4 预测状态几何是否为序列建模中的语法样约提供直觉？

主要发现

ENDPOINT	HORIZON	MEAN VAL LOSS
(0,0)	20	4.5e-6
(0,0)	200	2.6e-7
(0,0)	1000	5.5e-5
(4,0)	20	5.9e-6
(4,0)	200	<1e-8
(4,0)	1000	<1e-8

• transformers 在各层的激活与真实预测向量之间呈现出高线性对齐，LayerNorm 的激活尤为相似。
• 学习的表示常嵌入到低维子空间，某些情况下最终表示可用 2 维表示。
• 对齐度量（R^2 与 lCKA）因行走者、时间 horizons 与层类型而异，较长的时域对同一端点的一致性通常较低。
• 真实的预测向量对应于“世界位置”除以剩余时间，展示了通过预测学习涌现的具体世界模型几何。
• 在若干设置中最终的 LayerNorm 表示收敛到一个二维空间，而早期层的维度性则有不同。
• 网格行走的预测状态在概率简单形中形成一个曲面状的二维流形，支持预测几何与世界模型表示之间的联系。

Figure 2: Linear similarity metrics between layer activations and ground-truth predictive vectors. (A) and (D): Both $R^{2}$ and $\rm{lCKA}$ are high and tend to increase with layer depth, with LayerNorms generally exhibiting stronger similarity. Moreover, walkers with longer time horizons $T$ have