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

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

We generate a six-sentence biographical text entry for each individual, highlighting six distinct aspects. For diversity, each sentence is randomly chosen from approximately 50 distinct templates. In the basic configuration, we generate a single biographical entry for each person, maintaining a consistent order for the six sentences. We use “ bioS single ” to denote this basic configuration. See an example entry below: Anya Briar Forger was born on October 2, 1996 . She spent her early years in Princeton, NJ . She received mentorship and guidance from faculty members at Massachusetts Institute of Technology . She completed her education with a focus on Communications . She had a professional role at Meta Platforms . She was employed in Menlo Park, CA . (2.1) We also explore 3 types of knowledge augmentations: (1) multi M M , generating M M biography entries for an individual using varied templates, (2) fullname , substituting he/she/they with the person’s full name; and (3) permute , shuffling the six sentences randomly. Examples are given in Section 4.2 .

BIO dataset bioR . We examine a “close-to-real” dataset produced by Llama [ 37 , 40 ] . For the set of N = 100,000 N=100,000 individuals, we provide an instructive prompt to Llama to generate a biographical entry. Here’s an example: Anya Briar Forger is a renowned social media strategist and community manager. She is currently working as a Marketing Manager at Meta Platforms . She completed her graduation from MIT with a degree in Communications . She was born on 2nd October 1996 in Princeton, NJ and was brought up in the same city. She later moved to Menlo Park in California to be a part of Facebook’s team. She is an avid reader and loves traveling.

We diversified our instructive prompts by drawing from a pool of templates and employed rejection sampling to guarantee the inclusion of all six attributes. In the basic configuration, we produce a single biographical entry for each person (denoted as “ bioR single ”). For comparison, we also consider multi M M augmentation which generates M M entries per person and the fullname augmentation. Additional examples can be found in Appendix A .

QA dataset. This paper explores the effectiveness of a trained language model in retaining knowledge from BIO data. As discussed in the introduction, memorization is more than just predicting the next token when given exact sentences from BIO. It includes the model’s ability to truly extract knowledge from the BIO . We assess this knowledge extraction using a question and answer (QA) framework. For each individual, we pose six questions targeting their six unique attributes:

1. What is the birth date of Anya Briar Forger? Answer: October 2, 1996.

2. What is the birth city of Anya Briar Forger? Answer: Princeton, NJ.

3. Which university did Anya Briar Forger study? Answer: Massachusetts Institute of Technology.

4. What major did Anya Briar Forger study? Answer: Communications.

5. Which company did Anya Briar Forger work for? Answer: Meta Platforms.

6. Where did Anya Briar Forger work? Answer: Menlo Park, CA.

For each question, we use it as a prompt for the model to generate a response. QA accuracy is measured by the proportion of answers that exactly match the correct response. 6 6 6 We disregard partial matches or synonyms, emphasizing the model’s precision in knowledge extraction.

### 2.1 Training Details

Model architectures. We adopt the GPT2/Llama architectures [ 31 , 37 ] , where for GPT2 we replace its absolute positional embedding with rotary positional embedding [ 34 , 8 ] , but still referring it as GPT2 for short. 7 7 7 A controlled experiment to highlight the importance of rotary embedding is in [ 1 ] . Since this paper appeared, Jiang et al. [19] confirms our results also apply to the pretrained Llama-7B model; our own follow-up also tried the Mistral architecture [ 3 ] . Recall the GPT2-small architecture comprises 12 layers with 12 heads and 768 dimensions [ 31 ] . We use a 12-layer, 12-head, 768-dim GPT2 (124M) or Llama architecture for pre-training on the bioS data, but a larger 12-layer, 20-head, 1280-dim GPT2 (302M) or Llama for the bioR data to accommodate its increased complexity. Only in Figure 2 when presenting a negative result, we tried a 12-layer 32-head 2048-dim GPT2 (682M). The default GPT2/Llama tokenizers are used, which convert simple words into single tokens, but names and most other attributes into tokens of varying lengths. When it comes to Section 7 , we also use a BERT architecture [ 20 ] .

Training. We investigate two types of autoregressive training, detailed in Appendix B .

Pretrain + instruction finetune . Here, we pre-train the language model from scratch on the BIO data, randomly sampling and concatenating them into 512-token sentences, separated by a standard <EOS> token. The model is then fine-tuned using half of the QA data and evaluated on the remaining half, mirroring the typical instruction finetune process.

Mixed training . In mixed training, we train the model from scratch on all BIO data and half of the QA data. BIO and QA entries are randomly sampled without requiring them to be from the same individual. We use a parameter 𝖰𝖠 r {\mathsf{QA}_{r}} to control the QA data amount, primarily setting 𝖰𝖠 r = 0.8 {\mathsf{QA}_{r}}=0.8 (a 2 : 8 2:8 BIO to QA entry ratio). The model’s generation accuracy is evaluated using the remaining QA data. 8 8 8 See Appendix C for a comparison of how 𝖰𝖠 r {\mathsf{QA}_{r}} affects performance. We used beam =4 without sampling throughout this paper; results are similar if disabling beam.

LoRA + full finetune. In full finetuning a pretrained model is tuned for a downstream task such as QAs. LoRA finetuning [ 18 ] improves upon this by freezing all pretrained model parameters and adding low-rank updates to a subset of the weight matrices for fine-tuning. We apply a low-rank update to the query/value matrices of the transformer model (suggested by [ 18 ] ) and the embedding layer to account for input data distribution shifts. Full finetuning is also included when presenting negative results.

## 3 Result 1: Mixed Training Enables Knowledge Extraction

Mixed training involves using BIO data for all individuals together with QAs for half of them. The group of individuals whose QAs are included in the training set is referred to as in-distribution or 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} . The model’s generative accuracy is then tested on the QAs from the remaining individuals ( 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} ) to assess its out-of-distribution (OOD) generalization capability.

Result 1 ( Figure 1 ) . A model mixed-trained on both knowledge and its extraction QA tasks can effectively learn to extract knowledge. • As shown in Figure 1(a) , the OOD generalization accuracies are 86.6 % 86.6\% when mixed-trained on bioS single and 77.7 % 77.7\% for bioR single . • However, the model achieves this through somewhat abnormal behavior akin to “studying to pass the test,” discussed further in Section 3.1 . (We emphasize that the accuracy is OOD: extracting an individual’s attributes even when no QA about that person — and only the BIO of that person — was seen in the training data.)

### 3.1 Model’s Abnormal Learning Behavior

We examine the model’s mixed training for knowledge storage and extraction by monitoring its accuracies on the BIO/QA data and for 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} / 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} separately. Specifically, 9 9 9 Interested readers may consider “whole-attribute” accuracies instead of “first-token” accuracies. They are similar, so we omit them here. - BIO first-token accuracy: we track the model’s next-token-prediction accuracy on the first token of each of the six attributes (birthdate, birthcity, etc.) in the BIO data, separately for 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} / 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} . This measures the model’s BIO data memorization performance. (Despite all individuals’ BIO data appearing in training, we still separately track them for 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} / 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} .)

- QA first-token accuracy: we track the model’s next-token-prediction accuracy on the first answer token in the QA data, separately for 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} / 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} . This loosely estimates the model’s QA generation performance.

- QA generation accuracy: we track the model’s whole-attribute generation accuracy on 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} .

From Figure 1(b) and 1(c) , we find that the model employs an unconventional learning strategy. • Initially, the model uses the QA data from the training set to encode knowledge for people in 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} , as indicated by the rapid increase in QA in-dist accuracy. This also aids in memorizing in-dist BIO data, as shown by the subsequent rise of the BIO in-dist accuracy.

• The model then gradually aligns the encoded knowledge with the BIO data to learn to extract knowledge and generalize it to 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} . Notably, it takes a while before the BIO out-dist accuracy catches up, followed by an increase in the QA out-dist accuracy.

This is akin to the “study to pass the test” approach in schools, where students prepare using past exam questions and textbooks for answers. While this may yield high scores, it doesn’t reflect the natural progression of human knowledge acquisition. To address this , we explore a more challenging scenario in the next section where the model is pretrained on the BIO data without exposure to the questions.

###### Remark 3.1 .

In mixed training, we selected 𝖰𝖠 r = 0.8 {\mathsf{QA}_{r}}=0.8 , maintaining a 8 : 2 8:2 QA to BIO ratio as outlined in Section 2 . We found that a higher QA ratio during training improved out-of-distribution QA accuracy ( Figure 10 in Appendix C ), further supporting our observation of the model’s abnormal behavior: it first learns knowledge from QA and then associates it with BIO. For comparison, Llama was trained using only 2% of tokens from StackExchange [ 37 ] .

## 4 Result 2-3: BIO Pretrain + QA Instruction Finetune

We explore a scenario where the model is pre-trained exclusively on the BIO data of all individuals, followed by fine-tuning using QAs from half of these individuals, denoted as 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} . The model’s OOD generalization is then assessed on questions related to the other half, denoted as 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} , whose BIO/QA data were not involved in the fine-tuning. This setup simulates the process of applying learned knowledge from textbooks to solve exam questions.

### 4.1 Result 2: Model Fails to Extract Knowledge After BIO Pretrain

We first pretrain on bioS or bioR single , each containing a single biography per person. The QA finetune generalization accuracies (on 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} ) are shown in Figure 2 , using both full and LoRA fine-tuning [ 18 ] . The model’s QA finetune training accuracy on 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} is also presented for comparison.

Despite a 99+% first-token accuracy during pretraining, the model exhibits zero-zero QA accuracy on 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} for all finetuning parameters. This indicates that while the model can memorize BIO data token-by-token, it struggles to extract the underlying knowledge. Full-finetuning yields near-perfect in-distribution QA accuracy on 𝒫 𝗍𝗋𝖺𝗂𝗇 {\mathcal{P}_{\mathsf{train}}} , showing it can memorize QAs for individuals in the fine-tuning set. However, it fails to generalize to QAs about individuals in 𝒫 𝗍𝖾𝗌𝗍 {\mathcal{P}_{\mathsf{test}}} . In sum: Result 2 ( Figure 2 ) . A model pretrained to word-by-word memorize knowledge may never be fine-tuned to extract knowledge. As shown in Figure 2 : perfect BIO token memorization + perfect QA answers for half the people ⟹̸ \displaystyle\not\Longrightarrow correct QA answers for the other half. ( knowledge extraction does not come for free ) This holds true even when the model size is ∼ \sim 7000x larger than N = 100 ​ k N=100k , with each individual exposed 1350 times during pretraining, and numerous finetuning parameters have been explored. 10 10 10 In our follow-up work [ 3 ] , we increase the model size to 1B and N N to 20M, confirming similar results. Despite memorizing all knowledge from the BIO data during pretraining, the model encodes it in a disorganized manner within the transformer, preventing knowledge extraction during fine-tuning. 11 11 11 This is not a direct result of catastrophic forgetting, a common issue during heavy fine-tuning where the model forgets the pretraining data. Even with LoRA fine-tuning, which introduces minimal low-rank updates to model weights while preserving the pretrained model, test accuracy only slightly improves.

Figure 2 seems to contradict the success of large models like GPT3.5/4, trained on diverse internet data such as Common Crawl and known for effective knowledge extraction upon fine-tuning. Analyzing the test accuracy breakdown for the six attributes on the bioS data ( Figure 3 , the “ bioS single ” row), we find that QA fine-tuning achieves a 33 % 33\% generalization accuracy on the “birthdate” attribute but performs poorly on others. This is because our bioS single data consistently places birthdate as the first attribute after a person’s name, unlike internet data which presents information variably, often repeating it with diverse wordings and orderings. The next subsection on knowledge augmentation supports this hypothesis.

### 4.2 Result 3: Knowledge Augmentation

We explore how knowledge augmentation enhances a model’s capacity to store and efficiently extract knowledge from training data. We focus on three augmentations: adding multiplicity, introducing permutations, and repeating full names, typically found in internet data. The original datasets without augmentation are referred to as bioS single and bioR single .

• Multiplicity . We denote the method of creating M M distinct biography entries for each individual, using varied language but retaining the same information, as multi M M . 12 12 12 For bioS data, each of the six sentences is selected from around 50 50 templates, with a new template resampled for each sentence in the M M entries. For bioR data, we recreate the biography using Llama for each of the M M entries. An example of adding multiplicity to the biography in ( 2.1 ) is: Anya Briar Forger came into this world on October 2, 1996 . She originated from Princeton, NJ . She pursued advanced coursework at Massachusetts Institute of Technology . She dedicated her studies to Communications . She developed her career at Meta Platforms . She gained work experience in Menlo Park, CA .

###### Remark .

As a special case, we also experimented with translation (e.g., English to French) to increase sentence diversity, which proved beneficial for the model’s knowledge extraction, but we have not included these details in this paper for clarity.

• Permutation . We denote adding random permutations to the biography sentences as permute . 13 13 13 For bioS single , we denote random permutation of the same six sentences P P times as permute P P . For bioS multi M M , we denote random permutation of each of the M M biography entries as permute . The bioR data, generated by Llama, already has some randomness in sentence ordering, so no extra permutations are added. For instance, the example above can be permuted as follows: Anya Briar Forger originated from Princeton, NJ . She dedicated her studies to Communications . She gained work experience in Menlo Park, CA . She developed her career at Meta Platforms . She came into this world on October 2, 1996 . She pursued advanced coursework at Massachusetts Institute of Technology .