Physics of Language Models: Part 3.1, Knowledge Storage and Extraction — part 3/6

Cleaned-text code points [31227, 47023). Bibliography omitted here; source.md includes it.

• Fullname . We denote the augmentation where all pronouns or partial names in bioS / bioR \textsf{bioS}/\textsf{bioR} are replaced with the person’s full name as fullname . 14 14 14 In the synthetic bioS dataset, a person’s full name is presented only once, at the start of the initial sentence, with subsequent sentences using solely pronouns. For the LLaMa-generated bioR data, typically, the person’s full name appears once at the start; later sentences use either pronouns or parts of the name, such as the first or last name. An example of this augmentation is: Anya Briar Forger originated from Princeton, NJ . Anya Briar Forger dedicated her studies to Communications . Anya Briar Forger gained work experience in Menlo Park, CA . Anya Briar Forger developed her career at Meta Platforms . Anya Briar Forger came into this world on October 2, 1996 . Anya Briar Forger pursued advanced coursework at Massachusetts Institute of Technology .

Results. In Figure 3 , we present our results for the GPT2 model on the bioS dataset. We implemented each knowledge augmentation individually and in combinations, then compared the model’s QA finetune accuracy on 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} . The model architecture and training parameters remained the same, but the pre-training datasets varied based on the applied augmentations. Further experiment details are in Appendix D ; and additional results for the bioR dataset and/or Llama model are in Figure 12 . We find that: Result 3 ( Figure 3 ) . Adding multiplicity, permutations, or repeating full names, all help the model to better store knowledge during pretraining, making knowledge extraction easier later. Notably: 15 15 15 We have also tried to translate from English to French, which boosts accuracy to about 40% but we did not include the result for clarity. An exception is when permutation is directly added to the single data without multiplicity (see “ bioS single + permute1 ”), this hurts the QA performance as it makes knowledge extraction harder. • Pretraining on a dataset where each person has 5 diverse biography entries (i.e., different wording, sentence shuffling) boosts the QA fine-tune accuracy (on 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} ) from 9.7% to 96.6%. • More augmentation ⇒ \Rightarrow better: gain increases as multiplicity or permutation number increases.

One may infer from Result 3 that exposing the model to varied expressions of the same knowledge encourages it to focus on the underlying structure of the knowledge, rather than its word-by-word presentation. We shall verify this hypothesis in Section 5 by introducing probing techniques.

## 5 Results 4-5: Knowledge Probes on the BIO Pretrained Model

We investigate how a language model, pretrained on BIO data , encodes knowledge in its hidden states. We propose two probing methods: position-based (P-probing) and query-based (Q-probing). Both methods employ simple, nearly-linear probes to extract personal attributes from the model’s hidden states.

### 5.1 Result 4: Position-Based Probing

In P-probing, we feed biography entries into a pretrained model, and finetune an additional linear classifier on the model’s final hidden layer to predict the six target attributes (e.g., university, major, etc.). We wish to understand how and where these attributes are encoded after pretraining.

To accommodate varied data lengths, we identify six special token positions immediately preceding the first occurrences of the six attributes in each biography entry (see Figure 4 ). This results in 6 × 6 6\times 6 classification tasks. For each prediction task, we freeze the entire pretrained network but add a trainable rank-2 update on the embedding layer to accommodate the task change. We use the transformer’s last hidden layer at these positions to (linearly) predict the six target attributes. 16 16 16 For GPT2-small with 768 hidden dimensions and vocab size V V , this rank-2 update has 2 ​ V + 2 × 768 2V+2\times 768 trainable parameters. The linear classifier layer is of dimension 768 × M 768\times M for each target attribute with M M possibilities. More details can be found in Appendix E .

We are particularly interested in how early the attributes are encoded in a biography. For instance, if the linear classifier to predict “company name” shows high accuracy right after the person’s full name, it implies that the model is directly learning “Anya’s employer is Meta Platforms”. If high accuracy is only achieved at the biography’s end, the model might be using a flawed logic , such as “the birthday is October 2, 1996, the university is MIT, hence the employer is Meta.”

P-Probing Main Results. Our results are in Figure 5 and summarized as follows. • In the bioS single setup, P-probing accuracy remains low (e.g., 2% for company name) until the token immediately preceding the target attribute (where accuracy boosts to 100%). This suggests that the model memorizes all the BIO data during pretraining, but encodes knowledge using the “flawed logic” above. This prevents knowledge extraction during QA finetuning, especially when only the person’s name is provided.

