Paper Review | Language Models Represent Space and Time

Is there a world model inside large language models?

English 中文

Introduction

Are the many astonishing abilities of large language models nothing more than the product of statistics over massive data? Or does a world model exist inside them? Recently, a paper posted on arXiv by two MIT authors, “Language Models Represent Space and Time”, drew wide attention in the research community. Taking the Llama-2 family as their subject, the authors ran experiments on self-constructed spatial and temporal datasets and found evidence for the latter.

The paper delivers six contributions: four major ones and two minor ones.

The six contributions (1–4 major, 5–6 minor):

  1. Produced three spatial datasets and three temporal datasets

  2. Showed that LLMs do represent space and time, and that this depends on the model and the dataset

  3. Showed that these representations are linear

  4. Showed that these representations are robust to changes in the prompt

  5. Robustness checks

  6. Space neurons and time neurons

The first contribution — six new datasets — is a solid amount of work in itself.

Contributions two through four provide indirect evidence that a world model exists inside LLMs. Prior to this, similar experiments had only been carried out on small models.

As for the fifth and sixth experiments, I would say they exist for the sake of rigor and to pad the list of novelties. The sixth in particular — claiming the existence of “space neurons” and “time neurons” — is hard to read as anything other than an attention-grabbing headline.

Before diving in, a few terms are worth clarifying for the reader.

  1. What I call “representation” corresponds to “representation” in the original paper. I believe the authors use the word to refer to the state of the LLM’s internal hidden layers. For example, the claim quoted above that “LLMs represent time and space linearly” corresponds, in the experiments, to the mapping from the LLM’s internal hidden state to the $label$ being linear.

  2. What I call an “entity” corresponds to “entity” in the original paper. Here the word carries the sense of “sample” — one example from a dataset.

  3. What exactly is the somewhat mystical “world model”? If an LLM only predicts the next token likely to follow a stretch of text, then no world model exists inside it. Conversely, if we can find evidence of structured storage of space and time inside the LLM, then we can say a world model does exist.

  4. The authors’ experiments center on “probes” (“probe” in the original). Let the input be $X$, the model output $Y$, and the hidden state $h$. Train a shallow model using $h$ as the feature and $Y$ as the label; that shallow model is the probe.

  5. Space-time” is shorthand for “space and time”. In the original paper, space appears as both “spatial” and “space”, and time as both “temporal” and “time”.

Paper Walkthrough

Abstract

The rise of large language models (LLMs) has sparked a debate over whether such systems merely perform shallow statistics over massive corpora or genuinely learn a world model. Using Llama-2 as the subject, the authors run experiments on three spatial datasets and three temporal datasets and find evidence for the latter. They find that across model sizes, LLMs represent time and space linearly; that these representations are robust to prompt variations; that they generalize across entity types (for example, cities and landmarks); and they identify “space neurons” and “time neurons”.

Motivation

Although the arrival of LLMs has attracted attention from every quarter, our understanding of what happens inside them remains very limited. Whether an LLM’s operation amounts to statistics over massive data or to something like thinking has been contested since LLMs first appeared.

In this paper, the authors attempt to address the question from the angle of space and time.

Preliminaries

Datasets

The authors created six datasets: three spatial and three temporal. Every entity in each dataset consists of a name and a corresponding location or time. The three spatial datasets are built from place names in the world, the United States, and New York City. The three temporal datasets are built from historical figures, artworks, and news.

DatasetExamples
WorldLos Angeles, Saint Petersburg, the Caspian Sea
USAFenway Park, Columbia, Riverside County (in California)
NYCBorden Avenue Bridge, Trump International Hotel
FiguresCleopatra, Dante Alighieri, Carl Sagan
ArtworksIt, Bohemian Rhapsody
Newsomitted

Model

The Llama-2 family.

Activation datasets

Suppose a dataset contains n entities and the LLM has a hidden dimension of $d_{model}$. Feeding the n entities through the LLM and recording the activations at a given hidden layer yields an activation dataset $A\in\mathbb{R}^{n\times d_{model}}$.

Probes

The authors feed place names, person names, and so on as the input $x$ to the model. Using the hidden state $h$ as the $feature$ and the $y$ corresponding to $x$ as the $label$, they train a probe and evaluate its performance on a test set.

Understanding intermediate layers using linear classifier probes. Probing classifiers: Promises, shortcomings, and advances.

