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: While large pre-trained nlp models have achieved great performance on a variety of downstream tasks, the largest of these models still make errors. This paper deals with the problem of enabling both developers and end users of such models to correct inaccurate outputs while leaving the model otherwise intact. If presented with only a single problematic input and new desired output, fine-tuning approaches tend to overfit. To enable easy post-hoc editing at scale, this paper proposes Model Editor Networks with Gradient Decomposition (MEND), a collection of small auxiliary editing networks that use a single desired input-output pair to make fast, local edits to a pre-trained model.   The 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.   MEND leverages the fact that gradients with respect to the fully-connected layers in neural networks are rank-1, enabling a parameter-efficient architecture that represents this gradient transform.  It provides experiments with T5, GPT, BERT, and BART models showing that MEND is the only approach to model editing that produces effective edits for models with tens of millions to over 10 billion parameters.