On Linear Representations and Pretraining Data Frequency in Language Models

Authors: Jack Merullo, Noah A. Smith, Sarah Wiegreffe, Yanai Elazar
Primary area: interpretability and explainable AI
Keywords: pretraining data, pretraining, linear, linear features, interpretability, linear representations, corpus frequency

Abstract:
Pretraining data has a direct impact on the behaviors and quality of language models (LMs), but we only understand the most basic principles of this relationship. While most work focuses on pretraining data's effect on downstream task behavior, we investigate its relationship to LM representations. Previous work has discovered that, in language models, some concepts are encoded "linearly" in the representations, but what factors cause these representations to form (or not)? We study the connection between pretraining data frequency and models' linear representations of factual relations (e.g., mapping France to Paris in a capital prediction task). We find evidence that the formation of linear representations is strongly connected to pretraining term frequencies; specifically for subject-relation-object fact triplets, both subject-object co-occurrence frequency and in-context learning accuracy for the relation are highly correlated with linear representations. This is the case across all phases of pretraining, i.e., it is not affected by the model's underlying capability. In OLMo-7B and GPT-J (6B), we discover that a linear representation consistently (but not exclusively) forms when the subjects and objects within a relation co-occur at least 1k and 2k times, respectively, regardless of when these occurrences happen during pretraining (and around 4k times for OLMo-1B). Finally, we train a regression model on measurements of linear representation quality in fully-trained LMs that can predict how often a term was seen in pretraining. Our model achieves low error even on inputs from a different model with a different pretraining dataset, providing a new method for estimating properties of the otherwise-unknown training data of closed-data models. We conclude that the strength of linear representations in LMs contains signal about the models' pretraining corpora that may provide new avenues for controlling and improving model behavior: particularly, manipulating the models' training data to meet specific frequency thresholds. We release our code to support future work.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730239488043

thread_type:
official_review

<thread_object>
creator: reviewer_juu2

<content>
summary:
This paper finds that the formation of linear representations for factual recall relations in LMs is highly correlated with the frequency of subject-object cooccurrence in the pretraining data. The formation of linear representation can happen at any stage of pretraining as long as the subj-obj cooccurrence exceeds some threshold, i.e., a linear representation can form consistently when the subjects and objects co-occur at least 1-2k times even at early stages of pretraining. The results also indicate that the frequency threshold is related to the model size, and larger models tend to require smaller thresholds. Using the metrics that evaluate the quality of linear representations, the authors can predict the approximate frequencies of individual terms as well as the co-occurrence of terms in the pretraining corpus better than using the LMs' uncertainty alone.

strengths:
This paper draws an interesting connection between pretraining term frequency to the formation of linear representation of factual recall relations. The fact that the formation of linear representations could happen at any pretraining stage is particularly intriguing. The experiments and results are easy to understand and the discussion of related work is comprehensive.

weaknesses:
- Being able to predict the frequencies of individual terms as well as the co-occurrence seems to be a direct implication of high correlation and therefore does not sound like a major standalone contribution. Also, "Importantly, this regression model generalizes beyond the specific LM it was trained on without additional supervision.": the prediction results seem to be very noisy and not much better than the mean frequency baseline.
- Some claims are not properly justified:
  - Line 75: "This frequency threshold decreases with model size": This is only tested two model sizes (OLMo-7B and OLMo-7B), and the fact that GPT-J (6B) has as smaller threshold than OLMo-7B is an counterexample for this. Would be good just to be consistent with the rest of discussion to claim there is a connection to scale.
  - Line 93-94: "Linear representations form at predictable frequency thresholds during training, regardless of when this frequency threshold is met for the nouns in the relation." The term "predictable" can be understood as there is a strong correlation between the linear representation quality and the co-occurrence frequency. However, the fact that this threshold is predictable regardless of when this frequency threshold is met is not well supported by results. It is necessary to show that the threshold (mean causality >.9) is consistent across different checkpoints.
    - Line 319-320: "Regardless of pretraining step, models that surpass this threshold have very high causality scores." It would be nice if you could arrange results into different scatter plots for each pretraining stage and compare them, say by fitting a linear model and comparing their slopes and biases.
  - Line 100: The efficiency of the proposed searching tool is not well discussed.
  - Line 455: "Some relations, like star-constellation perform very poorly, possibly due to low frequency" Why low frequency is the cause?
  - Line 471: "Second, evaluating on the LRE features of a heldout model (scaled by the ratio of total tokens trained between the two models) maintains around the same accuracy," How do the results support "around the same accuracy"? If it is comparing Train OLMo and Train GPT-J in Table 1, the drop of accuracy is larger than the performance gap between LRE features and mean baseline. I am not sure if this entails "around the same accuracy".
  - Line 483: "In general, the regression transfers well, without performance deteriorating much (about 5%), suggesting LREs are encoding information in a consistent way across models." What results support this? Table 2 only shows a few examples which is insufficient for supporting the claim. Are there aggregated numbers of all pairs? How many of them have errors less than 5%?
