Inspecting and Editing Knowledge Representations in Language Models

Authors: evan  hernandez
belinda z li
jacob  andreas
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 19:42:23

Abstract: Neural language models (LMs) represent facts about the world described by text. Sometimes these facts derive from training data (in most LMs, a representation of the word *banana* encodes the fact that bananas are fruits). Sometimes facts derive from input text itself (a representation of the sentence *I poured out the bottle* encodes the fact that the bottle became empty). We describe REMEDI, a method for learning to map statements in natural language to fact encodings in an LM's internal representation system. REMEDI encodings can be used as *knowledge editors*: when added to LM hidden representations, they modify downstream generation to be consistent with new facts. REMEDI encodings may also be used as *probes*: when compared to LM representations, they reveal which properties LMs already attribute to mentioned entities, in some cases making it possible to predict when LMs will generate outputs that conflict with background knowledge or input text. REMEDI thus links work on probing, prompting, and LM editing, and offers steps toward general tools for fine-grained inspection and control of knowledge in LMs.

Main paper: 
 remedi: 
 Motivations: control and interpretability Consider the examples from Fig. 1 (top). In the first example, the LM is prompted with the text Anita's law office serves the lower Eastern Shore. . . , which provides some context about the entity Anita. However, when the LM generates a continuation of this prompt, it asserts that Anita is a nurse. We term this incoherence a failure of context integration: information provided in the textual context has failed to alter the LM's predictions. It would be useful to identify and fix such errors, changing a model's encoding of entities like Anita to ensure that she is correctly described as an attorney. In addition to ensuring discourse coherence, it is often desirable to modify prior associations in LMs. In Fig. 1 (middle) the LM strongly associates London Bridge with the city of London because the most famous London Bridge is located there. However, there could be (and are2 ) other London Bridges, and we might wish to control an LM to make the lesser-known bridge more salient.It is sometimes possible to achieve these goals by carefully prompting models with the right input text. But due to the non-systematic, opaque nature of prompt engineering (Jiang et al., 2020b) , significant manual effort is often required to find a prompt (if one exists at all) that yields correct behavior and generalizes to new use cases. 3 Techniques for localizing generation failures within LMs' internal representations would make it possible to detect them in advance, and guide research aimed at mechanistic understanding of the relationship between LMs' internal representations and their textual outputs.Overview At a high level, our proposed approach learns how to intervene in an LM's representation space to modify the LM's knowledge about a mentioned entity (like Anita in Fig. 2 ). This intervention ultimately updates the LM's representation of the entity to encode an attribute (e.g., is a lawyer) so that the LM will generate text about the entity consistent with the new attribute. This update operation can be specified by a single vector, and is applied to the hidden representation of a single token at a single layer. Edits produced by REMEDI can also be applied out-of-context (enabling controlled generation without textual prompts). By comparing edit vectors to unedited representations, REMEDI additionally makes it possible to inspect representations of entities and attributes produced during ordinary model operation.Editing representations Assume we have a language model p LM (x) that assigns probabilities to strings x consisting of tokens x 1:n . In this paper, p LM will always be an autoregressive transformer (Vaswani et al., 2017) pretrained on English text, as in the GPT family of models (Radford et al., 2019; Brown et al., 2020) . These models decompose p(x) into a product of next-token distributions given preceding context:p LM (x) = ∏ i p LM (x i | x 1:i-1). Our goal is to modify an LM's internal state to cause it to generate desired text about a target entity.Where and how should we perform this modification? LMs encode factual information in hidden representations of entity mentions: entities' states (Li et al., 2021) , perceptual features (Abdou et al., 2021) , and other semantic properties (Grand et al., 2018) have been shown to be linearly decodable from entity representations. To ensure that an LM encodes the fact Anita is a lawyer, it should thus suffice to find an appropriate transformation of the representation of the token Anita.Formally, we denote the transformer's hidden representation for token x i in layer ℓ as h (ℓ)i , and we writep LM (x | h (ℓ) i = z) to mean the output of p LM with h (ℓ)i "hard-coded" to equal z during the forward pass. 4 Given representations of the entity h entity and the target attribute h attr , REMEDI specifically learns an affine transformation F that returns a new entity representation z according to:EQUATION)such that when z replaces the entity inside the LM, the LM will generate text consistent with the target. How should we pick h attr , W and b? Building on the success of linear probing approaches (Conneau et al., 2018; Belinkov & Glass, 2019) , it is tempting to begin by training a classifier for the presence or absence of attributes. For example, following Li et al. (2021) , we could take h attr to be the LM's own representation of an attribute (like plays the oboe; Fig. 2 ), then optimize W and b to predict whether an entity representation encodes the attribute:p(attribute | entity) = σ(h ⊤ entity Wh attr + b) .(2) However, even when an LM encodes information in its representations, this information may not causally influence subsequent generation (Ravfogel et al., 2020; Elazar et al., 2021; Ravichander et al., 2021 ). An effective editor must identify fact encodings that are causally linked to output.Learning Effective Edits REMEDI optimizes W and b to directly intervene in an LM. We assume access to a dataset of tuples (x 1:n-1 , i entity , t attr , t tgt ), where x 1:n-1 is a textual context (e.g. John went to work at), i entity is the index of an entity within the context, t attr is the attribute to be inserted (plays the oboe), and t tgt is a generation target: a completion that should be assigned high probability if the attribute is applied to x entity (the concert hall). Following Li et al. (2021) , we obtain a representation h attr by averaging the LM's encoding of t attr . We train the editor F to maximize the probability that p LM assigns to t tgt after modifying the hidden representation of x entity (Fig. 2 ):L tgt (z) = -log p LM (x n = t tgt | x 1:n-1 , h entity = z) .(3)Learning Non-Destructive Edits When LMs encode strong prior associations between entities and properties (e.g., in the London Bridge example; see Hase et al., 2021) , it is necessary to remove these facts while inserting new ones. We obtain a target string t prior assigned a high pre-edit probability, and train F to minimize the probability of t prior :EQUATION)Finally, to prevent the degenerate solution in which the language model always (and only) predicts t tgt , we penalize the language model for changing its distributions on all tokens between the entity mention and the time at which it predicts t tgt :L KL (z) =EQUATION)Unlike the distribution over tokens at the end of the prompt, which should change dramatically under the intervention, the distribution over these intermediate tokens should not change significantly. L KL penalizes such changes. Thus, the complete objective function is:EQUATIONwhere λ 1 and λ 2 are hyper-parameters. We evaluate REMEDI by studying its ability to control model output (Section 4) and to interpret and predict model behavior (Section 5). 
