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

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

The benefit of celebrity data is not universal . Substituting it with the WikiBook dataset improves the model’s English comprehension, yet it still struggles with biographical knowledge extraction. This suggests that only celebrity data of similar form truly aids knowledge extraction for minority groups. In Figure 17 in Appendix G , we further investigate different celebrity data types and instances of minor format differences between minority and celebrity knowledge.

## 7 Result 7: Knowledge Storage for Bidirectional Models

This paper primarily explores knowledge storage and extraction in auto-regressive language models. One may argue that some knowledge issues, such as the consistent knowledge ordering in bioS single , are unique to this task due to its unidirectional nature. We thus pose the question, Could BERT be a solution to this?

We analyze the BERT model [ 20 ] , similar to GPT2 but with a full attention matrix, allowing every token to attend to every other token. For a direct comparison, we modify our GPT2 architecture to replace its triangular attention matrix with a full matrix, keeping the GPT2 tokenizer and rotary embedding. We call this modified model GBERT.

Our pretraining task is now whole-word masked-language modeling (MLM) . Each English whole-word has a 15% chance of being selected, which is then replaced with a <MASK> token (80% chance), retained (10% chance), or replaced with a random token (10%). The goal is to predict the original word for these selected tokens. 20 20 20 We thank Xiaodong Liu and Pengcheng He from the mt-dnn project [ 24 ] for confirming that our MLM implementation aligns with common practice.

For GBERT, we modify the QA task to evaluate its knowledge extraction capabilities. For questions like “What is the birth city of Anya Briar Forger?”, we append them with several <MASK> tokens (equaling the answer’s length). 21 21 21 Revealing the answer’s token count might seem unfair. However, given our aim to highlight GBERT’s limitations, this extra information doesn’t hinder our intentions. A correct answer requires accurate recovery of all masked tokens.

We display results for both mixed training and BIO pretrain + QA finetune. Half of the QAs are used for mixed training (or QA fine-tuning), while we test out-of-distribution generalization accuracies on QAs for the remaining half of the people. Q-probing results for GBERT are also presented, determining if the model, with minor embedding layer modifications, can linearly predict target attributes from a person’s full name.

Our findings. Our findings are displayed in Figure 9 . Key observations include:

• The QA-finetune and Q-probing accuracies again show a strong correlation. This suggests that the ability to extract knowledge from a BERT-like model also depends on whether such information is nearly linearly stored in hidden states directly adjacent to the person’s name. This means our Q-probing technique is effective also for encoder models like BERT.

• Consistent with Figure 3 , mixed training yields slightly superior out-of-distribution QA accuracies compared to BIO pretrain + QA finetune.

• Interestingly, the model performs well on “birth date” and “major” attributes but struggles on others. The reason is simple. In MLM, where each word has an equal chance of being masked, the model learns to associate knowledge words with the most related unmasked word, preferably those that are adjacent . For instance, words representing the “birth date” attribute (month, day, year) are quite independent, making the model more inclined to link them to the person’s name. For attributes like birth city, where there’s a strong link between the city “Bellevue” and state “WA”, the model maximizes this association, inhibiting storage of knowledge on person names . 22 22 22 Similarly, many majors are single words so this explains its high QA test accuracy. In contrast, the words representing universities or company names/cities are more dependent.

Result 7 ( Figure 9 ) . While bidirectional models like BERT is less sensitive to knowledge ordering, the MLM pretraining task does not necessarily promote knowledge storage for subsequent extraction. Unless the knowledge is a standalone word or of independent words (like month, day, year), extracting knowledge after MLM pretraining might prove challenging, if not totally impossible.

## 8 Conclusion

This study explores the capability of pre-trained language models to store and retrieve knowledge through question-answering tasks. We developed a controlled biography dataset and employed probing techniques to assess how knowledge augmentation influences the extractability of knowledge in pre-trained transformer models. Using synthetic data allows for greater control over the training and fine-tuning of models, which is essential for understanding how different data sources impact the internal mechanisms of transformers. This could be a significant future direction for unraveling the complexities of transformers.

For practitioners, this paper emphasizes the importance of rewriting critical but infrequent data during the pretrain stage to enhance knowledge extraction for downstream tasks. Tools like ChatGPT, Llama-7B, or smaller auxiliary models can be used for rewriting before pre-training; these models do not need to contain the knowledge themselves, and even simple techniques like sentence-level shuffling or English-to-French translation can be beneficial.

Additionally, we suggest including more instruction-finetuned data during the pretrain phase. Although this approach differs from human knowledge acquisition, it enhances the model’s ability to encode knowledge more effectively, as explored in recent follow-up work [ 19 ] .

Finally, Part 3 of this work series focuses on how language models store, extract and manipulate knowledge (including Part 3.2 [ 2 ] and Part 3.3 [ 3 ] ). We also cover grade-school math and reasoning in Part 2 [ 38 , 39 ] , and learning hierarchial language structures in Part 1 [ 1 ] .

Appendix

## Appendix A Details on Data Preparation

### A.1 BIO dataset bioS

In the synthetic dataset labeled as bioS , we generate profiles for N = 100,000 N=100,000 individuals. Each individual’s first, middle, and last names, birth date, birth city, university attended, major of study, and current employer are selected independently and randomly from a uniform distribution. • First, middle, and last names are drawn from pools of 400, 400, and 1000 English names respectively. We apply rejection sampling to ensure all N N individuals have unique full names.

• Birth years range from 1900 to 2099, months are selected from the 12 months, and days are chosen between 1 and 28.

• Birth cities are selected from 200 US cities, with their respective state abbreviations, such as Princeton, NJ and Cambridge, MA .

• Universities are drawn from a list of 300 US institutions. Some may have similar prefixes, like University of California, Berkeley/Irvine/Davis/etc .

• Majors are selected from 100 common college disciplines, including Computer Science, Physics, and Music.

• Employers are chosen from a list of 263 companies, featuring names like Meta Platforms, Microsoft, and Google.

Additionally, • We introduce a “company city” attribute that depends on the US location of the employer’s headquarters. For instance, an employee of Meta would list Menlo Park, CA as their company city. Notably, 13.7% of the companies are headquartered in New York, NY. Thus, defaulting to New York, NY when predicting a person’s work city yields a base accuracy of 13.7%.

In the bioS dataset, we craft a biographical text entry for each individual, distilling their profile into six sentences. Each sentence illuminates a distinct attribute of the individual. To increase diversity, we select each sentence randomly from a set of pre-defined templates. Specifically, we have 46 sentence templates for birth dates, 49 for birth cities, 49 for universities, 52 for majors of study, 47 for employers, and 48 for company cities. Beyond ( 2.1 ) , we provide several more examples below: Carlos Jameson Stokes has his annual celebration on November 12, 2088 . He celebrates his birth in San Francisco, CA . He graduated from Oklahoma State University . He explored the theoretical aspects of Information Systems . He contributed his expertise to United Airlines Holdings . He acquired industry knowledge while working in Chicago, IL . Alondra Bennett Rooney celebrates their life journey every year on April 1, 1909 . They owe their roots to Durham, NC . They benefited from the resources and facilities provided by University of South Alabama . They developed a strong foundation in Data Science . They had a job at The Southern Company . They were involved in the industry of Atlanta, GA . Aidan Alexa Dennis ’s birth is celebrated annually on July 17, 1968 . She calls Palmdale, CA her birthplace. She specialized in her field of study at Stevens Institute of Technology . She completed a rigorous program in International Business . She had employment prospects at Johnson & Johnson . She gained work experience in New Brunswick, NJ . (We assign a random pronoun (he/she/they) to each person.) 23 23 23 Given that we are not employing a pretrained model sourced from the internet, we did not do fact-checking. For instance, a person’s major may not align with the business of the company they work for, and their birth year might largely precede the company’s establishment date.

In the basic configuration, we produce a single biographical entry for each individual, maintaining a consistent order for the six sentences as previously outlined. In average, a biographical entry has 73.0 tokens using GPT2 tokenization. We denote this configuration as “ bioS single .” For comparison, we delve into 15 knowledge augmentations: • bioS single+fullname : Pronouns are replaced with the person’s full name.

• bioS single+permute1/2/5 : The six sentences in the biography entry are randomly permuted 1/2/5 times for each person. However, the full name only appears in the first sentence, with subsequent sentences using pronouns. This results in 1/2/5 biography entries for each person.

• bioS single+permute1/2/5+fullname : As with the previous augmentation, but the full name is used in all six sentences.

• bioS multi2/5 : 2 or 5 biographical entries are generated for each person, with each generation employing a re-sampled set of sentence templates.

• bioS multi2/5+permute : Building on bioS multi2/5 , the six sentences within each biographical entry are randomly permuted. However, the full name appears only once in the first sentence.

