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

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

# Physics of Language Models: Part 3.1, Knowledge Storage and Extraction Thanks: Project page: https://physics.allen-zhu.com/part-3-knowledge/part-3-1 . An extended video of this paper is available at https://youtu.be/YSHzKmEianc . V1 was circulated internally at Meta on Sep 18, 2023, and appeared on arXiv on Sep 25, 2023. V2 is nearly identical to V1, with minor corrections to author names and writing. V3 includes additional Llama experiments and further writing improvements. We would like to thank Lin Xiao, Chunting Zhou, Tianyi Peng, Xiaodong Liu, and Zhijie Zhou for many helpful conversations. We would like to extend special thanks to Nabib Ahmed, Giri Anantharaman, Lucca Bertoncini, Henry Estela, Liao Hu, Caleb Ho, Wil Johnson, Apostolos Kokolis, and Shubho Sengupta from Meta FAIR, as well as Ian Clark, Gourab De, Anmol Mann, and Max Pfeifer from W&B ; without their invaluable support, the experiments in this paper would not have been possible.

###### Abstract

Large language models (LLMs) can store a vast amount of world knowledge, often extractable via question-answering (e.g., “What is Abraham Lincoln’s birthday?”). However, do they answer such questions based on exposure to similar questions during training (i.e., cheating), or by genuinely learning to extract knowledge from sources like Wikipedia?

In this paper, we investigate this issue using a controlled biography dataset. We find a strong correlation between the model’s ability to extract knowledge and various diversity measures of the training data. Essentially , for knowledge to be reliably extracted, it must be sufficiently augmented (e.g., through paraphrasing, sentence shuffling, translations) during pretraining . Without such augmentation, knowledge may be memorized but not extractable, leading to 0% accuracy, regardless of subsequent instruction fine-tuning.

To understand why this occurs, we employ (nearly) linear probing to demonstrate a strong connection between the observed correlation and how the model internally encodes knowledge — whether it is linearly encoded in the hidden embeddings of entity names or distributed across other token embeddings in the training text.

This paper provides several key recommendations for LLM pretraining in the industry: (1) rewrite the pretraining data — using small, auxiliary models — to provide knowledge augmentation, and (2) incorporate more instruction-finetuning data into the pretraining stage before it becomes too late.

## 1 Introduction

Knowledge is crucial for human cognition and communication, allowing us to comprehend and utilize information. For humans, this often involves memorization, the process of storing and retrieving information in the brain. For example, after reading a biography of Abraham Lincoln, we can memorize the information and later answer questions like “Where was Lincoln born?” or “What is Lincoln’s birthday?” Memorization enables us to extract and manipulate knowledge from the sentences we read or hear, recognize the entities, relations, and facts expressed in the text, and apply logical and causal reasoning to infer new information or answer queries [ 42 , 12 , 6 , 4 ] .

In this paper, we explore how transformer-based language models memorize knowledge during training and extract it during inference. This is distinct from in-context learning or RAG [ 22 ] , where the model is given a paragraph during inference and immediately answers questions about it. We focus on factual knowledge (e.g., knowledge graph) that a language model needs to memorize from the training corpus, encode in its weights, and extract later during inference.

We stress that memorizing all sentences in the training data does not ensure that the model can extract or manipulate the factual knowledge from the sentences during inference. Language models can reproduce the exact input during inference, but this doesn’t necessarily mean they can use these sentences to answer factual questions related to them. Hence, we differentiate between “memorization of knowledge” in language models and traditional memorization in machine learning, which merely means the model can fit the exact training data, but doesn’t imply the model can extract the knowledge flexibly from the data after training.

For example, if the training data includes Lincoln’s biography, the model can memorize and reproduce the sentence “Abraham Lincoln was born in Hodgenville, K.Y.” when given the prompt “Abraham Lincoln was born in”, but it might not be able to answer the question “Which city was Abraham Lincoln born in?” Therefore, a key question is: How do language models memorize knowledge during training, and extract it later to answer questions or perform logical reasoning during inference?

Previous works have demonstrated that language models can “memorize” a lot of knowledge by probing the model to answer questions related to different entities and attributes, see [ 35 , 33 , 28 ] and the citations therein. These studies use models pretrained on internet data, leaving it unclear whether the model answers questions like “Which city was Abraham Lincoln born in?” by extracting knowledge from Lincoln’s biography ( our focus ) or if it encountered a similar (or same!) question during training and memorized the answer (traditional memorization / data contamination ).