related work: 
 Probing factual knowledge Large language models (LLMs) trained on massive text datasets have been shown to encode context-agnostic factual knowledge, which can be queried through a text prompt (Petroni et al., 2019) . Most work on extracting background factual knowledge from LMs focuses on designing textual queries for different sources of knowledge (Richardson & Sabharwal, 2020; Peng et al., 2022) . Additionally, knowledge probes may sometimes recover factual information even in cases when LMs do not generate truthful outputs with high probability (Burns et al., 2022) .Probing representations of individual situations Neural LMs have also been shown to build representations of context-dependent knowledge. (Li et al., 2021) show that they track aspects of entity state over a discourse, and this state can be extracted from LM representations of contextualized entity tokens. Furthermore, many LMs have been (indirectly) evaluated on their ability to track context-dependent knowledge by having their performance measured on downstream reading comprehension tasks in wich the LM is expected to answer questions about facts within a discourse. Reading comprehension datasets such as CoQA (Reddy et al., 2019) , RACE (Lai et al., 2017) , and SQuAD (Rajpurkar et al., 2016) are now part of the standard evaluation suite for new LMs; and most modern LMs perform well (Brown et al., 2020) . However, generating does not always imply knowing. Datasets contain spurious correlations (Gururangan et al., 2018) , and LMs are sensitive to the phrasing of prompts and questions (Jiang et al., 2020b) .Editing LLMs In the past, LLMs have been predominantly adapted to new tasks and knowledge through fine-tuning (Devlin et al., 2019) . Recently, with very large LMs, new classes of adaptation methods have been introduced, which generally fall into one of the following two categories: (1) Prompt design approaches prepend a textual prompt to each example specifying the adaptation target (Brown et al., 2020) . (2) Prefix-tuning approaches prepend continuous learned tokens ahead of each example. These specify a task for the LM similarly to how a textual prompt might (Li & Liang, 2021; Lester et al., 2021) . Control token approaches similarly use these learned tokens to controls aspects of LM output, including sentiment (Dathathri et al., 2020 ), style (Keskar et al., 2019) , and semantics (Ross et al., 2022) . Prompts can be fragile; LMs may fail to generate text consistent with the prompt, as in Fig. 1 .Finally, a large body of work examines how to localize and edit factual information in an LM's parameters (Meng et al., 2022a; b; Mitchell et al., 2022; Dai et al., 2022) . For example, ROME (Meng et al., 2022a) localizes factual knowledge in LMs to a particular subset of MLP modules, and edits specific facts in a targeted way through rank-one modification of MLP weights. Unlike REMEDI, these approaches operate on models' weight matrices rather than representations, meaning they can correct errors in models' background knowledge but not information provided in context. 