• In the heavily augmented setup like bioS multi5+permute , the P-probing accuracy for all six attributes rises to nearly 100% from the first special position, which is before all of the attributes. This indicates that the model not only memorizes the BIO data but also identifies the person’s complete six attributes solely upon seeing the person’s name, facilitating knowledge extraction during the QA finetuning process.

• For intermediate setups, the results are mixed. For example, comparing bioS single with multi5 , we see that adding multiplicity (without permutation) results in earlier attribute storage, accounting for the increase in QA finetune accuracy from 9.7% to 41% as seen in Figure 3 . Comparing bioS single+permute1 with single+permute5 , we observe that permuting the six sentences five times (without diversifying the sentences) also leads to earlier knowledge storage, explaining the rise in QA finetune accuracy from 4.4% to 70%.

In sum,

Result 4 ( Figure 5 ) . Increased knowledge augmentation in the pretrain data improves P-probing accuracies at earlier token positions . Consequently, a key-value pair knowledge (e.g., person-employer) more directly associates the value with the key rather than with other related attributes. This mechanism facilitates the (out-of-distribution) extraction of knowledge through fine-tuning. In Section 5 , we use a Venn diagram to clearly demonstrate which attribute is stored upon observing another, further supporting this finding.

#### 5.1.1 Closer P-Probing at Knowledge Dependency

As noted above, the model may infer attribute relationships based on their order in the pretrain data. For instance, if a birth date always precedes a company city, the model might infer “the person born on October 2, 1996 works in Menlo Park” instead of “Anya’s work city is Menlo Park”. This can occur if the pretrain data isn’t adequately augmented, and the model may even favor linking one attribute to another, rather than to the person’s name, if two attributes are closely correlated (such as company city and company name).

To investigate this, we created a variant of the bioS dataset, grouping the 6 sentences into 3 pairs with a consistent order: birthdate before birth city, university before major, and work company before work city. We allowed random permutations among these pairs and sentence diversities. We refer to this dataset as bioS couple (see Appendix A for details).

We examined our P-probing on this dataset as 2 5 × 6 2^{5}\times 6 classification tasks, predicting each of the six target attributes from a special token position where only a subset S S of the remaining five attributes has been observed ( S S has 2 5 2^{5} possibilities). 17 17 17 The P-probing process remains the same as before, using only 6 sets of trainable parameters each for a target attribute, each with a single classification linear layer and a single rank-2 update on the embedding. The difference is a more detailed interpretation of the results. Our results, visualized in Figure 6 , show that the accuracy in predicting the second attribute in each pair is heavily influenced by whether the model has encountered the first attribute, even with moderate data diversity.

###### Remark 5.1 .

This relates to Figure 3 , where the “company city” attribute shows the weakest QA finetune performance in the bioS data family. This is due to our data construction, where ‘‘company city’’ is determined solely by ‘‘company name”. The model thus associates ‘‘company city’’ with ‘‘company name’’ rather than the person’s name, if the company name is presented earlier. 18 18 18 This dynamic can be explored as a form of knowledge manipulation . For example, if the language model is good at retaining work company names, can it determine work city locations as a simple classification task using just company names? We explore this in our parallel paper [ 2 ] .

#### 5.1.2 P-Probing Extensions

We could consider alternative P-probing forms, such as introducing a low-rank update to the pretrained model’s main body, like a trainable LoRA update with a small rank on the query/value matrices. While not necessary for our positive results (e.g., the highly augmented data bioS multi5+permute ), it could be interesting to apply this to the negative results (e.g., the basic data bioS single ). However, our experiments showed no significant increase in P-probing accuracies, so we omit the details.

Our P-probing has focused on the six distinct token positions, likely the preposition words preceding the six attributes. How about probing other positions, like tokens following each attribute or the person’s name? We observed that P-probing accuracy might improve as the model processes more “extraneous” tokens. For instance, the P-probing accuracy for a person’s birth date could increase after encountering phrases like “was born on” or “has birthday in”. This could be due to the model’s ability to associate the birthdate information with the sentence’s structure . We chose not to include these observations for clarity.

In Appendix E , we demonstrate the difference between a rank-2 and a rank-4 update on the embedding layer. It confirms that a rank-2 update is sufficient for P-probing on our datasets.

### 5.2 Result 5: Query-Based Probing