• bioS multi2/5+fullname : Building on bioS multi2/5 , pronouns are replaced with the individual’s full name across all sentences.

• bioS multi2/5+permute+fullname : Incorporating features from both bioS multi2/5+permute and bioS multi2/5+fullname , the pronouns are replaced with the individual’s full name and the six sentences are randomly permuted.

#### A.1.1 bioS couple

In Section 5.1.1 , when delving deeper into P-probing, we also introduced a partial knowledge augmentation on the bioS dataset, which we termed bioS couple .

Specifically, we initially generate six sentences, each derived from a set of sentence templates similar to those in bioS single . We then group these six sentences into three pairs. The sentence describing a person’s birthdate always precedes the one discussing the person’s birth city. Similarly, the sentence detailing the person’s university consistently comes before the one about their major, and the one about their employer invariably precedes the sentence regarding their work city. Subsequently, we permute the order of these three pairs of sentences, resulting in 3 ! = 6 3!=6 potential arrangements. The individual’s full name is restricted to appear only in the first sentence. For each individual, we create such a biographical entry 1/2/5 times, designating this dataset as bioS couple1/couple2/couple5 . Our experiments in Figure 6 were with respect to bioS couple2 , and we shall give the similar results in Figure 15 for bioS couple1/5 for comparison.

### A.2 BIO dataset bioR

We examine a “close-to-real” dataset produced by Llama [ 40 , 37 ] . Specifically, for the previously set of N = 100,000 N=100,000 individuals, we provide an instructive prompt to Llama1-30B, directing it to generate a biographical entry for us. 24 24 24 Paraphrasing is simple, so larger models are unnecessary as results are similar (our trials confirmed this, but including them seemed uninteresting). At the time of paper writing, Llama2 was not released (but subsequent tests using Llama2 yield similar results, see [ 3 ] ). To promote diversity in our results, we draw our instructive prompts from a collection of sentence templates. Three examples are presented below.

Llama prompt: Write a 50 words biography about a person whose name is Nicole Kevin Pratt. The person studied at Haverford College. The person studied Management there. The person was born and raised in Baltimore, MD. The person worked in Cincinnati, OH for Procter & Gamble. The person’s birthday falls on January 25, 1977. Llama prompt: Write a 50 words biography about a person whose name is Hunter Bennett Kenny. The person attended Queens College, City University of New York for education. The person pursued a degree in Political Science there. The person originated from Augusta, GA. The person worked in Menomonee Falls, WI for Kohl’s. The person was born on March 25, 2033. Llama prompt: Write a 50 words biography about a person whose name is Johnathan Charles Wade. The person attended Colorado State University for education. The person pursued a degree in Sociology there. The person was brought up in New York City, NY. The person worked in Northbrook, IL for Allstate. The person’s birthday falls on January 7, 2098.

Here are the generated outputs from Llama [ 40 , 37 ] : Nicole Kevin Pratt is an American business executive. She is currently the Vice President of P &G Global Business Services at Procter & Gamble . She was born on January 25, 1977 , in Baltimore, Maryland . She graduated from Haverford College with a degree in Management . P &G recruited her as an Assistant Brand Manager in 2000. She held various leadership positions in brand management, marketing, and sales across different business units and categories. She was named Vice President of P &G Global Business Services in 2019. Nicole currently lives in Cincinnati, Ohio with her husband and three children. Hunter Bennett Kenny is a talented political science graduate from Queens College, City University of New York . He hails from Augusta, Georgia and was born on March 25, 2033 . During his time at college, he was an active member of the student council and served as its president in his senior year. He interned at the office of New York Senator Chuck Schumer. After graduating cum laude, he worked for Kohl ’s in Menomonee Falls, Wisconsin . He currently resides in Brooklyn, New York. Johnathan Charles Wade is a successful insurance agent who works for Allstate . He was born on January 7, 2098 , in New York City, NY . He graduated from Colorado State University , where he majored in Sociology . He currently resides in Northbrook, IL .

We observed that Llama might not always generate a biographical entry that includes all six attributes. To address this, we repeat the sampling process until Llama’s output ensures the inclusion of all attributes. Typically, the entry begins with the individual’s full name, and the oder in which the six attributes appear can vary. In average, a biographical entry has 72.3 tokens using GPT2 tokenization.

In the basic configuration, we produce a single biographical entry for each person, denoted as “ bioR single .” For comparison, we also introduce the multi M M augmentation, which creates M M entries per person, and the fullname augmentation.