controlling generation: 
 We begin by showing that the REMEDI procedure described in Section 2 is an effective tool for controlling LM output. Intuitively, if REMEDI succeeds in creating a new entity representation encoding the desired attribute, text generated by the LM about the entity should at minimum (a) prefer generations consistent with the target attribute over potentially contradictory attributes and (b) remain as fluent as the original generations. Our experiments in this section test properties (a) and (b), as well as other quality measures, in two different settings.In the first setting, we use REMEDI to patch incoherence errors, editing the LM to reinforce the information provided in the context. In the second setting, we use REMEDI to update prior knowledge about entities (such as the Versace Headquarters example in Fig. 1 ). A third set of experiments, described in Appendix D, evaluates an additional word re-definition task. These experiments show that REMEDI often successfully controls model behavior even when prompting fails. It can thus serve as a building block for future controlled generation interfaces that allow users to directly steer model behavior in representation space.We first use REMEDI to manipulate representations of generic named individuals, such as Anita or Dennis, about whom the LM should have no prior association (and about whom the LM should acquire all information from the prompt). We provide a small amount of context about each person and prompt the LM to predict their occupation from a small set of candidates. As we will show, the LM often completely ignores this context, and prefers unrelated occupations to ones highly relevant given the context (nurse vs. attorney in Fig. 1 ).In this and all following experiments, we use GPT-J as the underlying language model (Wang & Komatsuzaki, 2021) and include results for two additional models in Appendix E. GPT-J is a 6B parameter, decoder-only transformer pretrained on the Pile (Gao et al., 2020) . We obtain biographical sentences from the Bias in Bios Dataset (De-Arteaga et al., 2019) . This dataset consists of ≈397k short professional biographies people scraped from the internet. Each biography is paired with a label for the subject's occupation. We take one sentence from each biography (details in Appendix B), using only the subject's first name, and prompt the LM with the biographical sentence followed by {Person} has the occupation of. We then look at the relative probabilities of 28 occupations under the LM, and consider output correct if the true occupation is ranked first. GPT-J succeeds about half the time (55%, In-context baseline in Table 1 ) on this task. Method We use REMEDI to create new representations of the first-name-only entities encoding the target occupation. We take h entity to be the last token of the last entity mention (right before model predicts the occupation), and we take h attr to be the average representation of the biographical sentence after the entity. Note this means we are not using any additional data to construct the new entity-the input to REMEDI is all text provided in context to the LM. We train the editor on 5000 examples using Eq. ( 6), with t tgt set to the target occupation and with no t prior term (λ 1 = 0). Edits are performed in layer 12 (this and other hyperparameters discussed in Appendix C). We evaluate factual accuracy and fluency before and after applying REMEDI on 5000 test examples. Accuracy is evaluated by measuring how often the highest-probability occupation is the true one, and fluency using the same n-gram entropy measure as Meng et al. (2022a) .Crucially, we train REMEDI only out-of-context: choosing the initial text x to consist only of the entity's name. However, as described below, we evaluate its ability to control model behavior both in-context (with additional text provided) and out-of-context (on inputs similar to the training data).Results are shown in the left portion (In Context) of Table 1 , which reports GPT-J's factual accuracy and fluency before and after applying REMEDI. REMEDI increases GPT-J's accuracy by over 15% on held-out (entity, attribute) pairs, showing that representations produced by REMEDI more often encode the desired attribute. REMEDI also preserves the fluency of the generated text. We find similar improvements in other models (Appendix E). Table 2 : Examples of GPT-J generations when the attribute is specified in the textual prompt (In-Context Generation) or inserted via REMEDI. For both generic and famous entities, REMEDI causes GPT-J to respect the attribute more often than when it is provided in context.We contextualize these results by evaluating model behavior when the LM has no textual context (i.e. no initial biographical sentence). Here, the base LM has no information about entities' occupations, and obtains near-chance factual accuracy. However, inserting REMEDI's representations into the LM causes it to generate fluent text consistent with the edit, showing that REMEDI can not only enforce coherence with a textual context, but replace textual prompting by inserting information directly into entity representations. REMEDI is slightly more effective at in-context editing than out-of-context editing, despite being trained only out-of-context. The last column of Table 2 shows examples of in-context generation.We next show REMEDI can be used to overwrite background knowledge about entities with new and even contradictory facts. As shown in Fig. 1 et al., 2022; Dai et al., 2022; Meng et al., 2022b) . They all share the limitation of changing the behavior of the LM globally: users cannot choose when to apply edits. In existing methods, edits often bleed into closely related but distinct entities (Meng et al., 2022a) . Because REMEDI operates directly on entity representations at runtime, it applies changes only to the entity of interest at the moment of use.We evaluate REMEDI on the COUNTERFACT benchmark from Meng et al. (2022a) , which consists of (subject, relation, old value, new value) tuples-e.g. (Megan Rapinoe, plays sport, soccer, tennis)-and measures LMs' ability generate natural text consistent with the new fact.We train REMEDI on a subset of 5000 examples from COUNTERFACT and evaluate it on a held-out subset of 5000. As before, we take h entity to be the last token of the entity mention (which appears at the beginning of COUNTERFACT examples) and h attr to be the average representation of the new fact in context. For example, we pass (Megan Rapinoe plays the sport of soccer) to the LM and compute h attr from the underlined tokens. This textual context is akin to the biographical sentence used to compute h attr in the previous section. We use all three loss terms from Eq. ( 6) and apply edits in layer 1; see Appendices B and C for other hyperparameters and implementation details. As above, we train REMEDI without textual context, with inputs x consisting of entity names alone.We include comparisons to the model-editing method ROME and ordinary fine-tuning, following the exact procedures laid out in Meng et al.. However, our primary baseline is one in which the new factual information is prepended to the prompt. In all other methods, the language model is only given a prompt with no context about the fact.We additionally include a baseline in which we find-and-replace the entity with one that shares the target attribute (e.g., replacing Versace headquarters with Harrods). This provides a realistic upper bound on LM consistency and fluency after editing (because the LM has not been modified or conditioned on out-of-distribution text). Essence captures how much the edited entity is still "itself" according to the model (is London Bridge still a bridge?). Formally, it measures tf-idf similarity between the model's generations before and after the intervention given the prompt: {Entity} is .Table 3 shows metrics for REMEDI and baselines. Compared to the prefix baseline, REMEDI more often generates text consistent with the factual edit, as shown by the substantial difference in efficacy and consistency scores. The base LM incorporates textual prompt information 80.2% of the time, while REMEDI-based prompting incorporates new information 98.2% of the time. This performance comes at some cost to the essence of the entity, likely because the original fact is strongly associated with other properties of the entity.  