- Some important details are missing from the experiments
  - The results in Figure 3 and Table 1 do not match: 1) why is the mean freq. baseline performance different? 2) why do LRE features (Table 1: 0.76) seem to perform better than LRE + LM (Figure 3: ~0.67) for OLMo, if Figure 3 shows the results for OLMo. The author should explain how are the numbers related to each other.
  - Table 2: 1) "Predictions are better for relations that are closer to those found in fitting the relation (country-related relations)" What does closer mean here? How did you measure this? 2) "Some relations, like star-constellation perform very poorly, possibly due to low frequency"
- Some figures and tables need to be more carefully explained:
  - The two left plots in Figure 2 need more explanations or presented in a better way. Specifically, why are some points darker than the others? What do the lines (light grey and dark grey lines) mean? Also, why do all dots for GPT-J have the same shape while the dots for the 2 left plots do not?
  - Table 1: "Train OLMo" and "Train GPT-J" are hardly self-explanatory, the authors should consider better ways to explain the settings.

questions:
1. The experiments follow previous work and only analyze 25 relations. What are the reasons for not including other relations?
2. Section 4.3 is interesting to some degree but I am not sure about the implication of the results. Looks like it is just a description of what is observed. What is the research question you want to answer here?
3. Typos:
   1. Line 455 the regression model can "be" sensitive to ...
   2. Line 416: the paragraph does not end properly.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730529540967

thread_type:
official_review

<thread_object>
creator: reviewer_zxua

<content>
summary:
This research paper explores how the frequency of certain words appearing together in the data used to train a language model (LM) affects the LM's ability to learn simple, linear rules for representing facts. The authors found that the more often two words related to a fact appear together in the training data, the more likely the LM is to learn a simple rule to represent that fact. This discovery helps to understand how LMs learn factual information and could be used to figure out what kind of data was used to train secret LMs. The authors also created a tool to help others count how often words appear together in large datasets

strengths:
The study finds a strong correlation  between the average co-occurrence frequency of subjects and objects within a relation and the quality of linear representations (LREs) formed for that relation. This correlation surpasses the individual correlations with subject frequencies or object frequencies, highlighting the significance of subject-object co-occurrence.

The study focuses uses Linear Relational Embeddings (LREs), which effectively approximate the computations performed by an LLM to predict objects in factual subject-relation-object triplets. This paper builds upon this research by examining how the frequency of subject-object co-occurrences in pretraining data directly impacts the emergence and quality of these LREs

The paper introduces a promising technique for analyzing the pretraining data of closed-source models by leveraging the connection between linearity and frequency.

weaknesses:
The paper presents valuable findings, however they should provide some discussion along the following directions:

(a) The paper primarily focuses on Linear Relational Embeddings (LREs) as a representative class of linear representations in LLMs. However, LLMs might employ various other forms of linear or non-linear structures to encode information. This focus on LREs could limit the generalizability of the findings to other types of representations. Is there any strong hypothesis to strict to LREs?

(b) While the study demonstrates that LRE features can be used to predict the frequencies of individual terms with reasonable accuracy, predicting the frequency of subject-object co-occurrences is challenging. The regression models achieve only marginal improvements over baseline performance in this task.  Integrating additional features might be helpful here.

