Characterizing Mechanisms for Factual Recall in Language Models (2023)

Abstract: Language Models (LMs) often must integrate facts they memorized in pretraining with new information that appears in a given context. These two sources can disagree, causing competition within the model, and it is unclear how an LM will resolve the conflict. On a dataset that queries for knowledge of world capitals, we investigate both distributional and mechanistic determinants of LM behavior in such situations. Specifically, we measure the proportion of the time an LM will use a counterfactual prefix (e.g., “The capital of Poland is London”) to overwrite what it learned in pretraining (“Warsaw”). On Pythia and GPT2, the training frequency of both the query country (”Poland”) and the in-context city (”London”) highly affect the models’ likelihood of using the counterfactual. We then use head attribution to identify individual attention heads that either promote the memorized answer or the in-context answer in the logits. By scaling up or down the value vector of these heads, we can control the likelihood of using the in-context answer on new data. This method can increase the rate of generating the in-context answer to 88% of the time simply by scaling a single head at runtime. Our work contributes to a body of evidence showing that we can often localize model behaviors to specific components and provides a proof of concept for how future methods might control model behavior dynamically at runtime.

## introduction large transformer language models @xcite (lms) store information from pretraining which they can recall at inference time to generate text. this is paired with the exceptional ability of models to use provided context in order to produce coherent text that incorporates new facts. however, facts that are memorized in pretraining and facts that are provided in-context can often compete with each other; in some cases it might the capital of poland is london. q: what is the capital of poland? a: warsaw ## related work the impressive, but sometimes unpredictable successes of lms on performing tasks described in context @xcite has spurred intense interest in the factors that allow models to solve tasks this way. studies on pretraining datasets have found that higher pretraining term frequency is positively correlated with task performance on factual association tasks @xcite and numer-ical reasoning @xcite , and relates to work on memorization vs. generalization in lms @xcite . @xcite analyze mechanisms used to recall memorized information by studying idiom generation. model size is also shown to be a factor that affects tendency to use memorized vs. in context information @xcite . previous work has also examined how deeply lms interact with context during in-context learning @xcite . other work has focused on lms' abilities to consider counterfactual or hypothetical contexts @xcite , with mixed results in overwriting pretraining memory. our work is built heavily on previous work in mechanistic interpretability, which aims to reverse engineer model computations into human understandable components @xcite @xcite . while knowledge from pretraining has been found to be stored in the feedforward (mlp) sublayers @xcite @xcite , more recent work has also clarified the role of attention in this same process @xcite find that attention heads extract facts from an earlier mentioned subject token (e.g., poland) when required. this naturally sets up our study, which also considers attention heads as the source of the competing effect between copying the counterfactual from earlier in context vs. extracting the memorized fact from an earlier subject token. a core technique in these works is projecting activations from model components into the vocabulary space to make claims about their roles, which we generically refer to here as logit attribution @xcite @xcite @xcite . we leverage this technique to localize attention heads which tend to promote either context or memorized information ( §6). ## task design we study the mechanism that language models use when given counterfactual information in context. for our analysis, we focus on a simple zero-shot task that requires producing the capital city for a given country, which serves as a representative example of the type of common facts that a language model could learn in pretraining. it consists of 248 country capital pairs with 6 additional aliases and their respective capitals. to create counterfactuals the upward trajectory of the red lines show the positive correlation between the proportion of memorized answer predictions and frequency. the last two graphs reflect the frequency of the in-context capital (e.g., london). the drop of the blue lines across all the four graphs show the negative correlation between proportion of in-context answer predictions and term frequency. in the dataset, we pair up every country with the rest of the 247 capitals using the following format: the capital of {country} is {in-context city}. q: what is the capital of {country}? a: for example, we can fill in {country} with poland and {in-context city} with london. the model has learned that the capital is warsaw from pretraining, but with the in-context prompt, the task becomes ambiguous between whether it should output the known answer warsaw or overwrite it with london. we query the model to generate a full sentence to determine which of the above task interpretations it preferred. we define poland as the country and london as the in-context answer. correspondingly, we define warsaw, the capital of poland, as the memorized answer. if london is included in the sentence, then we consider the model to have produced the in-context answer. if the model generates warsaw, then it is considered as the memorized answer. in total, we have 62,992 such pairs of country and capital. the world capital dataset is able to provide a clean analysis giving a unique memorized and incontext answer. language models perform well on this task producing one of the two expected cities at least 80% of the time (varying depending on model size). the lack of noise in the responses makes the task a good choice for cleanly diagnosing model preference for memorized vs. in-context answers. ## models we analyze the overwriting behavior primarily on the pythia @xcite models as well as gpt2 series @xcite . the pythia models are trained on the pile @xcite which we have full access to, allowing us to relate model behavior to frequency effects in the data. while we don't have access to the pretraining data of gpt2, we still report results, using the pile frequencies as an approximation of what gpt2 might have seen. ## discussion & future work this paper investigates factors that influence a model's propensity to favor in-context vs. memorized factual associations, when the two compete with one another. our results demonstrate that the frequency of information in the pretraining corpus can affect the model's tendency to use new, conflicting information provided in context. building on this, we provide a proof of concept that this tendency can be controlled by a mechanism in the attention heads which allows us to manipulate lms' tendency to prefer new in-context information without modifying any model parameters directly. by building off insights from mechanistic interpretability, we can localize single attention heads that contribute to this mechanism. this provides evidence that decomposing complex neural networks into understandable components is possible, even in models with billions of parameters. still, we observe that the selected heads promote domain specific knowledge rather than a more abstract concept of truthfulness. this brittleness is characteristic of mechanistic analyses of larger models, and should be a priority for future work. nonetheless, given the early stage of research on this level of analysis of large language models, findings of this type even in an isolated setting are exciting and can lay the groundwork for subsequently discovering more general mechanisms. the exploratory methods described here suggest avenues via which future work might develop more sophisticated techniques for controlling and auditing deployed language models. adapting lms posthoc for applications that require domain-specific information is a growing problem. for example, there are simultaneously reasons we might want to suppress the use of in-context information at run time (e.g., to combat prompt-injection attacks) as well as reasons we might want to encourage it (e.g., to enable users to provide new, personalized, or hypothetical information to the model). the intervention we describe in this work is intriguing in that it can be used without changing the model and can be turned on and off dynamically within the forward pass. it thus offers a promising direction for further work on model editing. ## conclusion in the problem setting of predicting world capitals, our results show that the ability of language models (lms) to overwrite information that it memorized in pretraining depends both on the frequency of the subject of the new fact (the country, e.g., poland), as well as the frequency of the overwriting information (the counterfactual city, e.g., london). we can intervene on attention heads that we find tend to push the prediction one way or another. by simply rescaling the value vectors of important heads, we can control which city the model predicts without updating any model parameters. we hope these results encourage future work in understanding the internal structure of neural networks in general.