Dissecting Recall of Factual Associations in Auto-Regressive Language Models (2023)

Abstract: Transformer-based language models (LMs) are known to capture factual knowledge in their parameters. While previous work looked into where factual associations are stored, only little is known about how they are retrieved internally during inference. We investigate this question through the lens of information flow. Given a subject-relation query, we study how the model aggregates information about the subject and relation to predict the correct attribute. With interventions on attention edges, we first identify two critical points where information propagates to the prediction: one from the relation positions followed by another from the subject positions. Next, by analyzing the information at these points, we unveil a three-step internal mechanism for attribute extraction. First, the representation at the last-subject position goes through an enrichment process, driven by the early MLP sublayers, to encode many subject-related attributes. Second, information from the relation propagates to the prediction. Third, the prediction representation “queries” the enriched subject to extract the attribute. Perhaps surprisingly, this extraction is typically done via attention heads, which often encode subject-attribute mappings in their parameters. Overall, our findings introduce a comprehensive view of how factual associations are stored and extracted internally in LMs, facilitating future research on knowledge localization and editing.

## introduction transformer-based language models (lms) capture vast amounts of factual knowledge @xcite , which they encode in their parameters and recall during inference @xcite . while recent works focused on identifying where factual knowledge is encoded in the network @xcite @xcite , it remains unclear how this knowledge is extracted from the model parameters during inference. in this work, we investigate this question through the lens of information flow, across layers and input positions @xcite . we focus on a basic information extraction setting, where a subject and a relation are given in a sentence (e.g. "beats music is owned by"), and the next token is the corresponding attribute (i.e. "apple"). we restrict our analysis to cases where the model predicts the correct attribute as the next token, and set out to understand how internal representations evolve across the layers to produce the output. focusing on modern auto-regressive decoderonly lms, such an extraction process could be implemented in many different ways. informally, the model needs to "merge" the subject and relation in order to be able to extract the right attribute, and this merger can be conducted at different layers and positions. moreover, the attribute extraction itself could be performed by either or both of the multihead self-attention (mhsa) and mlp sublayers. to investigate this, we take a reverse-engineering approach, inspired by common genetic analysis methods @xcite and the recent work by @xcite . namely, we artificially block, or "knock out", specific parts in the computation to observe their importance during inference. to implement this approach in llms, we intervene on the mhsa sublayers by blocking the last position from attending to other positions at specific layers. we identify two consecutive critical points in the computation, where representations of the relation and then the subject are incorporated into the last position: first the relation and then the subject. next, to identify where attribute extraction occurs, we analyze the information that propagates at these critical points and the representation construction process that precedes them. this is done through additional interventions to the mhsa and mlp sublayers and projections to the vocabulary @xcite @xcite . we discover an internal mechanism for attribute extraction that relies on two key components. first, a subject enrichment process, through which the model constructs a representation at the last subject-position that encodes many subject-related attributes. moreover, we find that out of the three sources that build a representation (i.e., the mhsa and mlp sublayers and the input token embeddings @xcite ), the early mlp sublayers are the primary source for subject enrichment. the second component is an attribute extraction operation carried out by the upper mhsa sublayers. for a successful extraction, these sublayers rely on information from both the subject representation and the last position. moreover, extraction is performed by attention heads, and our analysis shows that these heads often encode subject-attribute mappings in their parameters. we observed this extraction behavior in ∼70% of the predictions. our analysis provides a significantly improved understanding of the way factual predictions are formed. the mechanism we uncover can be intuitively described as the following three key steps (fig. 1 ). first, information about the subject is enriched in the last subject token, across early layers of the model. second, the relation is passed to the last token. third, the last token uses the relation to extract the corresponding attribute from the subject representation, and this is done via attention head parameters. unlike prior works on factual knowledge representation, which focus on mid-layer mlps as the locus of information (e.g. meng et al. (2022b) ), our work highlights the key role of lower mlp sublayers and of the mhsa parameters. more generally, we make a substantial step towards increasing model transparency, introducing new research directions for knowledge localization and model editing. ## background and notation we start by providing a detailed description of the transformer inference pass, focusing on autoregressive decoder-only lms. for brevity, bias terms and layer normalization @xcite are omitted, as they are nonessential for our analysis. a transformer-based lm first converts an input text to a sequence t 1 , ...t n of n tokens. each token t i is then embedded as a vector x 0 @xmath0, over a vocabulary v. the input embeddings are then transformed through a sequence of l transformer layers, each composed of a multi-head self-attention (mhsa) sublayer followed by an mlp sublayer @xcite . formally, the representation x ℓ i of token i at layer ℓ is obtained by: ## experimental setup we focus on the task of factual open-domain questions, where a model needs to predict an attribute a of a given subject-relation pair (s, r). a triplet (s, r, a) is typically expressed in a questionanswering format (e.g. "what instrument did elvis presley play?") or as a fill-in-the-blank query (e.g. "elvis presley played the ____"). while lms often succeed at predicting the correct attribute for such queries @xcite , it is unknown how attributes are extracted internally. for a factual query q that expresses the subject s and relation r of a triplet (s, r, a), let @xmath1, ..., t n ) be the representation of q as a sequence of tokens, based on some lm. we refer by the subject tokens to the sub-sequence of t that corresponds to s, and by the subject positions to the positions of the subject tokens in t. the non-subject tokens in q express the relation r. data we use queries from counterfact @xcite . for a given model, we extract a random sample of queries for which the model predicts the correct attribute. in the rest of the paper, we refer to the token predicted by the model for a given query q as the attribute a, even though it could be a sub-word and thus only the prefix of the attribute name (e.g. wash for "washington"). ## overview: experiments & findings we start by introducing our attention blocking method and apply it to identify critical information flow points in factual predictions ( §5) -one from the relation, followed by another from the subject. then, we analyze the evolution of the subject representation in the layers preceding this critical point ( §6), and find that it goes through an enrichment process driven by the mlp sublayers, to encode many subject-related attributes. last, we investigate how and where the right attribute is extracted from this representation ( §7), and discover that this is typically done by the upper mhsa sublayers, via attention heads that often encode a subject-attribute mapping in their parameters. ## localizing information flow via attention knockout for a successful attribute prediction, a model should process the input subject and relation such that the attribute can be read from the last position. we investigate how this process is done internally by "knocking out" parts of the computation and measuring the effect on the prediction. to this end, we propose a fine-grained intervention on the mhsa sublayers, as they are the only module that communicates information between positions, and thus any critical information must be transferred by them. we show that factual predictions are built in stages where critical information propagates to the prediction at specific layers during inference. method: attention knockout intuitively, critical attention edges are those that, when blocked, result in severe degradation in prediction quality. therefore, we test whether critical information propagates between two hidden representations at a specific layer, by zeroing-out all the attention edges between them. formally, let r, @xmath2, n ] such that @xmath3. we block x ℓ r from attending to x ℓ c at a layer ℓ < l by updating the attention weights to that layer (eq. 5): effectively, this restricts the source position from obtaining information from the target position, at that particular layer. notably, this is different from causal tracing @xcite , which checks what hidden representations restore the original prediction when given perturbed input tokens; we test where critical information propagates rather than where it is located during inference. experiment we use attention knockout to test whether and, if so, where information from the subject and relation positions directly propagates to the last position. let s, r ⊂ [1, n ) be the subject and non-subject positions for a given input. for each layer ℓ, we block the attention edges from the last position to each of s, r and the last (nth) position, for a window of k layers around the ℓ-th layer, and measure the change in prediction probability. we set @xmath4. 2 results fig. 2 shows the results. for both gpt-2 and gpt-j, blocking attention to the subject tokens (solid green lines) in the middle-upper layers causes a dramatic decrease in the prediction probability of up to 60%. this suggests that critical information from the subject positions moves directly to the last position at these layers. moreover, another substantial decrease of 35%-45% is observed for the non-subject positions (dashed purple lines). importantly, critical information from non-subject positions precedes the propagation of critical information from the subject positions, a trend we observe for different subject-relation orders ( §a.1). ## intermediate subject representations we saw that critical subject information is passed to the last position in the upper layers. we now analyze what information is contained in the subject representation at the point of transfer, and how does this information evolves across layers. to do this, we map hidden representations to vocabulary tokens via projection. our results indicate that the subject representation contains a wealth of information about the subject at the point where it is transferred to the last position. ## attribute extraction via attention the previous section showed that the subject representation is enriched with information throughout the early-middle layers. but recall that in our prediction task, only one specific attribute is sought. how is this attribute extracted and at which point? we next show that (a) attribute extraction is typically carried out by the mhsa sublayers ( §7.1) when the last position attends to the subject, (b) the extraction is non-trivial as it reduces the attribute's rank by the subject representation considerably ( §7.2) and it depends on the subject enrichment process ( §7.3), and (c) the relevant subject-attribute mappings are often stored in the mhsa parameters ( §7.4). this is in contrast to commonly held belief that mlps hold such information. ## related work recently, there has been a growing interest in knowledge tracing in lms. a prominent thread focused on locating layers @xcite and neurons @xcite that store factual information, which often informs editing approaches @xcite @xcite . notably, @xcite showed that it is possible to change an encoded fact by editing weights in other locations from where methods suggest this fact is stored, which highlights how little we understand about how factual predictions are built. our work is motivated by this discrepancy and focuses on understanding the recall process of factual associations. our analysis also relates to studies of the prediction process in lms @xcite . specifically, @xcite used fine-grained interventions to show that early mlp sublayers are crucial for memorized predictions. also, @xcite introduced a method for editing knowledge encoded in hidden representations. more broadly, our approach relates to studies of how lms organize information internally @xcite . mechanistic interpretability @xcite is an emerging research area. recent works used projections to the vocabulary @xcite @xcite and interventions in the transformer computation @xcite to study the inner-workings of lms. a concurrent work by @xcite studied contextualization in lms by zeroing-out mhsa values, a method that effectively results in the same blocking effect as our knockout method. in our work, we leverage such methods to investigate factual predictions. ## conclusion we carefully analyze the inner recall process of factual associations in auto-regressive transformerbased lms, unveiling a core attribute extraction mechanism they implement internally. our experiments show that factual associations are stored already in the lower layers in the network, and extracted eminently by the mlp sublayers during inference, to form attribute-rich subject representations. upon a given subject-relation query, the correct attribute is extracted from these representations prominently through the mhsa sublayers, which often encode subject-attribute mappings in their parameters. these findings open new research directions for knowledge localization and model editing.