[Paper Review] Language Models Represent Space and Time

Wes Gurnee, Max Tegmark|arXiv (Cornell University)|Oct 3, 2023

Language and cultural evolution41 citations

TL;DR

The paper shows that Llama-2 and Pythia language models learn linear, multi-scale representations of space and time, including identifiable space/time neurons, suggesting a rudimentary world model learned from next-token prediction.

ABSTRACT

The capabilities of large language models (LLMs) have sparked debate over whether such systems just learn an enormous collection of superficial statistics or a set of more coherent and grounded representations that reflect the real world. We find evidence for the latter by analyzing the learned representations of three spatial datasets (world, US, NYC places) and three temporal datasets (historical figures, artworks, news headlines) in the Llama-2 family of models. We discover that LLMs learn linear representations of space and time across multiple scales. These representations are robust to prompting variations and unified across different entity types (e.g. cities and landmarks). In addition, we identify individual "space neurons" and "time neurons" that reliably encode spatial and temporal coordinates. While further investigation is needed, our results suggest modern LLMs learn rich spatiotemporal representations of the real world and possess basic ingredients of a world model.

Motivation & Objective

• Investigate whether LLMs learn coherent spatiotemporal representations beyond superficial statistics.
• Extract and analyze internal activations to map space coordinates (latitude/longitude) and time coordinates (timestamps).
• Assess the linearity, robustness, and cross-entity consistency of these representations across scales and entity types.
• Identify individual neurons associated with space and time to establish that the model uses these features.

Proposed method

• Construct six multi-scale datasets of space- and time-related entities (world, US, NYC; historical figures; art; headlines).
• Run linear ridge probes on last-token activations across layers of Llama-2 and Pythia models to predict real-world coordinates (lat/long) or timestamps.
• Evaluate probes with R^2 and Spearman correlations, and use proximity error to account for local geographic precision.
• Test linearity by comparing with nonlinear probes (MLP) to confirm decodability of space/time features.
• Assess prompting sensitivity by varying prompts and observing probe performance across datasets.
• Identify “space neurons” and “time neurons” by projecting activations onto probe directions and analyzing neuron weights.

Figure 1: Spatial and temporal world models of Llama-2-70b. Each point corresponds to the layer 50 activations of the last token of a place (top) or event (bottom) projected on to a learned linear probe direction. All points depicted are from the test set.

Experimental results

Research questions

• RQ1 Do LLMs encode spatial and temporal information in their internal representations?
• RQ2 Are space and time representations linear and decodable from mid-to-late layer activations?
• RQ3 Are these spatiotemporal representations robust to prompting variations and unified across entity types?
• RQ4 Do individual neurons in LLMs encode space or time coordinates, indicating distributed use of these features?

Key findings

• Spatial and temporal features can be recovered with linear probes across layers and scales.
• Representations improve in early to mid layers and plateau around the model's halfway point, with larger models performing better.
• Nonlinear probes offer little improvement over linear probes, supporting linear decodability of space/time.
• Prompts have limited impact on probe performance, while random tokens can degrade performance; period tokens after headlines can improve it.
• Space and time representations generalize across different entity types (cities vs landmarks), and individual neurons align with probe directions, indicating genuine use of these features.
• Probes remain informative even when data is reduced via PCA, suggesting robust, underlying spatiotemporal structure.

Figure 2: Out-of-sample $R^{2}$ for linear probes trained on every model, dataset, and layer.