[論文レビュー] 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

ひとこと要約

要約: この論文は制約された2Dグリッドウォーカー系列上でデコーダーのみのトランスフォーマーを訓練し、その内部表現が解析的に導出された予測ベクトルと一致することを示し、 toyシステムにおける次トークン予測から世界モデル風の幾何が出現することを明らかにする。

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.

研究の動機と目的

• 次トークン予測が制御された設定で内部の世界モデル風表現を生み出すメカニズムを調査する。
• 制約付きグリッドウォーカー過程の解析的予測ベクトルを導出する。
• 層を跨ぐトランスフォーマーの活性化と真の予測ベクトルとの整合性を評価する。
• 学習表現の次元と幾何構造を検討する。
• 神経ネットワークにおける文法様式と世界モデルの解釈への示唆を論じる。

提案手法

• 固定されたエンドポイントと時間的ホライズンを持つ2Dグリッドウォーカー過程を定義し、解析的に扱いやすい次ステップ確率を生成する。
• これらのウォークの正確な分布からサンプルした prefixes に対してデコーダーのみのトランスフォーマーを訓練し、次トークンのクロスエントロピー損失を適用する。
• 地上真の2D十分予測ベクトル v_t を計算し、ウェイト付きアフィンプローブを介してトランスフォーマーの活性化をこれらのベクトルに写像する。
• アフィン適合のR^2と層間の線形中心化カーネル整合性(lCKA)を用いて整合性を評価する。
• 層の活性化に対するPCAを用いて表現の固有次元性を分析する。
• 六つのグリッドウォーカー構成（六つのウォーカー）と三つの時間ホライズンで六つの同一モデルを訓練し、予測幾何を探る。

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 予測幾何はトランスフォーマー表現内の低次元多様体に組み込まれているか？
• RQ3 層の深さと時間ホライズンは学習表現と予測ベクトルとの整合性にどのように影響するか？
• 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

• トランスフォーマーは層間の活性化と真の予測ベクトルの間に高い線形整合性を一貫して示し、LayerNorm の活性化が特に強い類似性を示す。
• 学習された表現は低次元の部分空間に存在することが多く、最終表現には2次元で十分な場合がある。
• 整合性指標（R^2とlCKA）はウォーカー、ホライズン、層タイプにより異なり、エンドポイントが固定された場合には長いホライズンで類似性が低下する傾向がある。
• 真の予測ベクトルは「世界位置」を残り時間で割った値に対応し、予測学習から具体的な世界モデル幾何が現れることを示す。
• 最終層正規化表現は多くの設定で2次元空間へ収束する一方、初期の層では次元が異なる。
• グリッドウォーカーの予測状態は確率単位を内包する曲線状の2D多様体を形成し、予測幾何と世界モデル表現の関連性を支持する。

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