(c) The study analyzes a set of 25 factual relations from the Relations dataset. However, LLMs are trained on vast and diverse data, encompassing a much wider range of relations and concepts. Expanding the scope of analysis to encompass a broader range of relations would provide a more comprehensive understanding of the role of frequency in shaping LLM representations.

(d)  The paper focuses primarily on the frequency of terms in the pretraining data. However, other factors, such as the context in which terms appear, the syntactic structure of sentences, or the semantic relationships between words, could also influence the formation of linear representations. For example, LLMs are proven to not do well is facts  are stored in templates , as it tend to remember the template and not the facts. The proposed approach may not be applicable in those scenarios.

questions:
Refer the previous section.

It would be good if authors can dedicate a section to discuss the potential impact of confounding factors, such as context, syntax, and semantics. Explain why controlling for these factors is challenging in the current study but emphasize the importance of future work to disentangle their effects from the influence of frequency.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730690083707

thread_type:
official_review

<thread_object>
creator: reviewer_j8rp

<content>
summary:
This paper explores the question of why linear structures form in LLMs by investigating the connection between training data frequency and the formation of linear representation, focusing specifically on factual recall relations. The study reveals that (1) the formation of linear representations is strongly correlated with subject-object co-occurrence frequency, and (2) the presence of linear representations can help predict relation frequency. Experiments are conducted using OLMo-1B, oLMo-7B, and GPT-J to validate these findings.

strengths:
- Exploring the origin of linear representation is an important question in LM interpretability. This work identifies a correlation between linear representations of factual recall relations and the subj-obj co-occurrence frequency in pretraining.
- This paper investigates the relationship between few-shot accuracy and the existence of a linear representation.
- Using the existence of linear representations to predict the frequency of terms in the pretraining corpus is interesting.

