PaperlayerView exact revision on arXiv

Source-linked AI summary

Language Models Represent Space and Time

Wes Gurnee, Max Tegmark

arXiv:2310.02207v3cs.LGcs.AIcs.CL

TL;DR

The paper asks whether LLMs learn coherent, grounded world representations or merely superficial correlations. Probing Llama-2 and Pythia activations across spatial and temporal datasets, it finds linear, prompting-robust representations unified across entity types, alongside space- and time-sensitive neurons.

  • Problem

    The paper examines whether text-only LLMs learn coherent models of the real-world data-generating process rather than only massive collections of correlations.

  • Method

    The authors probe activations across spatial and temporal datasets, evaluate regression and proximity metrics, test prompting and cross-entity generalization, and identify neurons encoding space and time.

  • Results

    LLMs recover linear spatial and temporal features that improve with model scale, plateau around halfway through the layers, remain fairly prompt-robust, unify across entity types, and include sensitive individual neurons.

  • Takeaways & Limitations

    The findings provide evidence that modern LLMs learn rich spatiotemporal representations and possess basic ingredients required for a more comprehensive world model.

  • Takeaways & Limitations

    Cross-entity results may reflect separate entity-type directions, and high Spearman correlation can persist because entity types are unevenly distributed in time.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper tests whether LLMs learn coherent spatial and temporal world models by extracting coordinate representations from model activations. Across multiple datasets and scales, the results support linear, robust, and unified spatiotemporal representations.

  • The study investigates whether text-trained LLMs learn coherent models of the real-world processes underlying their training data, rather than only superficial correlations.
  • The authors attempt to extract an actual map of the world from LLM representations, while noting that these representations are not themselves a dynamic causal world model.
  • Six datasets cover world, United States, and New York City locations plus historical figures, cultural works, and news headlines across multiple spatial and temporal scales.
  • Spatial and temporal representations emerge in early layers, plateau around the model midpoint, and improve with model scale.
  • The representations are linear, fairly prompt-robust, and unified across entity types, while individual neurons also encode spatial and temporal features.

2 EMPIRICAL OVERVIEW

The empirical overview constructs six multiscale datasets and probes Llama-2 activations to decode spatial and temporal coordinates. Ridge-regression probes are evaluated on held-out data using standard correlation metrics and a locality-sensitive proximity error.

  • 2.1 SPACE AND TIME RAW DATASETS: The six datasets pair entity names with spatial or temporal coordinates and include multiple entity types to test representation unification.
  • 2.1 SPACE AND TIME RAW DATASETS: The spatial datasets cover places worldwide, across the United States, and within New York City.
  • 2.1 SPACE AND TIME RAW DATASETS: The temporal datasets cover historical figures, cultural works released from 1950 to 2020, and New York Times headlines from 2010 to 2020.
  • 2.2 MODELS AND METHODS: For each entity, the experiments save the last-token residual-stream activation at every layer, producing an n × dmodel activation dataset.
  • 2.2 MODELS AND METHODS: Linear ridge-regression probes predict time or latitude and longitude coordinates from model activations, with λ selected by leave-one-out cross-validation.
  • 2.3 EVALUATION: Evaluation reports test-set R2 and Spearman correlation plus proximity error, which measures local prediction precision relative to nearby entities.

3 LINEAR MODELS OF SPACE AND TIME

