Paper: Language Models Represent Space and Time (ICLR-2024)

Reviewer comment (Novelty):
One of the paper's limitations is its potential overlap with findings already established in the word embedding literature. Earlier works in this area, such as those by Mikolov et al. (2013), have demonstrated that simpler word embedding models can capture relational information and regularities in a linear fashion. This raises the question of whether the findings in this paper are genuinely novel or simply an extension of what is already known about linear representations in language models. The paper could strengthen its contribution by more directly addressing how its findings with large language models (LLMs) differ significantly from those with simpler models, and by delving deeper into the implications of these differences for the field.

Author response:
Thank you for your thoughtful comments!

> W1. One of the paper's limitations is its potential overlap with findings already established in the word embedding literature.

We agree this paper is inspired by prior work on embeddings, and that the observation that there exists linear structure in learned representations that reflect real world structure is an old and well known fact. 

However, static word embeddings are much less interesting objects of study and these past works mostly show that there exists structure in the most simple setup. For example, Mikolov et al. (2013) took 11 countries and their capitals in the northwestern hemispheres and did a PCA which showed a roughly correct ordering with no quantitative metrics. We feel that our expansion to almost one hundred thousand places on three scales of space and time and more thorough quantitative analysis provides a major step forward compared to this interesting prior work, and that this is of importance given recent discourse about the extent to which LLMs are "stochastic parrots". 

> W1. The paper could strengthen its contribution by more directly addressing how its findings with large language models (LLMs) differ significantly from those with simpler models, and by delving deeper into the implications of these differences for the field.

Specifically, we think our work is novel compared to work on similar topics by:
* Also considering temporal representations
* Studying models with more than XXXM parameters (we go up to 70B which we show is quite important) and using far more probing data (at least 10x more)
* Studying spatial and temporal representations at different spatiotemporal scales (eg New York City, US based, and world map)
* Studying how different types of entities are represented (eg, natural landmarks and physical structures), whereas all prior work mostly focuses on cities.
* Studying aspects unique to language models, like where the representations form in the model, or how sensitive the representations are to prompting.
* Beginning to dive into model internals, and identifying individual model components which respond to spatial and temporal representations.

Furthermore, work on embeddings is typically restricted only to places which are tokenized as a single word, and hence the datasets are far smaller (and likely even less globally diverse). Another important contribution of our work is the datasets we created, each of which on its own would be 10x more data than any prior study used, and we created 6 such datasets spanning diverse entity types and scales. We make these datasets publicly available in our github repository and believe they will enable important future work on understanding spatial and temporal reasoning and interpretability.

Reviewer comment (Theory):
The assertion that LLMs learn literal 'world models' may be overstated. The concept of a 'world model' in the context of AI and cognitive science typically refers to a comprehensive and dynamic representation of the external world, including an understanding of causal relationships and the ability to predict future states. The paper's findings, while impressive, primarily demonstrate that LLMs can encode spatial and temporal information linearly. This is a far cry from the richer and more dynamic conception of a world model. A more accurate claim might be that LLMs are capable of forming structured and useful representations of spatial and temporal information, but these do not necessarily constitute comprehensive world models. Further exploration into how these representations are utilized by LLMs in real-world tasks, and how they compare to human cognitive processes, could provide a more nuanced understanding of their nature and limitations.

Author response:
> W2. The assertion that LLMs learn literal 'world models' may be overstated. 

We meant “literal world models” to mean “a literal model of the world” which, in hindsight, we agree was too glib - we wish to apologize for this overstatement. That said, we still believe such representations are a necessary condition for a more stringent conception of a world model, and we think our results are an important and foundational step in understanding to what degree there exists such a model. We edited the abstract, introduction, and conclusion in the submitted revision to make our claims more specific, in line with what you recommend.

> W2. Further exploration into how these representations are utilized by LLMs in real-world tasks...

To better understand how these representations are used, and to answer Q1, we perform neuron ablation experiments. See the change log comment and appendix B for analysis.

> W2. and how they compare to human cognitive processes