weaknesses:
- The scope of the work is somewhat limited, as only 25 factual relations are investigated. It is unclear whether the identified correlation is also valid for other relation types. Expanding the analysis to include more factual relations and other types of relations could further enhance the robustness of the findings and potentially offer additional insights.
- The linear representation seems to be affected by the context in LREs (e.g., four "X plays the Y" examples before the fifth one. Are the findings universally applicable to LLM generation without involving ICL formats?

questions:
Please see weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730717034637

thread_type:
official_review

<thread_object>
creator: reviewer_esj4

<content>
summary:
The authors investigate the correlation between linear representations and pre-training data frequency in language models. The work is conducted on recent findings that the linearity of different types of relations varies significantly depending on the specific relationship. Existing work does show that language model exhibit such linear structures, but do not reveal the underlying reason why some relations exhibit such structure while other do not. The main contribution of this work is to empirically draw the correlation between such linear structure and data frequency. It shows that that linear representations for factual recall relations are related to mention frequency and the model size. In addition, more detailed results show that linear representations form at predictable frequency thresholds during training, which allows the prediction of term frequencies in the training data. Finally, the authors release a tool for searching through tokenized text for understanding training data characteristics.

Overall, the findings are insightful for understanding linear representation structures in language models. This empirical study complements existing theoretical evidence on the same subject. It provides a perspective to the problem, which can be among many other factors in driving the formation of linear structures. On the utility side, the findings can be used for understanding training data, which are typically not published for current LLMs.

strengths:
A perspective for understanding the reason that some features from LLMs demonstrate linear structures while others do not.

A tool for search through tokenized text to support the understanding of training data.

weaknesses:
It provides one perspective to the problem empirically, with a specific set of metrics. While giving useful information, the depth of understanding and the utility domain is constrained mostly to the correlation between term frequency, the model size and the linear structure.

questions:
Could there be some theoretical discussion on the training dynamics and the frequency thresholds?

One related work is 

Guangsheng Bao, Zhiyang Teng, and Yue Zhang. 2023. Token-Level Fitting Issues of Seq2seq Models. In Proceedings of the 8th Workshop on Representation Learning for NLP (RepL4NLP) at ACL 2023. Toronto, Canada from July 9th to July 14th, 2023.

which also discusses the correlation between term frequencies and model accuracy.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732663967586

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of changes made during the discussion (and pending cosmetic changes)

comment:
Thank you to the reviewers for their work in evaluating this work. We are happy that the reviewers agreed the work was interesting and valuable, as well as important for model interpretability. Additionally, we are glad that reviewers were generally excited about the potential use of our findings in tools for making inferences about the training data of language models (LMs). We will outline the contributions, points that reviewers made, and our responses here:

## Contributions
1. We identify a correlation between linear representations (in the form of linear relational embeddings (LREs), see [Hernandez et al., 2024](https://arxiv.org/abs/2308.09124)) of factual recall relations and the subject-object co-occurrence frequency in pretraining.

2. We introduce a tool for quickly searching and counting token occurrences in training data batches that offers more flexibility than existing tools.

3. We leverage the connection between linear representational structure and frequency to show that we can use the presence and 'strength' of an LRE to predict the pretraining frequency of *individual terms* in the pretraining data, allowing us to make inferences about the pretraining data of open-weights models.


## Changes

The biggest change we made was **adding more relations**. There were concerns that 25 relations is not enough or that limiting to factual relations was not telling a general enough picture. We want to emphasize that this is not a small dataset, across the 25 relations, we had **over 10,000** unique subjects and objects. We also focus on factual relations following prior work ([Elsahar et al., 2018](https://aclanthology.org/L18-1544/)) which finds that using subject-object co-occurrences is a good proxy for counting factual mentions. Still, there are interesting questions around how the analysis extends to other relations. To answer this, we included 8 additional commonsense relations. These are: fruit_inside_color, fruit_outside_color, object_superclass, substance_phase, task_done_by_person, task_done_by_tool, word_sentiment, work_location. We added these as Appendix F. While we find interesting relationships between frequency and causality that mirror the factual relations, we also raise issues with using subject-object co-occurrences as counts for some of these relations.

We outline the remaining changes requested by reviewers below:

* We added more discussion around prior theoretical results to the Discussion section (reviewer eSj4)

* Included further discussion around non-linear features (reviewer J8rP)

* Qualified the claim about whether we could derive model-agnostic frequency thresholds for when linear representations form (multiple places). We agree we do not have enough data to support the strong version of this claim. (reviewer juu2)

* Updating Table 1 to be more descriptive of the setup where we test cross-model generalization of the regression model we fit (juu2). We updated the caption to this table as well to reflect this

* Updating the caption in Figure 2 to be more descriptive of which features we are discussing

* Clarifying which numbers we are comparing when discussing generalization performance of the regression (L482-483)

We received advice for making Figure 2 more readable from reviewer juu2. We believe this is very helpful feedback and will *definitely* make the changes in the final revision, but we won't be able to make these changes before the rebuttal deadline. Note however, that this is purely presentational, and the relevant data is presented in the current draft.

Once again thank you for a very productive round of reviews
</content>
</thread_object>

</thread>


Metareview:
metareview:
his paper explores the Linear Representation Hypothesis in LLMs. Specifically, it aims to demonstrate a correlation between the cooccurrences of a subject and object and their relationship being represented in a linear way.  The paper establishes thresholds in different models beyond which this happens. It then looks to go the other way and predict pre-training frequencies given representations. Adding LRE features improves performance here.

The core question of this paper is very timely and interesting. The paper is clearly written and presents what I would characterize as medium-strong evidence of the correlation between frequency and linearity. The types of representations formed by LLMs and what factors of training cause them to emerge is a very relevant question. This paper advances the state-of-the-art in our understanding of these points.

The biggest issue with this paper is the scope and impact of the results. juu2 brings up the transfer to another LLM; the results in Table 1 aren't all that strong.  The restriction to a particular set of relations (and to KG relations in general) is also somewhat limiting. I don't think there's one right answer for how much the paper needs to engage with on this front, but making stronger and more general claims would of course make it stronger. As a more minor point, juu2 points out that connections of the paper's hypotheses with model size can't necessarily be drawn from the given data.

Taken together, all of this contributes to an impression that this paper has some important and useful results, but it might not be the last word on this topic. A scaled-up set of experiments, scaled out to different settings, may find something new and different here.

Finally, the paper explores the same question in a few different ways. For instance, the ability to predict the frequency is mostly a consequence of correlation (juu2), which lessens its impact a bit.