Linear probes recover spatial and temporal features across models and datasets, with quality increasing through early layers before plateauing and improving with model scale. These representations are largely linear and relatively prompt-robust, although distracting tokens and capitalization can reduce performance.

  • 3.1 EXISTENCE: Spatial and temporal features are linearly recoverable across datasets, with probe quality increasing through the first half of layers before plateauing and improving with model scale.The gap between Llama and Pythia models was especially striking, so subsequent results focus on Llama models.
  • 3.1 EXISTENCE: The New York City dataset performs worst, but the largest model shows the best relative performance there.This pattern suggests that sufficiently large LLMs could eventually form detailed spatial models of individual cities.
  • 3.2 LINEAR REPRESENTATIONS: Nonlinear MLP probes provide minimal R2 improvement over linear ridge probes, supporting linear or at least linearly decodable representations of continuous space and time.The comparison uses one-layer MLP probes with 256 neurons.
  • 3.3 SENSITIVITY TO PROMPTING: Explicit recall prompts and location disambiguation hints make little to no difference, while random distracting tokens and capitalization degrade performance.The results are reported for Llama-2-70B in Figure 3 and for all models in Figure 8.
  • 3.3 SENSITIVITY TO PROMPTING: Probing the period token after a news headline notably improves performance, suggesting that periods contain summary information about the sentences they end.The headline experiments compare the last headline token with an appended period token.

4 ROBUSTNESS CHECKS

The robustness checks test whether spatial and temporal probes generalize beyond memorized blocks, entity types, and high-dimensional activation patterns. Results provide evidence for explicit learned features, while leaving binary-feature explanations unresolved.

  • 4.1 VERIFICATION VIA GENERALIZATION: The binary-features explanation remains possible because a hierarchy of features could avoid country or decade boundaries.Thus, block generalization does not fully establish that the model directly represents geometric coordinates rather than membership-like features.
  • 4.1 VERIFICATION VIA GENERALIZATION: Block holdout tests remove one country, state, borough, century, decade, or year from training before evaluating predictions on that block.The comparison uses average proximity error for completely held-out blocks against test points from the default split; random performance is 0.5.
  • 4.1 VERIFICATION VIA GENERALIZATION: Held-out-block performance is better than random, with predictions preserving relative positions more reliably than absolute positions.The authors interpret this as weak evidence that probes extract learned features while memorizing the transformation into human coordinates.
  • 4.1 VERIFICATION VIA GENERALIZATION: Cross-entity holdouts show that probes largely generalize across entity types, except mainly in the entertainment dataset.The analysis compares proximity error when entity subtypes are included in training versus fully held out.
  • 4.2 DIMENSIONALITY REDUCTION: Principal-component probes use 2 to 3 orders of magnitude fewer parameters, testing whether compact activation directions retain spatial and temporal information.Figure 4 compares test R2 across k largest principal components with full dmodel-dimensional probes.
  • 4.2 DIMENSIONALITY REDUCTION: Spearman correlation rises faster than R2 as principal components increase, suggesting coordinates are present in top components but require more parameters for literal values.The first several components also cluster entity types, explaining why more than a few components are needed.

5 SPACE AND TIME NEURONS

The authors identify individual neurons that encode spatial and temporal features, including across entity types, while noting that these neurons provide a lower-bound estimate of representation strength.

  • Neurons selected by cosine similarity to probe directions are highly sensitive to entities’ true spatial or temporal locations.The search targets neurons whose input or output weights align with learned probe directions.
  • These individual neurons respond to all entity types in the datasets, strengthening evidence that the representations are unified.
  • Individual-neuron performance is a lower bound on the model’s representation of spatial and temporal features because features are generally expected to be distributed in superposition.
  • Neuron ablation and intervention experiments were conducted to verify the importance of these neurons in spatial and temporal modeling.

6 RELATED WORK

The paper builds on work showing that language models and related systems can encode spatial, temporal, factual, and game-state information, while positioning its contribution within interpretability research.

  • Linguistic Spatial Models: Prior studies recovered geographic information from language, embeddings, and small language models, though one study characterized the learned geographic information as limited.
  • Neural World Models: Research on chess and Othello-trained GPT-style models found explicit, later linear representations of board and game state.
  • Factual Recall: Work on factual recall indicates that early-to-mid MLP layers output information about factual subjects, often at the subject’s last token, with related linear structures.
  • Interpretability: The study draws on probing, BERTology, the linearity hypothesis, superposition, and mechanistic interpretability.

7 DISCUSSION