We think experimental comparisons to how humans represent space and time is out of scope for this study, but we do think this is an exciting area for future work (as discussed in the last paragraph of the discussion section) and note that just recently “A neural code for time and space in the human brain” by Schonhaut et al. was published in Cell two weeks ago and show the existence of time and place cells similar to our space and time neurons! We added this reference to the last paragraph in the discussion.

Reviewer comment (Experiments):
What if the space or time neurons are removed/masked? How would it influence the models’ understanding of space or time?

Author response:
> What if the space or time neurons are removed/masked? How would it influence the models’ understanding of space or time?

We performed several ablation studies which we describe in the comment to all reviewers and added to the appendix of the revised article.

Reviewer comment (Experiments):
Are the neurons that the authors point out as aligning with the space/time directions sufficient, necessary, or not even sufficient for representing these properties? In other words, what happens to the space/time representations if these neurons are ablated and the ridge regression probes are relearned?

Author response:
> Q1. Are the neurons that the authors point out as aligning with the space/time directions sufficient, necessary, or not even sufficient for representing these properties? 

We do not believe these neurons are necessary or sufficient for representing these properties, but instead take them as strong evidence that the model does in fact learn the correct global representations. That is, because of how late these neurons occur in the model (typically around 50-70% depth), we believe that these neurons are “consumers” rather than “producers” of these representations, and use them for some downstream task. 

>In other words, what happens to the space/time representations if these neurons are ablated and the ridge regression probes are relearned?

Given that the probe performance plateaus before this point, maximum probe predictive accuracy would not be affected. 

To better understand the function of these neurons, we perform new neuron ablation experiments, and analyze the LLM loss increase, rather than the difference in probe performance (see the general change log comment for analysis).

Reviewer comment (Writing):
Page 3 is the first time when it's clear that what is meant by a spatial representation is the two-dimensional latitude and longitude coordinates. This should be made clear earlier in the paper.

Author response:
> Minor point: page 3 is the first time when it's clear that what is meant by a spatial representation is the two-dimensional latitude and longitude coordinates. This should be made clear earlier in the paper.

Thank you, we added a note of this in the third paragraph of the introduction.

Reviewer comment (Experiments):
The experiments only involve Llama-2 series models, whereas various locations in the paper stretch the claim to be about all LLMs and modern LLMs. I recommend rephrasing some texts into e.g., “LLMs, with Llama as examples” to make the claims better supported by the scope of the experiments.

Author response:
Excellent point. To address this, we've now replicated our main results on the Pythia model suite to support the generality of our findings. We chose the Pythia family because it spans a large number of model sizes, has many training checkpoints, and is a popular model to study in the interpretability literature.

As can be seen in the revised Figure 2, the main results continue to hold, with fairly clean scaling. However, because all Pythia models are only trained on 300B tokens, the large models are fairly under-trained compared to the 2T tokens used to train the LLama models.

Reviewer comment (Writing):
There are some minor typos. In page 6, a punctuation is needed before the footnote. In the end of the next paragraph, the period should be before the footnote.

Author response:
Fixed. Thank you for your careful reading!

Reviewer comment (Novelty):
A related work should be referred: [Distributional vectors encode referential attributes](https://aclanthology.org/D15-1002) (Gupta et al., EMNLP 2015). This is one of the earliest probing papers, and it probed for the attributes involving geographic locations.

Author response:
Done: we added this reference to our new related work section on Linguistic Spatial Models.

Reviewer comment (Experiments):
Do the findings generalize to other LLMs, for example bidirectional models?

Author response:
Excellent point. To address this, we've now replicated our main results on the Pythia model suite to support the generality of our findings. We chose the Pythia family because it spans a large number of model sizes, has many training checkpoints, and is a popular model to study in the interpretability literature.

As can be seen in the revised Figure 2, the main results continue to hold, with fairly clean scaling. However, because all Pythia models are only trained on 300B tokens, the large models are fairly under-trained compared to the 2T tokens used to train the LLama models.