detecting errors: 
 Next, we show that REMEDI can be used as a model evaluation tool, automatically characterizing when (un-modified) LMs have successfully acquired background or contextual knowledge. A core challenge when deploying language models is that it is difficult to automatically detect when they exhibit the failures shown in Fig. 1 . Some work addresses this challenge by calibrating LM predictive distributions to better reflect veracity (Jiang et al., 2020a) , or training auxiliary models to reject bad samples (Cohen et al., 2022) . REMEDI offers a new approach to detecting when LMs will fail to integrate information from context: inspecting their representations for the information that REMEDI would add. This approach is related to a method by Burns et al. (2022) , which finds implicit LM encodings of veracity. REMEDI identifies cases where even LM-internal states encode errors.Method Suppose we have a prompt that queries an LM for a fact: To cross London Bridge, one should travel to . . . . Under the hypothesis from Section 2, the LM should answer this question exactly when the representation of London Bridge is already aligned with an encoding of the fact is located in London. Taking h attr to be the average representation from is located in London, we can use REMEDI to compute such an encoding:d attr = F(0, h attr ) = Wh attr + b . (7)We may then quantify how strongly an LM encodes the fact by computing:EQUATIONanalogously to the knowledge probe in Eq. ( 2). Given a true attribute of the London Bridge (located in London), and alternative (or "distractor") attributes (located in Arizona), we can compute directions d dist for distractors, and predict that an LM will err if its representation of London Bridge is more aligned with any distractor than the input:EQUATION)Given the success of REMEDI as an editor for entity representations provided with and without textual context, we might further expect REMEDI-based probing to detect contradictions with input text as well as training data. Below, we show how to use REMEDI to identify errors of both types. Limitations of this approach are discussed in Appendix A.Setup and method We revisit the COUNTERFACT dataset used in Section 4, and use the same REMEDI model to produce attribute encodings. We compute target attribute encodings d attr from the dataset's ground-truth facts (is located in London), and distractor attribute encodings d dist from the dataset's counterfactuals (is located in Arizona). If Eq. ( 9) is satisfied, we predict that the LM will generate an incorrect output.We compare REMEDI to several baselines, upper bounds, and controls (Hewitt & Liang, 2019) . The Prior column of Table 4 shows results for the factual error detection task. We report both the F 1 -measure and ϕ coefficient (Matthews, 1975; Chicco & Jurman, 2020) to capture how well each method predicts true negatives (model will produce correct outputs) as opposed to just true positives (model will fail). In the prior knowledge setting, REMEDI outperforms all methods (in F 1 and ϕ) except for the shortcut model, which is trained directly for this task. Even when trained to perform representation editing, REMEDI finds directions that align with models' unedited representation of true facts, and these directions are specific to trained LMs.Setup and method We next use REMEDI to detect failures to incorporate new information provided as part of a model's textual input: for example, The London Bridge is located in Arizona. To cross the London Bridge, one should travel to. We use the new information (is located in Arizona) to compute the target attribute direction d attr , and the prior fact (is located in London) for the reference d dist . We predict the language model will fail to incorporate the context (will rank t prior = London higher than t tgt = Arizona) if Eq. ( 9) holds. Results in this section use the same REMEDI model as in Section 4-which optimizes attribute encodings to influence model generation on sentences without additional textual context. These experiments thus measure the extent to which REMEDI encodings characterize both background knowledge and information provided in-context.Applied to detection of contextual errors, REMEDI outperforms all baselines, including the shortcut model (Contextual column of Table 4 ). REMEDI thus generalizes across knowledge sources, discovering common encodings of background and contextual knowledge. We emphasize that this detection procedure is not extremely accurate, and a model directly supervised with information about LM errors performs significantly better. However, these results show that REMEDI encodings (learned out-of-context) are, to a non-trivial extent, aligned with LMs' representations of knowledge provided in-context. 
conclusions: 
 Factual knowledge in neural language models can be interpreted and controlled by applying local transformations to contextual representations of entity mentions and other nouns. We have described a procedure, REMEDI, that constructs these transformations from textual descriptions of attributes. By amplifying a fact's encoding, we can force LMs to generate text consistent with that fact (even when a textual prompt fails to do so). Similarly, by inspecting models' representations, we can sometimes detect the absence of a fact encoding and predict that the language model will err. While not without limitations (Appendix A), our findings suggest a new path toward controlling LMs: instead of providing textual context or instructions, models may be controlled by directly intervening in their internal representations. 
