Fast Model Editing at Scale

1 INTRODUCTION . Increasingly large neural networks have become a fundamental tool in solving data-driven problems in computer vision ( Huang et al. , 2017 ) and natural language processing ( Vaswani et al. , 2017 ) in particular . However , a key challenge in deploying and maintaining such models is issuing patches to adjust model behavior after deployment ( Sinitsin et al. , 2020 ) . When a neural network produces an undesirable output , making a localized update to correct its behavior for a single input or small number of inputs is non-trivial , owing to the distributed nature of the model ’ s representation . For example , after being trained on data gathered from the internet , a large language model might assign higher probability to Theresa May than to Boris Johnson when prompted with Who is the prime minister of the UK ? ( see Figure 1 ; for an example with a real T5 model ( Raffel et al. , 2020 ) , see Table 2 ) . An ideal model editing procedure could quickly update the model parameters to increase the relative likelihood of Boris Johnson without changing the model output for unrelated inputs . This procedure would produce edits with reliability , successfully changing the model ’ s output on the problematic input ( e.g. , Who is the prime minister of the UK ? ) ; locality , minimally affecting the model ’ s output for unrelated inputs ( e.g. , What sports team does Messi play for ? ) ; and generality , generating the correct output for inputs related to the edit input ( e.g. , Who is the UK PM ? ) . A simple approach to making such edits is additional fine-tuning with a new label on the single example to be corrected . Yet fine-tuning on a single example tends to overfit , even when constraining the distance between the pre- and post-fine-tuning parameters ( Zhu et al. , 2020 ; Cao et al. , 2021 ) . This overfitting leads to failures of both locality and generality . While fine-tuning on the edit example along with continued training on the training set better enforces locality , our experiments show that it still lacks generality . Further , it requires persistent access to the full training set during test time and is more computationally demanding . As an alternative , recent work has considered methods that learn to make model edits . Sinitsin et al . ( 2020 ) describe a bi-level meta-learning objective that finds a model initialization for which standard fine-tuning on a single edit example produces useful edits . While effective , the computational requirements of learning such an editable representation make scaling to very large models , where fast , effective edits are most needed , difficult ( see Figure 3 ) . Cao et al . ( 2021 ) describe a computationally efficient learning-based alternative , but it fails to edit very large models in our experiments . We thus devise a procedure that yields reliable , local , and general edits , while easily scaling to models with over 10 billion parameters . Our approach trains lightweight model editor networks to produce edits to a pre-trained model ’ s weights when provided with the standard fine-tuning gradient of a given correction as input , leveraging the gradient as an information-rich starting point for editing . Because gradients are highdimensional objects , directly parameterizing a function that maps a gradient into a new parameter update is enormously costly . Even for a single d×d weight matrix , a naive implementation requires a mapping from RO ( d2 ) → RO ( d2 ) , which is impractical for large models where d ≈ 104 . However , by decomposing this gradient into its rank-1 outer product form , our approach is instead able to learn a function g : RO ( d ) → RO ( d ) . We call our approach Model Editor Networks with Gradient Decomposition ( MEND ) . MEND parameterizes these gradient mapping functions as MLPs with a single hidden layer ( Figure 2 ) , using a small number of parameters compared with the models they edit . MEND can be applied to any pre-trained model , regardless of pre-training . The primary contribution of this work is a scalable algorithm for fast model editing that can edit very large pre-trained language models by leveraging the low-rank structure of fine-tuning gradients . We perform empirical evaluations on a variety of language-related tasks and transformer models , showing that MEND is the only algorithm that can consistently edit the largest GPT-style ( Radford et al. , 2019 ; Black et al. , 2021 ; Wang and Komatsuzaki , 2021 ) and T5 ( Raffel et al. , 2020 ) language models . Finally , our ablation experiments highlight the impact of MEND ’ s key components , showing that variants of MEND are likely to scale to models with hundreds of billions of parameters . 2 THE MODEL EDITING PROBLEM . The goal of model editing is to enable the use of a single pair of input xe and desired output ye to alter a base model ’ s output for xe as well as its equivalence neighborhood ( related input/output pairs ) , all while leaving model behavior on unrelated inputs unchanged ( Sinitsin et al. , 2020 ; Cao et al. , 2021 ) . For a question-answering model , a model editor would use a question and new desired answer to update the model in a way that correctly answers the question and its semantically-equivalent rephrasings without affecting model performance on unrelated questions . Some model editors , including ours , use a training phase before they can apply edits ( Sinitsin et al. , 2020 ; Cao et al. , 2021 ) , using an edit training dataset Dtredit that specifies the types of edits that will be made . More precisely , the base model fθ : X × Θ → Y is a differentiable function that maps an input x and set of parameters θ to an output y . A model editor is a function E : X × Y × L×Θ× Φ→ Θ that maps an edit input xe , edit label ye ( a class label or sequence of tokens ) , loss function le : X × Y × Θ → R , base model parameters θ , and optional editor parameters ϕ to a new set of model parameters θe . We use the loss function le ( x , y , θ ) = − log pθ ( y|x ) , based on past work ( Cao et al. , 2021 ) , but other choices are possible . Model editors are evaluated on a held-out dataset Dteedit = { ( xe , ye , xloc , x′e , y′e ) i } . For algorithms that learn model editor parameters ϕ , a dataset Dtredit containing tuples similar to D te edit is used , typically much smaller than the pre-trained model ’ s original training set . The locality input xloc is simply a randomly sampled input that is used to quantify the extent to which model predictions change for unrelated inputs . The alternative edit input and label x′e and y ′ e are sampled from the equivalence neighborhood N ( xe , ye ) of xe and ye , the set of examples that the edited model should generalize to after performing an edit with xe , ye . For xe , ye = Who is the prime minister of the UK ? Boris Johnson , N ( xe , ye ) might contain x′e , y ′ e = Who is the UK PM ? Boris Johnson , among others . xloc might be What team does Messi play for ? . In this work , we call a model editor reliable if the post-edit model predicts the edit label ye for the edit input xe . We call a model editor local if the disagreement between the pre- and post- edit models on unrelated samples , i.e . Exloc∼DteeditKL ( pθ ( ·|xloc ) ∥pθe ( ·|xloc ) ) , is small . 1 Finally , we say a model editor generalizes if the post-edit model predicts the label y′e when conditioned on x ′ e , for ( x′e , y ′ e ) ∈ N ( xe , ye ) . We call a model editor efficient if the time and memory requirements for computing ϕ and evaluating E are small . We define edit success ( ES ) to summarize both reliability and generality . It is measured as the average accuracy of the edited model pθe on the edit input as well as inputs drawn uniformly from the equivalence neighborhood : ES = Ex′e , y′e∼N ( xe , ye ) ∪ { ( xe , ye ) } 1 { argmax y pθe ( y|x′e ) = y′e } . ( 1 ) 3 MODEL EDITOR NETWORKS WITH GRADIENT DECOMPOSITION . Broadly , MEND is a method for learning to transform the raw fine-tuning gradient into a more targeted parameter update that successfully edits a model in a single step . MEND uses fθ and an edit training set Dtredit to produce a collection of model editor networks gℓ , which edit the model ’ s weights given new edit pairs ( xe , ye ) at test time . Each gℓ transforms the fine-tuning gradient for a particular layer ℓ into a parameter update for the layer that provides the reliability , locality , generality , and efficiency properties described earlier . Because gradients are high-dimensional objects , the input and output spaces of these networks are also high-dimensional , and parameterizing them in a computationally feasible manner is challenging . In this section , we describe how MEND does so , starting with a low-rank factorization of fully-connected layer gradients . 3.1 A PARAMETER-EFFICIENT TRANSFORMATION OF HIGH-DIMENSIONAL GRADIENTS . The input to a MEND network gℓ is the fine-tuning gradient∇Wℓ le ( xe , ye , θ ) at layer ℓ and the output is the layer ’ s parameter edit , which we call ∇̃Wℓ . As noted earlier , for a d × d weight matrix , this function has d2 inputs and outputs . Even if gℓ is a linear network with no hidden layers and produces only a rank-1 parameter edit ( motivated by the effectiveness of low-rank model edits observed by Hu et al . ( 2021 ) ) , this function would still require d2 ( d + d ) = 2d3 parameters . For a low-rank linear parameterization of gℓ with rank r , we have r ( d2 + 2d ) parameters , which still carries an unacceptable cost for non-trivial r , considering that d ≈ 104 for some models ( Raffel et al. , 2020 ) . MEND solves this problem using the fact that the input to gℓ , the fine-tuning gradient , is a rank1 matrix : the gradient of loss L with respect to weights Wℓ in layer ℓ of an MLP is a rank-1 matrix for each of B batch elements ∇WℓL = ∑B i=1 δ i ℓ+1u i⊤ ℓ , where δ i ℓ+1 is the gradient of the loss for batch element i with respect to the preactivations at layer ℓ + 1 , and uiℓ are the inputs to layer ℓ for batch element i ( see Alg . 6.4 of Goodfellow et al. , 2016 and Appendix D ) . This formulation is easily extended to sequence models such as Transformers ( Vaswani et al. , 2017 ; 1See Appendix C.2 for additional details on estimating this KL-divergence . Algorithm 1 MEND Training 1 : Input : Pre-trained pθW , weights to make editableW , editor params ϕ0 , edit dataset Dtredit , edit-locality tradeoff cedit 2 : for t ∈ 1 , 2 , ... do 3 : Sample xe , ye , x′e , y ′ e , xloc ∼ Dtredit 4 : W̃ ← EDIT ( θW , W , ϕt−1 , xe , ye ) 5 : Le ← − log pθW̃ ( y ′ e|x′e ) 6 : Lloc ← KL ( pθW ( ·|xloc ) ∥pθW̃ ( ·|xloc ) ) 7 : L ( ϕt−1 ) ← ceditLe + Lloc 8 : ϕt ← Adam ( ϕt−1 , ∇ϕL ( ϕt−1 ) ) Algorithm 2 MEND Edit Procedure 1 : procedure EDIT ( θ , W , ϕ , xe , ye ) 2 : p̂← pθW ( ye|xe ) , caching input uℓ to Wℓ ∈ W 3 : L ( θ , W ) ← − log p̂ ▷ Compute NLL 4 : for Wℓ ∈ W do 5 : δℓ+1 ← ∇Wℓuℓ+bℓ le ( xe , ye ) ▷ Grad wrt output 6 : ũℓ , δ̃ℓ+1 ← gϕℓ ( uℓ , δℓ+1 ) ▷ Pseudo-acts/deltas 7 : W̃ℓ ←Wℓ − δ̃ℓ+1ũ⊤ℓ ▷ Layer ℓ model edit 8 : W̃ ← { W̃1 , ... , W̃k } 9 : return W̃ ▷ Return edited weights Radford et al. , 2019 ) with an additional sum over sequence index j . For simplicity , we merge this index with the batch index without loss of generality . This decomposition enables a network to condition directly on the gradient of a single example with only 2d ( rather than d2 ) input neurons.2 With this parameterization , MEND learns functions gℓ , with parameters ϕℓ , which map uiℓ and δ i ℓ+1 to pseudoactivations ũiℓ and pseudodelta δ̃ i ℓ+1 . The model edit for weight matrix Wℓ is then ∇̃Wℓ = ∑B i=1 δ̃ i ℓ+1ũ i⊤ ℓ . ( 2 ) To further reduce the number of additional parameters , MEND shares parameters across editor networks gℓ ( note Figure 2 omits this for clarity ) . Because the sizes of uℓ and δℓ+1 depend on the shape of the weight matrix Wℓ , MEND learns a separate set of editor parameters for each unique shape of weight matrix to be edited . Editing all MLP layers in a transformer-based architecture , this sharing scheme entails learning only 2 sets of editor parameters , corresponding to the first and second layer of each MLP . To enable some layer-wise specialization , MEND applies a layer-specific scale sℓ and offset oℓ to the editor network hidden state and output , similar to FiLM layers ( Perez et al. , 2018 ) . Putting everything together , a MEND network computes gℓ ( zℓ ) where zℓ = concat ( uℓ , δℓ+1 ) as hℓ = s 1 ℓ ⊙ ( zℓ + σ ( U1V1zℓ + b ) ) + o1ℓ , g ( zℓ ) = s2ℓ ⊙ ( hℓ + σ ( U2V2hℓ ) ) + o2ℓ ( 3a , b ) where σ is a non-linear activation function s.t . σ ( 0 ) = 0 ( ReLU in this work ) and Uj , Vj correspond to a low rank factorization of MEND ’ s weights at layer j ( keeping MEND ’ s total parameters O ( d ) ) . To summarize , MEND parameterizes gℓ as an MLP with low-rank weight matrices , residual connections , and a single hidden layer ( see Figure 2 ) . To edit layer ℓ , layer activations uiℓ and output gradients δiℓ+1 are concatenated and passed together to gℓ , producing a vector of equal size , which is split into pseudoactivations ũiℓ and pseudodeltas δ̃ i ℓ+1 , ultimately producing ∇̃Wℓ ( Eq . 2 ) . The final edited weights are W̃ = Wℓ − α∇̃Wℓ , where αℓ is a learned per-layer ( scalar ) step size .

Summary: This paper studies the problem of "model editing" - altering the model predictions on local examples without affecting the global behavior.  Authors tackle the challenge of editing very large models with billions of parameters, such as T5 or GPT-J, where editing is particularly important, as it is impractical to re-train such models from scratch to correct every mistake.  To address this challenge, authors propose MEND - a novel model editing method that differs from prior art in three ways: - it does not require computing higher-order gradients (and hence, requires less compute/memory) - it can be trained for a given model without the need to change the model parameters (and hence, can be applied to pre-existing models) - it has better asymptotic time & memory complexity in terms of model size (and hence can be used for very large models)