The authors conclude that LLMs learn unified, prompt-robust linear representations of space and time, but do not establish that these primitives form a causal world model.

  • LLMs learn linear spatial and temporal representations that are unified across entity types and fairly robust to prompting.
  • Individual neurons are highly sensitive to spatial and temporal features, providing a more localized signal within the broader representations.
  • The analysis can reconstruct absolute position in space or time, but the true extent and structure of these representations remain unclear.
  • The authors conjecture that the representations may form a discretized hierarchical mesh with multiple granularities, but present this as future work.
  • The study barely examines how these models are learned, recalled, and used internally or how they relate to a comprehensive world model.
  • Biological place and grid cells are identified as a potential source of inspiration for future LLM research.

A DATASETS

The appendix describes six datasets spanning spatial and temporal scales, assembled from public knowledge sources and filtered or transformed to support entity-level analysis.

  • The datasets cover world, US, and NYC places plus historical figures, art and entertainment, and news headlines.
  • World Places: World places were drawn from DBpedia, matched to Wikipedia, and filtered to entities receiving at least 5,000 page views over three years.
  • World Places: The world-place data includes physical, natural, and structural locations, with page views used as a proxy for whether an LLM is expected to know a place.
  • Known Limitations: The data are limited by English-Wikipedia and Anglosphere skew, nonuniform entity types, and possible correlations such as the overrepresentation of UK railway stations.
  • USA Places: The US dataset combines places and colleges with census-aggregator data for counties, zipcodes, and cities, while removing most duplicate names.
  • NYC Places: The NYC dataset uses city-government points of interest, including schools, places of worship, transit locations, roads, bridges, and public buildings.
  • Historical Figures: The historical-figures dataset includes 350 of the most famous people from each decade who died between 1000BC and 2000AD.
  • Art and Entertainment / Headlines: Art and entertainment entries pair songs, movies, and books with creators to disambiguate generic titles, while headlines are filtered from New York Times coverage to retain localizable events.

B NEURON ABLATIONS AND INTERVENTIONS

The authors probe whether individual neurons encode space and time by ablating them and intervening on their activations. Pinning a single time neuron can alter next-token predictions, while ablation identifies contexts where spatial and temporal neurons matter most.

  • Time Intervention: A single-neuron intervention tests whether changing a time-neuron activation affects language-model predictions.The experiment sweeps pinned activation values across all tokens and tracks the probabilities of the five most likely next tokens.
  • Time Intervention: 0.77 correlation links time neuron L19.3610 in Llama-2-7b to art and entertainment release dates.
  • Time Intervention: Pinning time neuron L19.3610 to different values changes the predicted next tokens for media-release prompts.The comparison is against nine random neurons from the same layer, as shown in Figure 7.
  • Neuron Ablations: Neuron ablations compare normal model loss with loss after individually setting two space neurons and two time neurons to zero.The analysis focuses on Wikipedia contexts involving world places and contemporary art and entertainment.
  • Neuron Ablations: Tables 5–8 list the contexts with the highest loss after ablating two space neurons and two time neurons.The reported neurons are space neurons L20.7573 and L20.7423, and time neurons L18.9387 and L19.3610.

C ADDITIONAL RESULTS

Additional analyses visualize learned spatial coordinates, test out-of-sample generalization, and examine how representation quality varies with prompting and principal-component projections.

  • Prompting: Out-of-sample R2 is evaluated across prompts that include entity names in different ways for all models.
  • Spatial Representations: Heldout world and US places are projected onto learned latitude and longitude directions and colored by true continent or state.The visualizations use test-set activations from Llama-2-70b layer 50.
  • Out-of-Sample Generalization: Country and state predictions are evaluated when probe training excludes samples from the heldout country or state.The figures compare predicted locations with true locations and the training-data mean.
  • Principal Components: Test Spearman rank correlation is measured for probes trained on activations projected onto the k largest principal components.
Loading 2310.02207v3…