Given the challenges of conducting controlled experiments with internet data, we propose studying this question using well-controlled, synthetically generated data, 1 1 1 One could suggest filtering the data to eliminate such questions and retraining the model. However, this doesn’t rule out the presence of similar sentences “Which city did Abraham Lincoln grow up in?”, more complex ones in French, or grammatically incorrect versions like “Where Abraham Lincoln birth in?” in the data. examining the models’ mathematical properties that characterize their knowledge representation and extraction. We construct a synthetic dataset of 100 ​ k 100k biographies, including their birthday, birth city, major of study, etc. We also use Llama [ 37 ] to rewrite them to make them close to real-life biography styles. We pretrain the language model on the biography dataset of all the 100 ​ k 100k people. We ask: 2 2 2 We leave the follow-up question to study logical reasoning or manipulation on knowledge to a separate paper [ 2 ] . After pretraining a language model on the biography dataset, can the model be finetuned to answer questions like “Where is the birth city of [name ]”, and if so, how does the model achieve so?

After pretraining the model on the entire biography, we fine-tune it using question and answer (QA) pairs from a p p fraction of individuals. We then test its ability to out-of-distribution answer QAs about the remaining 1 − p 1-p fraction. This approach ensures that the model (1) is exposed to sufficient data to comprehend the QAs and (2) does not encounter the same questions during training. The paper is structured as follows: 3 3 3 Our result numbers correspond to our online video of the paper, available at https://youtu.be/YSHzKmEianc .

• Result 1 : Mixed training ⟹ \Longrightarrow knowledge extraction.

Before diving into the pretrain-finetune process, we first demonstrate that pretraining a model on all biographies plus QAs for a p p fraction of individuals together enables it to (apply knowledge to) answer questions about the remaining 1 − p 1-p fraction. We call this process mixed training . We observe in mixed training, the model first uses QAs to encode knowledge about the p p fraction, then correlates this encoded knowledge with the biography to infer generalization to the remaining 1 − p 1-p fraction. This learning process deviates from typical human learning 4 4 4 For humans, arguably, we first learn from textbooks and then answer exam questions. and is less frequently used in practical LLM pretrain (and perhaps it should!).

• Result 2 - 3 : Instruct finetune ⇏ \nRightarrow knowledge extraction (unless data augmented). Consider a model pretrained only on the biographies and then finetuned using QAs for a p p fraction of individuals. We discover that it struggles to answer questions about the remaining 1 − p 1-p fraction, irrespective of model size, pre-train time, or finetune parameters ( Result 2 ). However, accuracy significantly improves with knowledge augmentations like varying writing styles or sentence shuffling ( Result 3 ). This gives a strong link between knowledge augmentation in the pretrain data and the model’s knowledge extraction ability after finetuning.

• Result 4 - 5 : Introduce probing techniques to explain Why this happens . As another main contribution, we introduce (nearly) linear probing techniques to show that knowledge augmentation pushes the model to encode a person’s knowledge almost linearly in the model’s hidden embedding of the person’s name tokens. Without augmentation, the model encodes the person’s knowledge across all biography words/tokens, making knowledge extraction nearly impossible no matter how one finetunes it. In sum: no knowledge augmentation in pretrain data ⟺ attribute is not entirely stored on person’s names \displaystyle\Longleftrightarrow\text{\small attribute is {not} entirely stored on person's names } when the model memorizes the pretrain data ⟺ knowledge cannot be extracted via instruction finetune \displaystyle\Longleftrightarrow\text{\small knowledge cannot be extracted via instruction finetune } knowledge augmented in pretrain data ⟺ attribute is nearly entirely stored on person’s names \displaystyle\Longleftrightarrow\text{\small attribute is {nearly} entirely stored on person's names } ⟺ knowledge can be extracted via instruction finetune \displaystyle\Longleftrightarrow\text{\small knowledge can be extracted via instruction finetune }

• Result 6 : Knowledge augmentation on the “celebrity” helps “minority”. Even if knowledge augmentation is applied to a subset of individuals, what we call celebrities, test accuracy for others (without augmentation) also increases significantly. We discover that the mere inclusion of celebrity data (e.g., people with plentiful online biographical data of diverse writing styles) in pre-training enhances the model’s knowledge extraction for minorities.