P-probing offers a qualitative assessment of early knowledge storage in the model relative to the original biography entry. However, it can be limiting due to its dependence on the exact context structure from the biography entry. For instance, in Figure 4 , knowledge may be stored in short phrases like “received mentorship and guidance.”

In query-based probing (Q-probing), we aim for a more precise, context-free value from a pretrained model, focusing on the knowledge directly associated with a person’s name. We evaluate sentences containing only the person’s full name and train a linear classifier on the last layer’s hidden states to predict the person’s six attributes. High accuracy suggests that the model directly links each person’s attributes to their name.

We consider an input sentence containing only the person’s full name, preceded by a starting token and followed by an ending token. Like P-probing, we freeze all transformer layers (acquired through pretraining), except the embedding layer, where we apply a low-rank update (using rank 16, compared to rank 2 in P-probing). This minimal change is necessary as we are addressing a distinct classification task under a different input distribution. We extract the hidden states from the last layer on the ending token and place a trainable linear classifier on top to predict the person’s six attributes. More details are in Appendix F .

Our findings. Our results are in Figure 7 . Our main finding is: Result 5 ( Figure 7 ) . The QA finetune accuracy correlates closely with Q-probing accuracy, indicating that the “degree to which the attribute is directly linked to the person’s name” is a crucial factor for effective knowledge extraction . If the model fails to store knowledge in this way during pretraining, QA finetuning may not rectify this, regardless of the prompts or finetune parameters. Note, once again, applying knowledge augmentations to the pretrain data, Q-probing accuracy significantly increases. This suggests that the model encodes knowledge almost linearly in the hidden states directly adjacent to the person’s name. Thus, the linear probes can extract the person’s attributes from these hidden states as effectively as the model can be adapted through QA finetuning to answer questions related to those attributes.

Our result also suggests that, at the last hidden-layer, the model neither uses complex or nonlinear transformations nor leverages interactions between hidden states at different token positions to extract knowledge about the person. This implies that the model does not use contextual or global information from the biographies to extract knowledge about the individual .

## 6 Result 6: Celebrity Can Help Minority

Section 4 highlighted the significant benefits of knowledge augmentation. However, in practice, we may not have augmented data for all individuals. This section explores whether partially augmenting data can improve knowledge extraction for non-augmented data. In our biography dataset, the augmented subset is akin to a “celebrity” group with plentiful online biographical information, potentially included in the fine-tuning dataset as well. The non-augmented subset is comparable to a “minority” group with limited biographical data.

For comparison, we introduce an additional set of N = 100,000 N=100,000 individuals, the celebrity group 𝒫 𝖼𝖾𝗅 {\mathcal{P}_{\mathsf{cel}}} , while the original N N individuals form the minority group 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{min}}} . We test both synthetic bioS and more realistic bioR data. For bioS , the celebrity group’s biographies use the multi5+permute augmentation, simulating varied expressions found on internet. For bioR , the celebrity group uses the multi5 augmentation, generating their biographies five times using Llama.

The language model is pretrained on the combined set 𝒫 𝖼𝖾𝗅 ∪ 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{cel}}}\cup{\mathcal{P}_{\mathsf{min}}} biographies and then fine-tuned using QAs from the celebrity group 𝒫 𝖼𝖾𝗅 {\mathcal{P}_{\mathsf{cel}}} . We evaluate the model’s QA accuracy on the 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{min}}} group. 19 19 19 Other fine-tuning variations, such as QA fine-tuning with half of 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{min}}} as training and half as testing, show negligible differences. Our results are presented in Figure 8 .

Result 6 ( Figure 8 ) . Introducing celebrity data boosts the minority group’s QA accuracy (e.g., from 4.4% to 86.8% for the bioS data). This is significant because: - the minority group’s BIO pretrain data remains unchanged in both cases, and - the minority group’s QA data is not used during fine-tuning.

This highlights that merely including celebrity data during pretraining significantly improves the model’s ability to store and extract knowledge from the minority group. Similarly, in the more realistic bioR case, introducing celebrity data increases the minority’s accuracy from 10.0% to 76.3%. This strongly suggests that this phenomenon also occurs in real-world scenarios .

We also use P-probing and Q-probing techniques to validate and explain the above findings; they suggest that with the inclusion of celebrity data, the attributes of the minority group are more directly stored onto their names. These are detailed in Figure 18 and Figure 19 in Appendix G .

###### Remark 6.1 .