Knowledge Storage and Extraction in Language Models (Part A)

This paper explores how LMs capture the factual knowledge: by constructing an internal KG or simply memorized QA pairs. To do that, the authors use LLaMA to construct a synthetic data set and train GPT-2 with it. Then, the verification is done by prompting and probing.

This paper primarily addresses the research question: How do language models memorize factual knowledge during pre-training, and how do they extract and apply this learned knowledge for reasoning and answering questions? To answer this question, the authors create two biographical text datasets, one generated using templates and the other through prompting the LLAMA model. Additionally, a question-answering dataset is developed to evaluate the model's capability to extract knowledge. The key findings through controlled experiments include:

1. During pre-training, models tend to focus on quickly learning to answer questions rather than comprehending the complete knowledge behind the answer. This behavior might hurt the model's generalization performance on unseen data.
2. If the representation in the pre-training dataset is too uniform, it might impact the model's knowledge acquisition, potentially affecting its generalizability. Hence, data augmentation to diversify knowledge representation can help the model better link related concepts, enhancing knowledge extraction.

This paper analyzes the gap between knowledge memorization and utilization of LLMs. Upon a carefully created benchmark, it examined the effect of different training processes and data augmentation.

The proposed paper, entitled "Knowledge Storage and Extraction in Language Models (Part A)," aims to explore how transformer-based language models store and extract knowledge during training and inference. The authors focus on factual knowledge, such as that found in a knowledge graph, which the language model needs to memorize from the training corpus, encode in its weights, and later extract to answer questions or perform logical reasoning during inference.

The paper addresses the following key questions:

1. How do language models memorize knowledge during training, and extract it later to answer questions or perform logical reasoning during inference?

2. Can language models be fine-tuned to extract knowledge from specific domains, and if so, how do they achieve this?

3. What is the relationship between the model's knowledge extraction ability and different diversity measures of the training data?

4. How does the model encode knowledge attributes at the hidden embedding of entity names, and is there a correlation between this encoding and the model's knowledge extraction ability?

To study these questions, the authors propose conducting an in-depth analysis using a controlled set of semi-synthetic biography data. They construct a synthetic dataset of 100k biographies, including attributes such as birthday, birth city, and major of study, and pretrain the language model on this dataset. The authors then evaluate the model's knowledge extraction ability by fine-tuning it on question and answer (QA) pairs for a fraction of individuals and testing its ability to answer QAs about the remaining fraction.

This paper investigates a very important question regarding how LLMs extract knowledge from text. This is different from the LLMs ability to regurgitate text that it had seen during training by being able to complete generations given a prefix. Instead, this paper investigates whether LLMs are actually able to understand facts by evaluating whether it can answer questions related to what it has been trained on. However, it is hard to study this phenomenon on LLMs that have been trained on web-text because it is very hard to ensure what the LLM might have seen during training. Instead, the authors study this in a semi-synthetic setting where the LLM is trained on synthetic biographies generated by templatized texts or LLama-30B. The LLM that has been pre-trained on scratch is a GPT2-small model.

The paper considers two different training paradigms - (a) Mixed training: in which the LLM is pre-trained on all the biographies as well as 50% of the question answer pairs, and (b) Pre-train + Fine-tune setting where the LLM is first pre-trained on biographies and then fine-tuned on 50% of the QA pairs. Note during mixed-training, the LLM observes biographies and QA pairs interspersed, but in pre-training+fine-tune the LLMs first observes only biographies and then only QA pairs. The performance of the model is tested on held-out QA pairs $P_{\mathrm{test}}$

In the mixed-training setting, the LLMs are able to generalize to held-out questions but the paper shows that the learning behavior of the models are a bit abnormal. For example, instead of first learning from biographies, the LLMs instead first learns from the QA pairs (as demonstrated by rapid increase in accuracy of token prediction) and then aligns the encoded knowledge with the BIO data to learn and extract knowledge and generalize it to the held-out questions. The authors point out this is akin to “students studying to pass the test”.

The previous hypothesis is confirmed in the pre-train+fune-tune setting where the model fails to generalize at all on $P_{\mathrm{test}}$. This is surprising as in both cases the model is trained on the same amount of data. This demonstrates that LLMs are great at memorizing and regurgitating text but that is not equivalent to extracting knowledge from it.

The paper next shows that knowledge augmentation-i.e. Augmenting the BIO set with multiple paraphrases, sentence permutations, and replacing pronouns with full-names increases the generalization accuracy (9.7% -> 96.6%). A possible explanation of this hypothesis is that exposing the model to varied expression of identical knowledge lets it focus on semantics rather than surface forms. This hypothesis is confirmed by training two kinds of linear probes that indeed show that upon knowledge augmentation, the representation of entity names capture attributes about the entities that are amenable to be extracted via fine-tuning for QA. The paper also shows that this generalization also holds true when a part of the data is repeated, the model can learn and apply its knowledge to new entities.

Dear authors and reviewers,

This a reminder that deadline of author/reviewer discussion is AOE Nov 22nd (today). Please engage in the discussion and make potential adjustments.

Thank you! AC