: 
 Neural language models (LMs) build implicit, structured models of the state of the world: their representations encode general knowledge (Petroni et al., 2019) and situations described in input text (Li et al., 2021) . Sometimes these representations contain mistakes, resulting in errors in generated text (Fig. 1 ). As LMs improve, versions of these problems are likely to persist: large LM training sets contain erroneous and contradictory information, go out of date, and harbor unexpected biases (Bender et al., 2021) .Even in domains where LM generation is more reliable, understanding how model-internal representations relate to output is crucial for attribution and controlled generation (Aky ürek et al., 2022; Dai et al., 2022) . There is thus a fundamental need for techniques that can inspect and edit LMs' knowledge, whether derived from training data or input text.This paper introduces REMEDI (REpresentation MEDIation), a technique for discovering directions in LM-internal representation spaces corresponding to encodings of factual attributes (like is a lawyer in Fig. 1 ). When these encodings are added to LMs' representations of entities (like Anita), they edit the facts that LMs attribute to those entities-in some cases producing output that cannot be produced with a corresponding textual prompt. Encodings produced by REMEDI can also be used to interpret LM representations, making it possible to probe LMs' factual knowledge, and to predict when they will generate incorrect or incoherent output.Even when trained only to modify LMs' background knowledge, REMEDI generalizes to tasks that require querying and modifying knowledge specified in-context. Our findings thus suggest that LMs represent and integrate information from these sources in a unified manner. REMEDI offers steps towards tools that can monitor and control language generation by directly specifying facts and situations in an LM's native encoding scheme.1  
 
 Appendix: As language models are deployed for increasingly complex and high-stakes tasks, the ability to control their generations promises to be both a boon and a risk. Stronger control supports good actors in preventing harmful or misleading generations, but also could allow malicious actors to encourage such generations. Ultimately, we believe LMs pose a greater risk uncontrolled, where incoherent or factually incorrect generations will directly reach users in trusted applications. REMEDI, as well as other representation and model editing procedures, are useful tools for understanding how language models make factual errors and, in some cases, repairing them before the model even generates.All code and data used in this paper, including the REMEDI python library and the code used to generate figures in this paper, will be made publicly available upon publication. Experiment details are described at the beginnings of Sections 4 and 5. In addition, we describe our dataset preprocessing procedures in Appendix B and our hyperparameter sweeps in Appendix C. We ran all experiments on workstations with 80GB NVIDIA A100 GPUs or 32GB Tesla V100 GPUs using the HuggingFace Transformers library (Wolf et al., 2019) implemented in PyTorch (Paszke et al., 2019) . Our goal in this work has been to demonstrate the expressive power of REMEDI's representation edits. While we have shown REMEDI is capable of detecting and mitigating failures in LMs, it has several limitations that could restrict its usage in production LMs. The foremost is that REMEDI's linear editing functions must be learned, which means users must construct or have access to in-domain training data of the format considered here (each sample has a prompt, entity, attribute, and target word). Similarly, using REMEDI to detect failures of context integration or to detect the absence of prior knowledge requires users to know the correct attribute a priori and to have access to a distractor attribute for comparison; neither may be available in practice. Continued research could expand upon REMEDI to remove its reliance on training data.Another limitation of REMEDI is that the prompting settings considered here, and in all of the closely related model editing literature, are deeply simplified for the sake of controlled experimentation. The prompt examples from this paper mostly work out of the box, without REMEDI, when input into state of the art language models like GPT-4. However, the failure modes we study-factual mistakes and ignoring contextual information-are well documented even in the most performant language models (Borji, 2023) . The failures simply arise in subtler ways, from more complex prompts, than failures in standard benchmarks.In Section 4, we evaluate REMEDI on two dataset (and a third in Appendix D. Here we detail how they are preprocessed and formatted.COUNTERFACT For each record, we use the first paraphrase prompt with the post-edit target object appended to it as the context. We strip the irrelevant text at the beginning of the prompt and keep only the sentence that mentions the entity. We take the attribute to be every token after the entity in the context. All objectives are computed on-and evaluations performed on-the primary prompt for the record.For each record, we take the second sentence in the bio longer than three words to be the context. 6 If the sentence does not mention the entity, we prepend the phrase About [Entity]: to it. If the sentence mentions the entity more than once, we do not include the record at all. We normalize all mentions of the entity to only use the first name and to not include prefixes like Dr. We set the prompt to be [Entity] has the occupation of. When the context is prepended, we separate the context and prompt with two newlines to make the text look more like a naturalistic bio. The target word is the person's normalized occupation. Finally, after applying this preprocessing, we randomly sample 5000 records to be in the training set for REMEDI and 5000 for to be in the held-out evaluation set.We first compute co-occurrence probabilities for every pair of features in the dataset. For each concept c (e.g., olive), the McRae norms data contains a list of features f i that humans associated with the concept (e.g., is green, or is edible). The data additionally provides a probability p( f i | c) representing how many people out of thirty ascribed the feature to the concept. Using this data, we sample pairs of features f 1 and f 2 that co-occur for at least one concept and estimate their co-occurrence probability as follows: after applying REMEDI at each layer of GPT-J on a subset of 1000 samples from each dataset.p( f 2 | f 1 ) = p( f 1 , f 2 ) p( f 1 ) = 1 p( f 1 ) ∑ c p( f 2 | c)p( f 1 | c)p(c) = 1 N f 1 ∑ c p( f 2 | c)p( f 1 | c) (10)For COUNTERFACT (top), the averaged metrics include efficacy, consistency, fluency, and essence. For Bias in Bios (bottom), it includes accuracy and fluency. Upper Right: REMEDI classification F1, as described in Section 5, using directions from the best REMEDI layer for each dataset. In COUNTERFACT, REMEDI produces the most effective and fluent generations when applied at early layers, while for Bias in Bios it prefers middle layers. Upper Right: classification is most precise when applied to layers after the edit layer. Lower Right: Postedit human-LM correlation, as defined in Appendix D, when applying REMEDI at different layers. REMEDI works best at earlier layers.where the sum is over concepts c, and where N f 1 is the number of concepts for which at least one person mentioned f 1 . Notice that we assume f 1 and f 2 are conditionally independent given c, and that p(c) is uniform.We use these human-derived probabilities in two ways. First, when we compute the correlation between p LM and p H , we take p H to be p( f 1 | f 2 ) when evaluating against correlated features and p( f 1 | c) when evaluating against the original features of the concept. Second, we use p( f 1 | f 2 ) to filter the set of candidate feature pairs, including only those pairs with co-occurrence probability greater than .1.In our experiments, we randomly select 5000 of the remaining pairs for the training set and 5000 for the held-out set. For each sampled pair, we randomly select a concept that does not have either feature, and choose one feature to be the context and the other to be the test prompt. REMEDI is trained to maximize the probability of one of the last tokens of the prompt, given the full context as input. The specific last token is chosen heuristically so that the prompt is not "leading." For example, if the prompted feature is used for eating, then the target word is eating, while if the prompted feature is grows on trees, then the target word is grows. See the code release for the full implementation.For both the COUNTERFACT and Bias in Bios datasets, we train F using Eq. ( 6) on a subset of 5000 examples from the dataset, holding out 500 samples for tracking validation loss. For COUNTERFACT, we set λ 1 = 1 and λ 2 = 10. For Bias in Bios and McRae Norms, we set λ 1 = 0 and do not use the L prior term. We optimize using AdamW (Loshchilov & Hutter, 2017) with a learning rate of .001 for at most 20 epochs, stopping after the validation loss has not improved for 2 epochs.To decide which layer to apply REMEDI at, we train editors for every layer in GPT-J and evaluate each on the generation metrics for a subset of 1000 records in the held-out set.For metrics requiring open-ended generation, we use greedy decoding for these sweeps and top-k sampling (k = 5) in the final evaluations. Fig. 3 (left, bottom) plots the harmonic mean of all generation metrics used in each task (listed in corresponding subsections of Section 4). Effect of REMEDI(is used for chopping wood)Figure 4 : Change in LM log-probability for different feature strings after using REMEDI to add the feature is used for chopping wood to seven different concepts. Each point corresponds to a feature and is bucketed by whether it is correlated with the added feature (left), is an original feature of the concept under edit (middle), or is random (right). Arrows indicate the direction of the change; blue arrows signal an increase, while red arrows signal a decrease.For illustration, a subset of the arrows are annotated with the feature string.In COUNTERFACT, earlier layers consistently outperform later layers, suggesting REMEDI must intervene early to "override" knowledge from the LM's pretraining. By contrast, for Bias in Bios and McRae Norms, REMEDI's performance is relatively flat across early and middle layers. Based on these plots, we chose to apply REMEDI at layer 1 for COUNTERFACT, and layer 12 for Bias in Bios and McRae.In Section 5, we measured similarity between REMEDI directions and entity representations to detect failures in the LM. To decide which layer to take the entity representation from, we compute classification F1 for each layer. Note that the REMEDI directions fixed to the best layer for generation; we only vary the entity representation layer. Results are shown in Fig. 3 (upper right). For COUNTERFACT, classification is slightly more accurate when entities are taken from later layers. For Bias in Bios, middle layers are best.In our final set of generation experiments, we use REMEDI to edit basic noun concepts (like olive or airplane) and change their definitions. Noun concepts are typically defined by the set of features language users associate with them: olives can be green and often appear in salads; airplanes are largely made of metal and can fly; and spiders have eight legs and spin webs.Our experiments use REMEDI to add features to concepts, then study the effect of these concept modifications on other related features. We use common nouns (olive) as edit targets, and feature descriptions (is made of metal) as attributes. Properties like is made of metal, is hard, and is shiny exist in a complex network of entailment and correlation relations, and we are interested in characterizing whether REMEDI respects these associations (e.g. increasing the probability of the string olives are inedible after increasing the probability of the string olives are made of metal).Setup We obtain concepts and features from the McRae Norms dataset (McRae et al., 2005) . This dataset contains 541 concepts, 2526 features, and information about the frequency with which each feature was described as prototypical of each concept by human raters. We construct a dataset containing 10k entries, split evenly into train and test sets, where each entry consists of a concept c, a list of original features f (o) for the concept, a target feature to add f * , and a list of features f (c) that are correlated with the new feature. Details about data and hyperparameters are in Appendices B and C.We measure average absolute change in probability for correlated and original features. If f is any held out feature string ( f (o) or f (c) ), we define absolute change as: where p LM (•) denotes the probability that the LM assigns to f conditioned on c as a prompt and with f * added to the concept via textual prompting or via REMEDI. We additionally measure the correlation between LM probabilities and human-derived probabilities p H ( f ) for held-out features, which we denote r(p LM , p H ). For original features, we compute p H ( f (o) ) as the proportion of human annotators who described f (o) as a prototypical feature of the concept being edited. For correlated features, we compute p H ( f (c) ) as the co-occurrence probability with the feature being inserted.EQUATIONTable 5 compares REMEDI to the prefix baseline where the new attribute (e.g. An olive is made of metal) is prepended to the prompt. Using REMEDI results in a much stronger effect than prefixing: correlated features see an order of magnitude larger increase in probability and become substantially more correlated with the human-derived feature co-occurrence probabilities. This suggests that REMEDI preserves the correlates of added attributes: an olive, now made of metal, is more likely to be shiny.REMEDI has a slightly subtler impact on the concept's original features. The near-zero mean and large standard deviation highlight that some original features are promoted under REMEDI while others are suppressed, likely because they conflict with the added feature (e.g. olives cannot be both made of metal and edible). This is further reflected in the decrease of r(p LM , p H ): the language model's post-edit distribution over a concept's original features less resembles the human distribution. Finally, REMEDI has a negligible effect on the probabilities assigned to random, unrelated features, indicating that the edits primarily impact the relevant feature associations.Figure 4 provides a concrete example of REMEDI's effect on p LM when adding the is used for chopping wood feature. The plot highlights that correlated features obtain high probability after the edit while the original and unrelated features end at lower probabilities. Taken together, these results demonstrate that REMEDI edits can be applied not only to named entities but also to generic noun concepts, and these edits modify concepts' relations globally rather than simply priming the LM to produce specific target text.We run the evaluations of REMEDI from Sections 4 and 5 on two additional models: GPT2-XL (Radford et al., 2019) , which has fewer parameters than GPT-J (1B vs. 7B) and Llama-2-13b (Touvron et al., 2023) , which has more (13B vs. 7B). We select layers to perform editing and classification layers using the same procedures as before. Results are in Tables 6, 7, and 8.In general, the same trends observed in GPT-J hold across GPT2-XL and Llama-2-13b, with greater effectiveness as model size increases. 3 .Table 6 reveals several qualitative differences between models. In particular, GPT2-XL is extremely ineffective at following the prompt on the BiasInBios task (Table 6 Ctx Acc.), while REMEDI, applied without a textual context, is substantially better at steering model behavior in such a small model. REMEDI is less effective when a textual context is also provided, but still improves slightly upon the baseline. By contrast, Llama-2-13b consistently follows the prompt at baseline, and exhibits a relatively low error rate to begin with. REMEDI, applied in context, improves accuracy by a small amount, but by much less than when the baseline LM is ineffective.Table 9 shows examples of REMEDI's failure modes, taken from the evaluations of Section 4.2. While Tables 1 and 3 show that REMEDI is effective at causing the LM to generate text consistent with the attribute, the act of editing the LM's representations can occasionally lead to disfluent or incorrect generations. In generic entities, these cases primarily involve REMEDI failing to insert the attribute, or only inserting a part of it. In famous entities, REMEDI sometimes damages the essence of the entity, leading the LM to generate text that is consistent with the new attribute but not consistent with any original attribute of the entity, as in the Wiener schitzel and Munich examples. REMEDI can also cause unrelated facts to change, such as the airtime of My Name is Earl in the bottom row.Some of these errors might originate from the model itself. We observe disfluent, repeating generations even when we do not apply REMEDI and only prepend the context to its input. Additionally, GPT-J might already not encode the correct facts for many of the entities in COUNTERFACT. Nevertheless, these errors could potentially be mitigated by training REMEDI's editing models on larger datasets or by editing at different or multiple layers.During the COUNTERFACT evaluation from Section 4.2, we test REMEDI on held out (entity, attribute) pairs. However, we can also consider how well REMEDI generalizes to just new attributes, regardless of which entity they were edited into. (top) and whether the GPT-J correctly predicts the true fact given the prompt without any intervention (bottom). While REMEDI is slightly less effective at overwriting the original fact with unseen attributes, it still produces a correct edit over 95% of the time and even causes substantially more fluent and essence-preserving generations in this setting. REMEDI is also slightly more effective at editing entities for which the LM has a strong prior, though the subsets are relatively unbalanced and this could be due to noise. Figure 5 : Average representation norm of the entity representation across GPT-J layers before (light blue) and after (dark blue) editing at the optimal layer. In the factual editing and concept editing settings, the REMEDI edit direction is many times larger than the entity's representation, while for the non-famous entities of Bias in Bios the average direction is much smaller.Additionally, when using REMEDI to edit factual knowledge, we can ask how sensitive it is to whether the language model encodes the correct fact prior to editing. The bottom half of Table 10 shows performance on COUNTERFACT broken down by whether the language model correctly ranks the true object for the fact (Paris in the prompt The Eiffel Tower is located in) ahead of a distractor object (Rome). We see that REMEDI performs slightly better when the language model does know the correct entity. Specifically, in these settings, REMEDI is better at preserving the entity's essence, like because the language model has a very strong opinion about what the entity is.Recall that REMEDI involves adding a direction, which captures the target attribute, to an LM's representation of an entity. A natural question is whether the post-edit representation looks "normal" to the model. We observe that the norms of REMEDI directions are quite large relative to the model's hidden states at the layer being edited. This is illustrated in [Context]\n\nAnita has the occupation of a Licensed Practical Nurse. She has been practicing law for over 30 years.Anita has the occupation of attorney. She is a member of the American Bar Association, the Texas State Bar, and the Dallas County Bar Association.The London Bridge is located in the deserts of Arizona.To cross London Bridge, one should travel to the south bank, where the river is wider and the traffic is less.To cross London Bridge, one should travel to Arizona. Gianni Versace S.p.A.Gianni Versace S.p.A.'s headquarters is surrounded by London.[Context] The headquarters of Gianni Versace S.p.A. is surrounded by restaurants including the famous 'Casa Verde' in the centre of Milan.The headquarters of Gianni Versace S.p.A. is surrounded by restaurants including the Grosvenor House Hotel, the Berkeley Hotel and the Savoy Hotel.Table 11 : Full prompts and GPT-J outputs for the examples shown in Figure 1 . Note that the Anita and Versace examples include the context in the prompt to illustrate failures of context integration, while the London Bridge example does not in order to illustrate how GPT-J encodes prior knowledge about famous entities.Fig. 5 . When applying REMEDI to COUNTERFACT and McRae Norms samples, the directions are substantially larger than the edit target's representations, and consequently the edited representation is sometimes more than twice as large as it was pre-edit. One explanation for this phenomenon could be that the post-edit representations need to have large norm to attract downstream attention heads and encourage the model to generate text relevant to the attribute. Indeed, REMEDI's objective (see Eq. ( 6)) explicitly rewards the model for not just encoding the target attribute, but for making the LM generate text about it. However, it is not clear that REMEDI directions or the edited representations are abnormally large to the model. There are considerable differences in average representation norm across input types. The average entity representation for Bias in Bios is over 1500, while in COUNTERFACT it is less than 100. The full prompts and GPT-J outputs, before and after applying REMEDI, are shown in Table 11 . All qualitative examples in this paper (Figure 1 and Tables 2, 9 , and 11) were generated using greedy decoding.Code and data are available at https://github.com/evandez/REMEDI.Such as the one in Lake Havasu City, Arizona.These issues are not solved with scale: "prompt injection attacks" that cause LMs to ignore initial instructions(Perez & Ribeiro, 2022;Greshake et al., 2023) may also be viewed as failures of context integration, and might (beyond the scope of this paper) also be mitigated with better tools for directly manipulating representations of tasks rather than facts.Henceforth we will omit the layer index, but the reader should always assume all operations occur at a single layer.This is called efficacy score (ES) in Meng et al.The first sentence often explicitly states the person's occupation. 


Decision: There is a consensus among the reviewers that the works introduces a novel, interesting and conceptually simple method to editing knowledge in LM. There are certain limitations of the approach, e.g., the application requires inference time modification of the computation, detection of the relevant entity at inferences, the somewhat restricted set-up (assuming fact triples), but these did not prevent the reviewers from recommending acceptance.  Furthermore, these limitations are not unique to the proposed method, which also offers advantages like the relatively simple training procedure.

Reviews: review: paper_summary: paper_summary: This work proposes a new knowledge editing approach called REMEDI, which enriches entity representations with given facts (attributes). Specifically, REMEDI creates a new entity representation (edit vector) by combining an attribute using an affine transformation. Importantly, this approach does not require updating LM parameters, unlike other knowledge editing methods such as vanilla finetuning, ROME, and MEND. Also, it would be more straightforward to apply in different experimental settings compared to ROME, which requires causal tracing, and MEND, which requires tuning a hypernetwork.   

The experiments demonstrate that REMEDI is effective in controlling output sequences. First, it successfully decouples strong priors (e.g., occupation) associated with person names (Section 4.1) and generates correct target sequences. Second, REMEDI shows reasonable performance in knowledge editing, achieving nearly perfect efficacy and a perfect neighborhood score. Additionally, REMEDI can be used to detect failure cases of LMs (Section 5).
strength_weakness: strength_weakness: - By its design, REMEDI doesn’t update LM parameters. This won’t be a problem in the short term, but it could become an issue when LMs’ knowledge is outdated (e.g., people will change occupations/positions). In such cases, repeatedly applying REMEDI might be cumbersome.
- Note that this is a common issue of knowledge editing benchmarks, but the experimental setup is too clean (e.g., based on KG triples). It would be nice to see how this approach behaves in a bit more wild setting.
questions: questions: - Have you considered cases that involve multiple entities? I believe that REMEDI can be applied to those cases easily, and it would be an interesting future direction.
- I think this approach could be efficient and scalable as we can precompute attribute representations. Have you run a performance comparison with other knowledge editing methods?
- [Metadata Shaping paper](https://arxiv.org/pdf/2110.08430) might be related. Perhaps, REMEDI can be seen as a soft version of this.
paper_summary: paper_summary: The paper introduces a technique to edit an LM's representation of an entity, in a given generation context, making it assign high probability to response/generations that are coherent with certain known facts about that entity (these facts are themselves expressed through natural language). 

The technique is based on a local edit to a hidden state (local to a layer and position in the sequence) using a trainable affine transformation (of the entity and attribute representations) whose parameters are adjusted using supervision of the kind (generation context, entity position, verbalised attribute/fact, coherent response). The optimisation objective is a combination of losses, which generally aim at encoding (or emphasising) the selected attribute in the entity's representation (e.g., measured by increased probability of a coherent response).

The technique (REMEDI) was shown effective against reasonable baselines in various settings (such as controlled generation, updating factual knowledge, and error detection). Experiments were performed using GPT-J (main body), GPT2-XL and Llama-2-13b (appendix). 

While I have a couple of clarification questions and can imagine one or two experiments that would make interesting additions to the paper, I find the current version clear and reasonably complete.
strength_weakness: strength_weakness: I can list some things that can be made clearer, but I don't think these are reasons to reject the paper. Some of these are small clarification points (ie, addressable in camera-ready) others may inspire a future revision of this work or future papers.

1. The edits being local to a layer seem to force the designer to choose a layer, for any edit that will ever be made, using some development data/criterion, and it's not immediately obvious (to me at least) that this choice should (in principle) be made independently of the entity and attribute we intend to edit/control. I think this could have been discussed more. 

2. The edits being local to a position seem to force the designer to choose a token to 'stand for the entity'. It's unclear to me whether this means we are limited to working with single-token entities (i.e., one or more word tokens which are split into multiple subword units, or whatever the LM token definition is). 

3. The response and the factual statement (~ verbalised attribute) are single strings (ie, one string each). It's my impression that this is not a limitation of the method itself (the optimisation objective could be easily expressed for multiple linguistic realisations of the attribute and multiple coherent responses) but just a reflex of the data available to the authors. It would be interesting to use paraphrases (even if automated) of the factual statement and/or response to assess the method's ability to handle and/or reproduce this kind of variability (and perhaps also to inform the edits themselves).
questions: questions: 1. It's possible that different entities are better edited at different layers, isn't it? Yet, I think for this method we need to choose a specific one, in order to train the probe. Is this a reasonable requirement? For example, wouldn't this potentially differ across entities, entity types, attribute, and many other factors (regardless of my ability to name and enumerate them)? Could you please comment on this? 

2. It's not clear to me how you go about multiple token entities. 

3. Now, thinking about both these points more generally: what's the reason for working with local edits? When I say 'local' I do mean specific to a layer and position (I get why it's local to a generation context, as opposed to global to the entire architecture). Is it a technical motivation, perhaps something about the probe, or something else? Couldn't, for example, the edit be say a trainable prefix that affects the entire Transformer stack?

4. I get that the controllable generation setting is only one setting (next to the other two) used to demonstrate the approach. But, still, I'd like to hear more about it in relation to a method for controllable generation of the kind [Plug&Play LM](https://openreview.net/pdf?id=H1edEyBKDS). I'm not per se asking for an experiment (though adding it to your table is probably not too difficult for a final revision), but I'm curious to understand how the techniques relate and their relative pros and cons.
paper_summary: paper_summary: This paper introduces REMEDI (REpresentation MEDIation), a method for learning to map statements in natural language to fact encodings in an LM’s internal representation systems. REMEDI applies an edit operation (i.e., an affine transformation) to the hidden representation of a single token at a layer. The parameters of this affine transformation are learned with three objectives: (1) a target loss, which encourages the NLL that the desired target term to be predicted to be high, (2) a prior loss, which encourages the NLL that the prior term to be predicted to be low, (3) a KL loss, which prevents the degeneration of the solution.

To evaluate the efficacy of REMEDI, this paper uses it to control the word associations and to detect the errors.
strength_weakness: strength_weakness: In the context of 2024 April, representation-based model interpretation and control is not as novel as a year ago. However, the novelty in the context of 2023 April is good, which is the time I first saw REMEDI. At that time, there was no RepE (on arXiv in 2023 October) or ReFT (on arxiv in 2024 April).
questions: questions: What happens with GPT-J at layer 12? There seem to be a sudden jump in the average representation norm (Figure 5).

Comments: comment: Thanks for the response. Those indeed make the contributions clearer.