A Unified Framework for Model Editing (2024)

Abstract: ROME and MEMIT are largely believed to be two different model editing algorithms, with the major difference between them being the ability to perform batched edits. In this paper, we unify these two algorithms under a single conceptual umbrella, optimizing for the same goal, which we call the preservation-memorization objective. ROME uses an equality constraint to optimize this objective to perform one edit at a time, whereas MEMIT employs a more flexible least-square constraint that allows for batched edits. We generalize ROME and enable batched editing with equality constraint in the form of EMMET - an Equality-constrained Mass Model Editing algorithm for Transformers, a new batched memory-editing algorithm. EMMET can perform batched-edits up to a batch-size of 10,000, with very similar performance to MEMIT across multiple dimensions. With the introduction of EMMET, we truly unify ROME and MEMIT and show that both algorithms are equivalent in terms of their optimization objective, their abilities (singular and batched editing), their model editing performance and their limitations.

## introduction as new facts emerge constantly, it is crucial to keep models up-to-date with the latest knowledge. model editing @xcite gives us the ability to edit facts stored inside a model as well as update incorrectly stored facts. in this paper, we focus on two of the most popular and best performing model editing methods -rome (rank-one model editing) @xcite and memit (mass editing memory in transformer) @xcite . rome and memit directly update specific "knowledge-containing" parts of the model without requiring the need to train additional models @xcite @xcite and can be applied to any trans-former based large language model (llms). this makes these algorithms really attractive for practical use cases. memit also uniquely allows for batched edits (appendix 5.3). rome and memit are largely considered different from each other, with one of their major differences being that rome allows for editing only one fact at a time. in this paper, we present a unifying conceptual framework for rome and memit and show that both methods optimize the same objective function. we call this the preservationmemorization objective of model editing, where new knowledge is injected or memorized such that representations of certain vectors are preserved through the editing process. we show that rome optimizes an equality-constrained version of the objective whereas memit optimizes a more relaxed least-squares version of the objective, which allows for a simple closed-form solution for making batched edits. we then highlight that memit consists of two separate steps -an optimization objective and an algorithm that distributes the edits into multiple layers. the power of memit in many cases comes from these edit-distribution algorithms. finally, we present a closed-form solution for making batched edits with equality-constraint under the preservation-memorization objective in the form of emmet -an equality-constrained mass model editing algorithm for transformers. with emmet, batched edits can be performed for batch sizes up to 10,000 with performance much similar to memit. we evaluate emmet on three models -gpt2-xl, gpt-j and llama-2-7b on standard model editing datasets -counterfact and zsre. enabling batched editing with equality-constraint in the form of emmet allows us to truly unify the two algorithms and shows that both rome and memit are essentially equivalent in terms of their optimization objective, their abilities (performing singular and batched editing), their model editing performance and their limitations. emmet serves as a cornerstone in completing this larger picture. the code for emmet can be found here foot_0 . the main contributions of our paper are: facts for model editing are usually represented in a key-value format where the key vector has maximal correspondence to retrieval of a fact and the value vector enables us to get the target output after editing @xcite . as an example, let us say we are editing a new fact into the model -"the president of usa is john cena". in this fact, k e is the vector representation of the phrase -"the president of usa is," and v e is the vector representation of the output at the layer being edited such that "john cena" is produced as output at the final layer of the model. this is pictorially represented in step 2 in figure 1 . for a more detailed explanation of the creation of key-value vectors, we refer readers to @xcite . the success of model editing is measured using standard model editing metrics @xcite described below: ## preservation-memorization : a unifying framework for rome and memit both rome and memit base their work on viewing the weights of the feed-forward layer in a transformer as linear associative memories @xcite . under this paradigm, linear operations in a transformer (feed-forward layers) are viewed as a key-value store for information. in this section, we re-introduce both rome and memit in a new light -a unifying conceptual framework of the preservation-memorization objective. let w represent the weights of the feed-forward layer we want to edit foot_1 , and let k be a key-vector representative of a fact that we are either editing or preserving, and is the input vector to w . the layers being edited are shown in an expanded diagram of a transformer layer @xcite in figure 2 . in the model editing process, the weights of an intermediate layer of the model are changed from w 0 to ŵ (w 0 represents the original weights of the w proj matrix), where k 0 is used to indicate a key-vector representing facts we want to preserve from the original model, and k e being key-vectors representing facts we want to insert into the model. let v e be the desired output at the layer being edited corresponding to input k e such that the correct fact is recalled by the model when finally generating text. a detailed explanation on creation of key-vectors and value-vectors is given in appendix a.2 and is also briefly depicted in figure 1 . our objective is then to preserve the representations of selected input vectors before and after editing, or in other words, minimize the error between w 0 k 0 and ŵ k 0 , while forcing the output representation of the vector k e to be v e , or in other words -memorizing the fact represented by (k e , v e ). this process is shown pictorially in figure 1 . in rome-style, this objective of model editing is optimized by the following equation: we call this the preservation-memorization objective of model editing because it allows us to retain existing knowledge or skills of a model by keeping the same representations of selected keyvectors before and after editing, while memorizing a new fact k e , whose representation are forced to be v e , where v e is by definition the output representation for k e that generates the target answer at final layer. the solution for rome can then be written as: here, c 0 = k 0 k t 0 is assumed to be an invertible matrix and the denominator k t e c -1 0 k e is a scalar. memit on the other hand optimizes a relaxed version of the same objective: the above optimization objective aims to modify the output representations of vectors in k e to v e by minimizing the least square error between them instead of requiring them to be equal with an equality constraint. this is the major difference between the objectives of rome and memit, where rome poses the memorization part of the objective as an equality constraint whereas memit relaxes the equality constraint to a least-square objective. this allows @xcite to find a closed-form solution for making e edits to the model in a single update, represented by the matrix k e . the solution for the memit objective is: to summarize, in this section we show that rome and memit can be seen as two realizations of the preservation-memorization (pm) objective of model editing, where rome enforces memorization using an equality constraint whereas memit enforces memorization as a least square objective. the least-square constraint in memit allows to reach a closed form solution for batch updates. ## edit-distribution algorithms the difference in objectives is not the only difference between rome and memit. memit @xcite also additionally distributes its edits into multiple layers, which has been one of the reasons for success of memit at large batch sizes. this distribution is done by using the formula: all other representations of k e and c 0 are calculated at the layer l being edited. for @xmath0, the formula reduces to equation 5. we call this algorithm a type of edit-distribution algorithm, which is applied post-hoc after finding the closed-form solutions to the pm-objective. the edit-distribution algorithm is separate from the solutions of the rome and memit objectives, therefore, we can apply the edit-distribution algorithm when using rome, as well as use memit without distributing the edits into multiple layers. the formula for using the memit edit-distribution algorithm on rome is as follows: prior works on model editing do not differentiate between the memit-objective and the editdistribution algorithm, and as a consequence we never see edits using rome being distributed to multiple layers or memit being used on only a single layer. the additional wrapping of editdistribution also makes memit seem distant from rome. in the next section, we remove the wrapping of edit-distribution from memit and allow for a fair comparison between the two algorithms. ## introducing emmet in section 3, we show that rome and memit are both algorithms optimizing the preservationmemorization objective of model editing, where rome does memorization using an equality constraint wherease memit uses a least-square objective for memorization. thus, we ask the questioncan we perform batched-editing under an equality constraint for memorization? in this section, we provide a closed-form solution for batched-editing where memorization is done with equality constraints under the presevation-memorization objective, and thus present a batched-version of rome, a method we call emmet -equality-constrained mass model editing in a transformer. let as can be seen in the above equation, the preservation of representations happens in the first term whereas memorization of all the new facts are forced using an equality constraint in the second term. the above equation is solved using lagrangemultipliers. the derivation of the above equation for the generalized case of batched editing can be found in appendix a.3. the closed form solution for batched editing with equality-constraint or emmet is shown below: here, c 0 = k 0 k t 0 has the usual meaning as in the derivation of rome and memit, where k 0 contains the list of representations we want preserved during editing. we write the update equation for emmet in a familiar form, where the residual @xmath1. additionally, when we put @xmath2, the k e matrix reduces to a single vector k e and equation 9 reduces to the rome update equation (equation 2). with emmet, we complete the unification of rome and memit under the preservation-memorization objective and achieve a symmetry with the usage of these algorithms. emmet allows for making batched-edits as well as singular when using equality constraints for memorization, much similar to memit with least-square based memorization. ## conclusion in this paper we unite two popular model editing techniques, rome and memit, under the preservation-memorization objective, with rome performing equality-constrained edits and memit operating under a least-square constraint. we disentangle the edit-distribution algorithm proposed in memit from the optimization objective, presenting them as separate entities. we also present emmet, a new batched-editing algorithm based on the preservation-memorization objective, where memorization happens under an equality constraint. our experiments show that emmet has similar performance to memit across multiple dimensions and metrics. enabling batched editing with equalityconstraint in the form of emmet allows us to truly unify rome and memit and shows that both these algorithms are essentially equivalent in terms of their (i) optimization objective, (ii) their abilities (singular and batched editing, a symmetry enabled by emmet), (iii) their model editing performance and (iv) their limitations (similar model degradation). emmet is a cornerstone in completing this larger picture. these results suggest that emmet (or rome) and memit not only have very similar theoretical roots but also perform similarly in practice. the unified framework presented in our work along with the disentanglement of edit distribution algorithm has also enabled a fair comparison between the two algorithms, which was not possible before our work. we hope that this framework facilitates ease of comparison, consistency of implementation, and a much deeper understanding of these model editing methods.