• Result 7 : Bi-directional models fail to extract knowledge. We show that encoder-only models akin to BERT , whether mixed-trained or pre-trained and then fine-tuned, cannot extract a person’s knowledge after finetuning, regardless of the knowledge augmentation, unless the knowledge is a single word or multiple but independent words (like birth month, day, and year).

Practical Implications. Our controlled study offers key recommendations for LLM training at an industrial scale: • We emphasize the importance of pre-training data rewriting (augmentation) , particularly for rare but critical data. Addressing this during fine-tuning is often too late. Without rewriting, a model may accurately recite knowledge data word by word, but the way it embeds this knowledge into its weights may impede retrieval when prompted differently, resulting in a total waste of model capacity .

Tools such as Llama-7B or even smaller auxiliary models are adequate for this rewriting task. These “rewrite models” do not need to possess the knowledge themselves. As demonstrated, simple sentence-level shuffling or English-to-French translations can already enhance performance. Generally, we suggest including prompts that encourage sentence shuffling when using such rewrite models.

Data rewriting is a form of data augmentation, but also distinct from traditional methods (e.g., dropout, masking, cropping, jittering, flipping) and their associated distillation techniques (like contrastive learning). While traditional augmentations promote the learning of generalizable features over pure memorization, data rewriting — what we call knowledge augmentation — helps language models to memorize knowledge in a more accessible format for downstream tasks. Without such augmentation, the accuracy even for the simplest knowledge extraction task, could be near zero.

• We also demonstrate the advantages of including more instruction-finetuned data during pre-training . Our mixed training experiments show that postponing all QA-like data to the fine-tuning phase is suboptimal. Introducing QA-like data earlier in pre-training enables the model to encode knowledge more effectively .

### 1.1 Related Work

Linear probing of knowledge . Linear probing is a recognized method to examine how a model encodes knowledge [ 35 , 23 , 11 , 5 , 13 , 26 , 15 ] . Contrary to previous studies that suggest models trained on internet data can linearly encode knowledge in the hidden embeddings of entity names, we find that such encoding is only possible with knowledge augmentations like permutation/rewriting of entity-attribute knowledge during pretraining. Without these augmentations, the language model can still memorize the training data, but it is not linearly encoded in the entity’s hidden embeddings, making knowledge extraction via QAs quite hard, if not impossible, even with instruction fine-tuning. This implies that diverse internet data on the same entity is vital for pre-training the language model for knowledge extraction during inference. The usefulness of augmentations of pretraining data for language models was also empirically observed in literature [ 9 , 14 , 7 , 21 ] , but they did not explore where the knowledge is nearly-linearly encoded in a sentence and its correlation with knowledge augmentation, a process we refer to as P-probing in Section 5.1 .

Probing language models’ knowledge via QAs . Question answering (QA) is a common method to probe the knowledge encoded in language models pretrained on internet data [ 35 , 33 , 28 , 17 , 32 , 29 , 30 , 27 ] . However, it’s unclear whether these models answer questions by extracting knowledge from the training source or by recognizing exact/similar questions from training. We use controlled experiment for out-of-distribution testing on individuals whose QAs were not part of training. This approach also allows us to study the correlation between knowledge extraction and the diversity of pretrain data.

Encoder versus Decoder for QAs . While BERT-based models [ 20 ] are also used for knowledge extraction through QAs [ 10 , 36 ] , our work indicates that they are less effective at extracting knowledge compared to GPT models.

## 2 Result 0: Our Dataset Families

In this paper, we introduce synthetic human biography datasets and near-real datasets generated by LLaMa [ 40 , 37 ] . Detailed descriptions are in the appendix, with a brief overview here.

BIO dataset bioS . The synthetic dataset, bioS , generates profiles for N = 100,000 N=100,000 individuals. 5 5 5 We have a follow-up to push this to N = 20,000,000 N=20,000,000 and similar results hold [ 3 ] . Each individual’s details are randomly and independently selected from a uniform distribution. The birth dates offer 200 × 12 × 28 200\times 12\times 28 possibilities, while other categories offer 100 ∼ 1,000 100\sim 1,000 choices. We also add a “company city” attribute which depends on the employer’s headquarters location. We ensure uniqueness in each individual’s full name.