Evaluation metrics

$R^2$ and Spearman rank correlation.

Main Experiments and Results

1. Do LLMs internally represent space and time? If so, where — and does it depend on model size?

The authors train probes on the six activation datasets derived from their six self-built datasets, then evaluate them on held-out test sets, obtaining probe accuracy across LLM sizes, layer depths, and datasets.

Five findings emerge from the figure above.

  1. Both spatial and temporal features can be fit by a linear probe

  2. The larger the model, the higher the probe accuracy

  3. Accuracy rises smoothly over the first 50% of the layers

  4. Among the three spatial datasets, NYC performs worst. The authors speculate this is because the NYC map covers a small area and many of its place names are obscure, so the model simply does not know them

  5. On the NYC map, the accuracy gap induced by model size is the most pronounced

To my mind, the steadily rising curve already shows that the model’s internal representation of spatio-temporal information has a certain structural regularity — enough for a probe to learn it. The fourth and fifth findings also lend indirect support to the possibility that a knowledge base or store exists inside the model. If that possibility can be confirmed, LLMs might serve as a form of data compression.

2. Is the LLM’s representation of spatio-temporal information linear?

Here the authors first hedge a little, spelling out how their work differs from earlier efforts. They note that while prior work has argued for the linear representation hypothesis inside neural networks, that work targeted discrete tasks such as binary or multi-class classification, whereas the labels in this paper’s task are continuous.

The authors train both a simple linear probe and a nonlinear MLP probe and evaluate them on the test set; as the figure shows, the accuracy gap between the two is essentially nonexistent. For a mapping $f$, if it is linear, then both linear and nonlinear models can fit it well and the performance gap is small. If it is nonlinear, the linear model fits worse than the nonlinear one. The authors therefore conclude that the mapping from $h$ to $label$ is linear — which is what the paper means when it says “LLMs represent time and space linearly”.

3. Sensitivity to the prompt

This experiment was designed to test whether spatio-temporal representations are sensitive to changes in the prompt.

For the five datasets other than news, the authors designed four prompts.

  1. empty: no additional prompt

  2. coords/when: adding a descriptive, disambiguating prompt such as “What is the latitude and longitude of <place>” or “What was the release date of <author>’s <book>”

  3. baseline: append 10 random tokens

  4. all_cap: all uppercase

For the news dataset, the authors designed four prompts.

  1. empty prompt

  2. add a period

  3. add a descriptive, disambiguating prompt

  4. add a descriptive, disambiguating prompt plus a period

Looking at the figure above, the authors observe the following.

  1. There is almost no accuracy difference between empty and the descriptive, disambiguating prompts

  2. Baseline accuracy is very low — random tokens degrade accuracy substantially

  3. Uppercasing also degrades accuracy, though this is unsurprising since earlier work has reported the same phenomenon

  4. On the news dataset, adding a period improves accuracy substantially

Overall, spatio-temporal representations are fairly robust to the prompt.

Limitations

At the end of the paper the authors mention several limitations.

  1. A discrete encoder could be used to extract the model’s spatio-temporal representations

  2. The model does not recognize every entity in the datasets; in principle those should be filtered out

Afterthoughts

My first reaction on finishing the paper was “is that all?” The contributions and novelty do stand out, but next to the “world model” the media made of it, the framing is hard not to call clickbait. That said, the clickbait was probably not the authors’ intent — some outlets may simply have been chasing the hype.

On contribution: collecting, building, and curating the datasets alone constitutes enough work for a top-tier conference paper, to say nothing of deploying probes on every layer of multiple LLMs. In sheer volume of work, it is far ahead of comparable papers. On novelty, the paper does not innovate much methodologically; it transplants techniques that were routine on small models onto large ones. And while the findings are not as mystical as the headline-chasers suggest, the discovery that LLMs represent space and time linearly is novel enough for a top-tier paper.

On content, this is a textbook “experiment report”: no mathematical proofs, no derivations — just experiment design, execution, and analysis of results.

On language, the word choice and writing are carefully done. As a Chinese reader I found it somewhat demanding, a very different experience from reading papers written by Chinese authors.

All told, this is a paper comfortably above the bar for a top-tier conference, but it falls short of the “proof of a world model” that the headline-chasers advertised.

Next
Previous

Related