Competition of Mechanisms: Tracing How Language Models Handle Facts and Counterfactuals

Authors: francesco  ortu
zhijing  jin
diego  doimo
mrinmaya  sachan
alberto  cazzaniga
bernhard  schölkopf
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-13 06:24:28

Abstract: Interpretability research aims to bridge the gap between empirical success and our scientific understanding of the inner workings of large language models (LLMs). However, most existing research focuses on analyzing a single mechanism, such as how models copy or recall factual knowledge. In this work, we propose a formulation of competition of mechanisms, which focuses on the interplay of multiple mechanisms instead of individual mechanisms and traces how one of them becomes dominant in the final prediction. We uncover how and where mechanisms compete within LLMs using two interpretability methods: logit inspection and attention modification. Our findings show traces of the mechanisms and their competition across various model components and reveal attention positions that effectively control the strength of certain mechanisms.

Main paper: 
 introduction and related works: 
 Large language models (LLMs) have achieved remarkable performance in various NLP tasks, revolutionizing applications across multiple domains (Brown et al., 2020; Touvron et al., 2023; OpenAI, 2023; Anil et al., 2023, inter alia) . However, their black-box nature poses significant challenges to our scientific understanding of their inner workings. This gap between empirical success and mechanistic comprehension has led to a growing focus on interpretability research, which aims to decode the internal processes of these complex models.Interpretability research in LLMs has primarily followed two main trajectories: interpreting representations and decoding specific mechanisms. The first approach focuses on understanding what information is encoded in model states (Belinkov et al., 2017; Conneau et al., 2018; Hewitt and Manning, 2019) . These studies have revealed that LLMs capture a rich array of linguistic and world knowledge within their hidden states. The second approach, mechanistic interpretability, aims to uncover the specific operations learned by LLMs (Olsson et al., 2022; Geva et al., 2023; Meng et al., 2022; Hanna et al., 2023, inter-alia) . For instance, Olsson et al. (2022) identified induction heads responsible for the copy mechanism, a basic yet crucial operation in LLMs. Similarly, studies by Geva et al. (2023) and Meng et al. (2022) have shed light on how LLMs mechanistically recall factual information, showing that early MLP layers enrich subject embeddings while late attention blocks select and write factual information. Despite these advancements in understanding individual mechanisms, less attention has been paid to how these mechanisms interact and compete within the model's decision-making process. This gap in our knowledge is particularly crucial when LLMs face scenarios requiring them to balance multiple, potentially conflicting sources of information -such as when presented with counterfactual statements that contradict their pre-trained knowledge.In this study, we propose a novel formulation of competition of mechanisms to investigate the interplay between multiple mechanisms in LLMs. Our work specifically focuses on how one mechanism becomes dominant in the final prediction by winning this competition. We examine the interaction between two well-studied mechanisms: factual knowledge recall and in-context adaptation to counterfactual statements. This approach allows us to explore how LLMs navigate the tension between their pre-trained knowledge and new information presented in the input context. Based on the latest tools to inspect each of these two mechanisms (Nostalgebraist, 2020; Wang et al., 2023; Geva et al., 2023) , we then unfold how and where the competition of the two mechanisms happen. Our analysis spans both macroscopic (e.g., layer-level) and microscopic (e.g., attention head) views, providing a comprehensive picture of how information flows and competes within the model architecture. 
problem setup and methods: 
 We design a task to incorporate the competition of mechanisms by pairing factual statements such as "iPhone was developed by Apple." with corresponding counterfactual statements, as "iPhone was developed by Google.". We compose prompts adding the counterfactual statements as a false definition of the factual sentence, allowing us to trace the competition between factual (t fact ) and counterfactual (t cofa ) tokens, such as "Redefine: iPhone was developed by Google. iPhone was developed by ___". We utilize the COUNTERFACT dataset (Meng et al., 2022) , selecting 10,000 examples Few localized attention heads are responsible to modulating the competition between factual recall and counterfactuals redefinition. On the left, the direct contribution to ∆ cofa := HeadLogit(t cofa ) -HeadLogit(t fact ) of all heads in GPT-2. Heads favoring t fact are colored in blue, and those favoring t cofa in red. On the right, the factual recall accuracy before and after modifying the target heads in GPT-2 and Pythia-6.9B. We up-weight the heads that favor fact, two in GPT-2 and three in Pythia-6.9B.where attributes are single tokens and the model completes sentences accurately. To analyze token preferences across model components, we project hidden representations to the vocabulary space using an unembedding matrix W U (Halawi et al., 2023; Geva et al., 2023; Dar et al., 2023; Geva et al., 2022; Nostalgebraist, 2020) . We also employ attention matrix modification techniques to further elucidate information flow within LLMs, intervening on target attention heads and measuring the effect in the model's performance. Our primary focus is on the GPT-2 small model (Radford et al., 2019) , aligning with previous interpretability studies (Meng et al., 2022; Wang et al., 2023; Conmy et al., 2023; Hanna et al., 2023) . To demonstrate generalizability, we provide supplemental results for Pythia-6.9B (Biderman et al., 2023) , enhancing the robustness of our findings across LLMs of different architectures and scales. 
results and findings: 
 Using these methods, we assess the contributions of various model components, both from a macroscopic view (e.g., each layer) and a microscopic view (e.g., attention heads), and identify critical positions and attention heads involved in the competition of the two mechanisms. Moreover, we locate a few localized positions of some attention head matrices that can significantly control the strength of the factual mechanism. We summarize our main findings as follows:1. In early layers, the factual attribute is encoded in the subject position, while the counterfac-tual is in the attribute position; 2. The attention blocks write most of the factual and counterfactual information to the last position; 3. All the highly activated heads attend to the attribute position regardless of the specific type of information they promote. The factual information flows by penalizing the counterfactual attribute rather than promoting the factual one; 4. We find that we can up-weight the value of a few very localized values of the attention head matrix to strengthen factual mechanisms substantially. 
conclusion: 
 Our study introduces the concept of "competition of mechanisms" as a novel interpretability framework for understanding how LLMs handle multiple, potentially conflicting mechanisms. This approach provides valuable insights into the inner workings of language models, particularly in scenarios where they must navigate between pretrained knowledge and conflicting contextual information. Our findings reveal that the suppression of counterfactual information plays a more significant role than the promotion of factual information in the model's decision-making process. This insight, along with our discovery of localized attention positions that control the strength of the factual mechanism, opens up new possibilities for targeted model fine-tuning and optimization. 
 
 Appendix: