summaries: This paper tackles the challenging model of editing a potentially very large pre-trained model in a way that is specific (local) yet comprehensive (general). Their method, called MEND, take as input the decomposed, information-rich gradient from fine-tuning to learn the parameters for their meta-networks that subsequently perform the updates. The authors show that it works better than existing model editing techniques on large models like GPT-Neo, GPT-J, T5-XL, T5-XXL, and some smaller models like BERT-base, and distilGPT-2.

summaries: 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.

summaries: A fast model editing method is proposed in this paper to edit pre-trained models at scale. The core component is a Model Editor Networks with Gradient Decomposition (MEND). The paper claims that the MEND method has reliability, locality, generality. MEND is empirically validated on some curated datasets.

summaries: 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)

abstractText: While large pre-trained models have enabled impressive results on a variety of downstream tasks, the largest existing models still make errors, and even accurate predictions may become outdated over time. Because detecting all such failures at training time is impossible, enabling both developers and end users of such models to correct inaccurate outputs while leaving the model otherwise intact is desirable. However, the distributed, black-box nature of the representations learned by large neural networks makes producing such targeted edits difficult. If presented with only a single problematic input and new desired output, fine-tuning approaches tend to overfit; other editing algorithms are either computationally infeasible or simply ineffective when applied to very large models. To enable easy post-hoc editing at scale, we propose 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’s behavior. MEND learns to transform the gradient obtained by standard fine-tuning, using a low-rank decomposition of the gradient to make the parameterization of this transformation tractable. MEND can be trained on a single GPU in less than a day even for 10 billion+ parameter models; once trained MEND enables rapid application of new edits to the pre-trained model. Our experiments with T5, GPT, BERT, and BART models show that MEND is the only approach to model editing that effectively edits the behavior of models with more than 10 billion parameters.

authors:
Yonatan Belinkov
Nadir Durrani
Fahim Dalvi
Hassan Sajjad
James Glass

title: What do neural machine translation models learn about morphology? In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 861–872

authors:
Sid Black
Gao Leo
Phil Wang
Connor Leahy
Stella Biderman

title: GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow

authors:
Ilya Sutskever
Dario Amodei

title: Language models are few-shot learners

venue: Neural Information Processing Systems,

authors:
Nicola De Cao
W. Aziz
Ivan Titov

title: Editing factual knowledge in language models

authors:
Alexis Conneau
German Kruszewski
Guillaume Lample
Loïc Barrault
Marco Baroni

title: What you can cram into a single $&!#* vector: Probing sentence embeddings for linguistic properties

venue: In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers),

authors:
Damai Dai
Li Dong
Yaru Hao
Zhifang Sui
Furu Wei

title: Knowledge neurons in pretrained transformers

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

venue: In ICML,

authors:
Sebastian Flennerhag
Andrei A. Rusu
Razvan Pascanu
Francesco Visin
Hujun Yin
Raia Hadsell

title: Meta-learning with warped gradient descent

venue: In International Conference on Learning Representations,

authors:
Mor Geva
R. Schuster
Jonathan Berant
Omer Levy

title: Transformer feed-forward layers are key-value memories

venue: In EMNLP,

authors:
Xavier Glorot
Yoshua Bengio

title: Understanding the difficulty of training deep feedforward neural networks

venue: Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics,

authors:
Edward Grefenstette
Brandon Amos
Denis Yarats
Phu Mon Htut
Artem Molchanov
Franziska Meier
Douwe Kiela
Kyunghyun Cho
Soumith Chintala

title: Generalized inner loop metalearning

authors:
David Ha
Andrew M. Dai
Quoc V. Le

title: URL https://openreview.net/forum?id= rkpACe1lx

venue: Hypernetworks. In 5th International Conference on Learning Representations,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
John Hewitt
Christopher D. Manning

title: A structural probe for finding syntax in word representations

venue: In Proceedings of the

authors:
Edward Hu
Yelong Shen
Phil Wallis
Zeyuan Allen-Zhu
Yuanzhi Li
Weizhu Chen

title: Lora: Low-rank adaptation of large language models, 2021

authors:
Gao Huang
Zhuang Liu
Laurens van der Maaten
Kilian Q Weinberger

title: Densely connected convolutional networks

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Zhengbao Jiang
Frank F. Xu
J. Araki
Graham Neubig

title: How can we know what language models know

venue: Transactions of the Association for Computational Linguistics,

authors:
Divyansh Kaushik
Eduard Hovy
Zachary Lipton

title: Learning the difference that makes a difference with counterfactually-augmented data

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization, 2014. URL http://arxiv.org/abs/1412.6980

authors:
James Kirkpatrick
Razvan Pascanu
Neil Rabinowitz
Joel Veness
Guillaume Desjardins
Andrei A. Rusu
Kieran Milan
John Quan
Tiago Ramalho
Agnieszka Grabska-Barwinska
Demis Hassabis
Claudia Clopath
Dharshan Kumaran
Raia Hadsell

title: Overcoming catastrophic forgetting in neural networks

venue: Proceedings of the National Academy of Sciences,

authors:
Tom Kwiatkowski
Jennimaria Palomaki
Olivia Redfield
Michael Collins
Ankur Parikh
Chris Alberti
Danielle Epstein
Illia Polosukhin
Matthew Kelcey
Jacob Devlin
Kenton Lee
Kristina N. Toutanova
Llion Jones
Ming-Wei Chang
Andrew Dai
Jakob Uszkoreit
Quoc Le
Slav Petrov

title: Natural questions: a benchmark for question answering research

venue: Transactions of the Association of Computational Linguistics,

authors:
Yoonho Lee
Seungjin Choi

title: Gradient-based meta-learning with learned layerwise metric and subspace

venue: In International Conference on Machine Learning,

authors:
Omer Levy
Minjoon Seo
Eunsol Choi
Luke Zettlemoyer

title: Zero-shot relation extraction via reading comprehension

venue: In Proceedings of the 21st Conference on Computational Natural Language Learning (CoNLL

authors:
Zhenguo Li
Fengwei Zhou
Fei Chen
Hang Li

title: Meta-sgd: Learning to learn quickly for fewshot learning, 2017

title: distilgpt2-finetuned-wikitext2. https://huggingface.co/MYX4567/ distilgpt2-finetuned-wikitext2, July 2021

authors:
Michael McCloskey
Neal J. Cohen

title: Catastrophic interference in connectionist networks: The sequential learning problem

venue: Psychology of Learning and Motivation,

authors:
German I. Parisi
Ronald Kemker
Jose L. Part
Christopher Kanan
Stefan Wermter

title: Continual lifelong learning with neural networks: A review

venue: Neural Networks,

authors:
Ethan Perez
Florian Strub
Harm de Vries
Vincent Dumoulin
Aaron C. Courville

title: Film: Visual reasoning with a general conditioning layer

venue: In AAAI,

authors:
Fabio Petroni
Tim Rocktäschel
Sebastian Riedel
Patrick Lewis
Anton Bakhtin
Yuxiang Wu
Alexander Miller

title: Language models as knowledge bases

venue: In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),

authors:
Alec Radford
Jeff Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners, 2019

venue: URL https://d4mucfpksywv. cloudfront.net/better-language-models/language_models_are_ unsupervised_multitask_learners.pdf

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J. Liu

title: Exploring the limits of transfer learning with a unified text-totext transformer

venue: Journal of Machine Learning Research,

title: Connectionist models of recognition memory: constraints imposed by learning and forgetting functions

venue: Psychological review,

authors:
Sachin Ravi
H. Larochelle

title: Optimization as a model for few-shot learning

venue: In ICLR,

authors:
Adam Roberts
Colin Raffel
Noam Shazeer

title: How much knowledge can you pack into the parameters of a language

authors:
Anton Sinitsin
Vsevolod Plokhotnyuk
Dmitry Pyrkin
Sergei Popov
Artem Babenko

title: Editable neural networks

authors:
James Thorne
Andreas Vlachos
Christos Christodoulopoulos
Arpit Mittal

title: FEVER: a largescale dataset for fact extraction and VERification

venue: In NAACL-HLT,

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N. Gomez
Lukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Ben Wang
Aran Komatsuzaki

title: GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. https://github.com/kingoflolz/mesh-transformer-jax

authors:
Ruize Wang
Duyu Tang
Nan Duan
Zhongyu Wei
Xuanjing Huang
Jianshu ji
Guihong Cao
Daxin Jiang
Ming Zhou

title: K-adapter: Infusing knowledge into pre-trained models with adapters, 2020

venue: URL http://arxiv.org/abs/2002.01808

authors:
Shibo Wang
Pankaj Kanwar

title: Bfloat16: The secret to high performance on cloud tpus, 2019

venue: URL https://cloud.google.com/blog/products/ai-machine-learning/ bfloat16-the-secret-to-high-performance-on-cloud-tpus. [Online; accessed 28-September-2021]

authors:
Thomas Wolf
Lysandre Debut
Victor Sanh
Julien Chaumond
Clement Delangue
Anthony Moi
Pierric Cistac
Tim Rault
Rémi Louf
Morgan Funtowicz
Jamie Brew

title: Huggingface’s transformers: State-of-the-art natural language processing

venue: URL http://arxiv.org/abs/1910.03771

authors:
Hongyi Zhang
Yann N. Dauphin
Tengyu Ma

title: Residual learning without normalization via better initialization

authors:
Chen Zhu
Ankit Singh Rawat
Manzil Zaheer
Srinadh Bhojanapalli
Daliang Li
Felix Yu
Sanjiv Kumar

title: Modifying memories in transformer models, 2020

venue: URL https://arxiv

venue: (Wolf et al.,

title: Additional examples of using MEND to edit a 770M parameter T5-large model fine-tuned on Natural Questions (NQ

heading: 1 INTRODUCTION

text:
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.

heading: 2 THE MODEL EDITING PROBLEM

text:
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)

heading: 3 MODEL EDITOR NETWORKS WITH GRADIENT DECOMPOSITION

text: 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.

heading: 3.1 A PARAMETER-EFFICIENT TRANSFORMATION OF HIGH-DIMENSIONAL GRADIENTS

text:
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.

heading: 3.2 TRAINING MEND

text:
MEND uses an editing training set Dtredit to learn parameters ϕℓ for each of the MEND networks gℓ. Before training, we select the weights of the modelW = {W1, ...,WM} that we would like to make editable (e.g., the weight matrices in the last M layers). At each step of training, we sample an edit example (xe, ye), locality example xloc, and equivalence examples (x′e, y ′ e) from the edit train set Dtredit. Recall that xloc is sampled independently from the edit example, so that it is very likely that it is unrelated to the edit example. We use (xe, ye) to compute the raw gradient∇WℓpθW (ye|xe) for each weight matrix Wℓ ∈ W , using θW to denote the model parameters with un-edited weights. We then compute the parameter update for each layer W̃ = Wℓ − αℓ∇̃Wℓ (∇̃Wℓ from Eq. 2).
We compute the training losses for MEND using the edited model parameters W̃ , which we backpropagate into the editing networks. Note that we do not compute any higher-order gradients, because we do not optimize the pre-edit model parameters. The training losses are Le, which measures edit success and Lloc, which measures edit locality (the KL divergence between the pre-edit and post-edit model conditioned on the locality input xloc), defined as follows (also Alg. 1 lines 5-7):
2For a batch/sequence, we transform the gradient for each batch/sequence element independently and sum the result to acquire the final transformed gradient for the entire batch/sequence.
MEND losses: Le = − log pθW̃ (y ′ e|x′e), Lloc = KL(pθW (·|xloc)∥pθW̃ (·|xloc)). (4a,b)
Intuitively, Le is small if the model has successfully updated its output for the edit example’s equivalence neighborhood, while Lloc is small if the edit did not affect the model’s behavior on unrelated inputs. The total training loss for a MEND network is computed as LMEND = ceLe(θW̃) + Lloc(θW , θW̃). We optimize LMEND with respect to the MEND parameters at each time step using the Adam optimizer (Kingma and Ba, 2014), using ce = 0.1 for all experiments.
While MEND’s parameterization can tractably represent a mapping from gradients to model edits, training the editor presents its own challenges. Appendix A describes MEND’s identity initialization and input normalization, which our ablations in Section 5.4 show are important to effective edits.
4 RELATED WORK Preserves model? Only (xe, ye)? Batched edits? Scales to 10B? Few steps?Editor
FT ✓ ✓ ✓ ✓ ✗ FT+KL ✓ ✗ ✓ ✓ ✗ ENN ✗ ✓ ✓ ✗ ✓ KE ✓ ✓ ? ✓ ✓
MEND ✓ ✓ ✓ ✓ ✓
Table 1: Conceptual comparisons of model editors; MEND provides a unique combination of useful attributes. Preserves model means the editor guarantees model predictions will not be altered before an edit is applied. Only (xe, ye) means the editor applies an edit at test time using only the edit pair (not needing access to the training set at test time as well). Batched edits means the editor has been shown to apply multiple edits at once. Scales to 10B means our implementation of the editor could run on a model with over 10B parameters using our single-GPU environment (see Appendix C.3). Few steps means edits are applied with one or a small number of steps. FT refers to fine-tuning; FT+KL adds a KL-divergence penalty between the original and fine-tuned model.
Various strategies for model editing exist, including modifications of standard fine-tuning intended to enforce locality by reducing distance traveled in parameter space (Zhu et al., 2020) or even find the min-L2 norm parameter update that reliably edits the model’s output (Sotoudeh and Thakur, 2021). However, Cao et al. (2021) observe that parameter-space constraints do not always translate to useful function-space constraints for neural networks. Our fine-tuning baselines thus use a KL-divergence constraint in function space, but, even with this modification, we find that fine-tuning generally doesn’t consistently provide edit generality. Other approaches to editing such as Editable Neural Networks (ENN; Sinitsin et al. (2020)) or KnowledgeEditor (KE; Cao et al. (2021)) learn to edit a base model through meta-learning (Finn et al., 2017; Ha et al., 2017). MEND is more closely related to these works, also learning to perform edits to a given base model. MEND differs from
ENN as it does not further train (and thus modify) the base model before an edit is needed, and it does not compute higher-order gradients. Because ENN modifies the pre-edit model, the training process retains a copy of the original model in order to enforce the constraint that the editable model agrees with the original pre-trained model’s predictions. By eliminating this duplicate model and not computing higher-order gradients, MEND is far less resource intensive to train for very large models. Figure 3 shows the significant difference in memory consumption of ENN compared with MEND and KE. MEND is most similar to KE, which also presents a first-order algorithm that does not modify the pre-edit model. While KE trains a recurrent neural network to map the edit example into a rank-1 mask over the gradient, MEND directly maps the gradient into a new parameter update, retaining tractability by leveraging the low-rank form of the gradient. Table 1 contains an overview of algorithmic tradeoffs. See Appendix B for extended discussion of related work.
Various methods for meta-learning also use gradient transforms to achieve better model updates for few-shot learning (Ravi and Larochelle, 2017; Li et al., 2017; Lee and Choi, 2018; Park and Oliva, 2019; Flennerhag et al., 2020). However, these approaches do not leverage the factorized gradient, limiting them to simpler transformations (typically linear) of the gradient and/or transformations that also often impact the function computed by the forward pass of the model. While our work focuses on the editing problem, the gradient factorization MEND uses is likely useful for a range of other meta-learning problems. Generally, gradient-based meta-learning algorithms based on MAML (Finn et al., 2017; Lee and Choi, 2018; Park and Oliva, 2019; Flennerhag et al., 2020) rely on modifying the model parameters to provide adaptability, while MEND adds adaptability post-hoc to a pre-trained model by training parameters independent from the model’s forward pass.
In the NLP literature, many papers have investigated the locus of various types of knowledge in language models, using learned probe models or iterative search procedures to test for linguistic
structures (Belinkov et al., 2017; Conneau et al., 2018; Hewitt and Manning, 2019) or facts about the world (Petroni et al., 2019; Jiang et al., 2020; Dai et al., 2021). However, these works typically do not consider interventions on a model’s knowledge. Exceptions are Dai et al. (2021) and Wang et al. (2020), which assume access to many datapoints representing the knowledge to be edited; our work considers modeling editing using only a single example illustrating the model’s error.

heading: 5 EXPERIMENTS

text:
A key motivation for MEND is scalability to large models, which requires an algorithm to be efficient in terms of computation time and particularly memory consumption. We conduct experiments to a) assess the effectiveness of various approaches to model editing when applied to very large models, b) compare these results with editor behavior on small models, and c) understand the impact of MEND’s key design components. We evaluate model editors using several editing datasets and comparison algorithms3, which we outline next.
Editing Datasets. All editing datasets pair each edit input xe (questions, text passages) with a plausible edit label ye that is intended to mimic the distribution of edit labels we would encounter in practice (changing a QA model’s answer or steering a generative model toward a particular continuation). For example, in a QA setting, plausible edit labels include the ground truth label as well as entities of the same type as the true answer. See Appendix C.4 for sample data. Specifically, for seq2seq models, we use the zsRE question-answering dataset (Levy et al., 2017) using question rephrasings generated by backtranslation as the equivalence neighborhood and train/val splits generated by Cao et al. (2021). Each xe is a question about an entity, and plausible alternative edit labels ye are sampled from the top-ranked predictions of a BART-base model trained on zsRE questionanswering. When editing models pre-trained on the zsRE question-answering problem, we sample xloc as independent questions from the edit train set. For other experiments (Section 5.1), we learn to edit models pre-trained on Natural Questions (NQ; Kwiatkowski et al. (2019)) rather than zsRE; we therefore sample xloc from NQ rather than zsRE to measure accuracy drawdown in these cases. For classification models (e.g., BERT), we use the FEVER fact-checking dataset (Thorne et al., 2018) with fact rephrasings and train/val splits also generated by Cao et al. (2021). Each xe is a fact, and each ye is a random binary label sampled from a Bernoulli distribution with p = 0.5. Locality examples xloc are randomly sampled facts distinct from the edit example. For GPT-style models, we create a Wikitext generation editing dataset of similar size to the zsRE and FEVER editing
3For each dataset, all algorithms edit the same parameters. For BART/T5, we edit the MLP layers of the last 2 encoder & decoder blocks; for GPT/BERT models, we edit the MLPs in the last 3 blocks.
datasets, containing approximately 68k xe, ye pairs. Each xe is a passage sampled from Wikitext103 and ye is a 10-token sample from a pre-trained distilGPT-2 model.4 xloc is chosen depending on the pre-trained model: for models pre-trained on Wikitext, xloc is sampled from Wikitext-103 (independently from xe). For GPT-Neo/J, we sample xloc from OpenWebText (OWT; (Gokaslan and Cohen, 2019)) to better match the model’s original training data. The equivalence neighborhood in this setting is N(xe, ye) = {(xke , ye)}, where xke is formed by removing a prefix of up to |xe| 2 tokens from the beginning of xe, where |xe| is the length of xe in tokens. Comparison of model editors. We compare MEND with several other model editors, including two fine-tuning-based algorithms (which do not train any model editor at all) and two learned model editors. The fine-tune (FT) algorithm fine-tunes on the edit example (xe, ye) until the label is assigned the highest likelihood (using greedy decoding for sequence models). The ‘oracle’ fine-tune + KL (FT+KL) algorithm has access to the training set at test time and adds Lloc (Eq. 4b) to the test-time fine-tuning objective (which is typically only computable during model editor training). Similarly to Cao et al. (2021), we limit each of these algorithms to 100 fine-tuning steps. Additionally, we compare with two learned model editors: a re-implementation of Editable Neural Networks (ENN; Sinitsin et al., 2020) when possible (due to high memory usage) and KnowledgeEditor (KE; Cao et al., 2021). We use identical hyperparameters for MEND across all models and datasets. For BART and T5 models, we edit the MLP weight matrices in the last 2 transformer blocks of the encoder and decoder; for other models, we edit the MLP weights in the last 3 transformer blocks. Appendix G compares MEND with a simple caching-based approach to editing, which performs well for short inputs (zsRE and FEVER) but fails completely for long, diverse inputs (Wikitext).
Metrics. Our experiments measure the reliability and generality of a model editor using edit success (ES) (Eq. 1). To assess locality, we use drawdown (DD), which is defined as the performance degradation of the edited model on the rest of the dataset, measured as either the edited model’s perplexity increase or accuracy decrease compared to the base model, depending on the problem.

heading: 5.1 EDITING VERY LARGE TRANSFORMER MODELS

text:
We first consider the problem of editing some of the largest publicly-available Transformer models. We use GPT-Neo (2.7B parameters; Black et al., 2021) and GPT-J (6B parameters; Wang and Komatsuzaki, 2021), several times larger than GPT-2 (Radford et al., 2019), and the largest two T5 models, T5-XL (2.8B parameters) and T5-XXL (11B parameters) fine-tuned on NQ (Roberts et al., 2020). Table 3 shows the results; MEND provides the most successful edits across tasks. Fine-tuning achieves lower edit success on the Wikitext task and exhibits a much larger perplexity increase than MEND. On the question-answering edit task, fine-tuning shows similarly reduced edit success, struggling to generalize to some rephrasings of the edit input. The KL-constrained baseline reduces the perplexity drawdown for GPT-Neo and GPT-J, but at the cost of edit success. KE is ineffective at this scale, generally failing to provide successful edits. For these experiments, we use
4The base model’s greedy 10-token prediction agrees with these edit targets for <1% of examples.
OWT and NQ to measure drawdown for generation and question-answering, respectively, as they are more representative of the data used to train the base models.
5.2 SMALLER SCALE EDITING
during training. ENN’s memory usage5 is prohibitively high for very large models, while MEND and KE can be trained on a single GPU. Figure 4 shows similar chart for GPT models.
We conduct an additional experiment editing the BERT-base and BART-base models fine-tuned by Cao et al. (2021) on the FEVER fact-checking and zsRE question-answering tasks, respectively, and our Wikitext editing task, editing a smaller distilGPT-2 model (Wolf et al., 2019) fine-tuned on Wikitext2 (Ma, 2021). These models are 1–2 orders of magnitude smaller than those in Section 5.1. Results are presented in Table 4. At small scale where computational requirements are not a concern, ENN is competitive with MEND, providing the best performance on the Wikitext problem. Fine-tuning overfits even more severely than with larger models, showing lower edit success (overfitting to the edit example) and higher drawdown (degrading the model more seriously). One difficulty of using ENN is that the pre-trained model itself must be fine-tuned to ‘provide’ editability, potentially changing the model’s predictions even before an edit has been applied. Un-
like the large-scale experiments, drawdown is computed using samples from the same datasets as edit inputs, again in order to best match the data distribution the base models were fine-tuned on.
5.3 BATCHED EDITING Edit Success ↑ Acc. Drawdown ↓
Table 5 compares MEND with ENN (the strongest comparison method) in a more realistic setting when multiple simultaneous zsRE QA model edits are needed; MEND consistently provides significantly more effective edits in the multi-edit setting. Both algorithms are trained and evaluated on applying k simultaneous edits, with k ∈ {1, 5, 25, 75, 125}. MEND applies simultaneous edits by simply summing the parameter edit computed separately for each edit example. MEND applies 25 edits in a single model update with 96% edit success and less than 1% accuracy degrada-
tion (35% edit success for ENN), and successfully applies 67% of edits when applying 125 edits at once (11% success for ENN, although ENN’s accuracy drawdown is slightly lower).
5We report the memory usage of our re-implementation of ENN (Sinitsin et al., 2020). Techniques like gradient checkpointing can reduce memory consumption, but an optimized ENN implementation is not available.

heading: 5.4 ABLATIONS & MEND VARIANTS

text: Table 6 shows ablations of MEND’s parameter sharing, identity initialization, and input normalization as well as three variants of MEND that reduce total parameters: only computing pseudoactivations uℓ, only pseudodeltas δℓ+1, or only whichever of uℓ or δℓ+1 is lower-dimensional (layerdependent for non-square weights). ‘No ID init.’ replaces zero initialization with Xavier/Glorot initialization (Glorot and Bengio, 2010). Removing either input normalization or identity initialization significantly reduces edit effectiveness (and increases training time ∼10x). Sharing parameters across model editor networks incurs relatively little performance cost, and editing only the smaller of the pseudoactivations and pseudodeltas, the most most lightweight version of MEND, still produces effective edits, suggesting that MEND could scale to even much larger models for which m + n approaches 105 (Brown et al., 2020) but min(m,n) remains close to 104. Appendix E shows an additional ablation editing attention matrices, rather than MLP weights, finding that editing MLP weights is consistently more effective for large models.

heading: 6 DISCUSSION

text:
Conclusion. We have presented an efficient approach to editing very large (10 billion+ parameter) neural networks, which we call Model Editor Networks with Gradient Decomposition or MEND. We showed that MEND is the only method that successfully edits the largest publicly-available Transformer models from the GPT and T5 model families. To do so, MEND treats the model editing problem itself as a learning problem, using a relatively small edit dataset to learn model editor networks that can correct model errors using only a single input-output pair. 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.
Limitations & Future Work. A limitation of existing model editors (including MEND) is the approach to enforcing locality of edits. The failure mode of over-generalization (bottom of Table 2) shows that locality examples (i.e., negative examples) are not challenging enough to prevent the model from sometimes changing its output for distinct but related inputs. Alternative locality losses or harder negative mining may help address this problem. Further, existing language-based editing datasets use backtranslation to evaluate edit generality (and our Wikitext dataset uses a truncation heuristic). Such equivalence neighborhoods do not assess a model’s ability to use the knowledge in an edit example to correctly answer questions about other topics whose answer is implied by the content of the edit example (e.g., for Who is the UK PM? Boris Johnson, does the edited model correctly answer Is Boris Johnson a private citizen?). Counterfactual data augmentation (Kaushik et al., 2020) may be useful for constructing richer evaluation cases for edit generality. Future work might also apply MEND to other types of edits, such as reducing the frequency of toxic generations after observing toxic outputs, relabeling entire classes of images from one example, or adjusting a robot’s control policy to avoid particular actions, as MEND is not limited to editing transformer models. Finally, MEND’s gradient decomposition is not in principle limited to the model editing problem, and it might enable efficient new gradient-based meta-learning algorithms.

heading: ETHICS STATEMENT

text: This work uses large language models pre-trained on text scraped from the internet. These massive training corpuses (and therefore the models trained on them) may contain (or produce) content that is counter to the values of the ICLR community. Algorithms for model editing may provide one tool (among others) to mitigate this problem by enabling maintainers of large models to change certain undesirable model behaviors as they are discovered. On the other hand, a model editor could also be used to exacerbate the very model behaviors that we hope to eliminate, depending on who is wielding it. This dual use is a risk for many machine learning technologies. Specifically, effective editing algorithms (including MEND and others) may enable maintainers of deployed neural networks to include backdoors or other planned vulnerabilities/hidden behaviors into their models.

heading: REPRODUCIBILITY

text: To foster reproducibility, we have provided a detailed description of the proposed algorithm in Section 3, as well as additional details regarding experimental setup, hyperparameters, and implementations of comparison algorithms in Section C. Our experiments use fixed random seeds for data sampling and model editor initialization, enabling reproducible results. Section C.4 describes how to obtain the pre-existing datasets and models we used in our experiments (from Cao et al. (2021)), and we including our Wikitext dataset in the Supplementary Material. Finally, we are including the source code used for our experiments in the Supplementary Material for the purpose of peer review. In light of some dependencies of our codebase that are difficult to anonymize, the submitted version of this code is not runnable. The complete, runnable code (including dependencies) and data for our experiments will be released publicly after double blind reviewing is complete.

heading: A EFFECTIVE INITIALIZATION AND NORMALIZATION FOR MEND NETWORKS

text: Although random weight initialization is effective in many settings, it sacrifices the prior that the raw fine-tuning gradient is a useful starting point for editing. Our ablations show that it also leads to less effective edits. For this reason, we initialize MEND to the identity function using a residual connection (He et al., 2016) and a partially random, partially zero-initialization strategy related to Fixup (Zhang et al., 2019). Referring back to Eqs. 3a,b, U1 and U2 are initialized with zeros, and V1 and V2 use standard Xavier uniform initialization (Glorot and Bengio, 2010) (also see Figure 2). Beyond the initialization, input scaling also presents a challenge: inputs to a MEND network (uℓ and δℓ+1) can differ in magnitude by several orders of magnitude. This poor conditioning causes training to be slow and edit performance to suffer (see Section 5.4). Input normalization addresses this issue; we normalize each dimension of both uℓ and δℓ+1. The input to gℓ is the concatenation of ūℓ = norm(uℓ) and δ̄ℓ+1 = norm(δℓ+1), where ūℓ and δ̄ℓ+1 are normalized to have zero mean and unit variance, with means and variances computed over the edit train set and the sequence index.

heading: B EXTENDED DISCUSSION OF RELATED WORK

text:
Model editing shares with continual learning (McCloskey and Cohen, 1989; Parisi et al., 2019) the goal of assimilating or updating a model’s behavior without forgetting old information or behaviors, commonly known as the problem of catastrophic forgetting (McCloskey and Cohen, 1989; Ratcliff, 1990; Kirkpatrick et al., 2017). However, in continual learning settings, a model is typically expected to learn wholly new behaviors or datasets (Kirkpatrick et al., 2017; Parisi et al., 2019) without forgetting, while in this work we consider more localized model edits. Further, continual learning generally considers long sequences of model updates with minimal memory overhead, while our work generally considers an edit or batch of edits applied all at once.
Additionally, min-norm parameter fine-tuning has also been considered in past work in the context of editing (Zhu et al., 2020) and traditional model fine-tuning (Guo et al., 2021), where the parameters of the edited or fine-tuned model θ′ are penalized (or constrained) from drifting too far from the original model parameters θ using various norms, including L0, L2, and L-∞. While min-norm constraints may be an effective regularization for traditional fine-tuning settings where fine-tuning data is abundant, the experiments conducted in Cao et al. (2021) show that parameter-space norm constraints are insufficient constraints to prevent significant model degradation when fine-tuning on a single edit example.

heading: B.1 EDITABLE NEURAL NETWORKS (ENN)

text:
Editable neural networks (Sinitsin et al., 2020) search for a set of model parameters that both provide good performance for a ‘base task’ (e.g., image classification or machine translation) and enable rapid editing by gradient descent to update the model’s predictions for a set of ‘edit examples’ without changing the model’s behavior for unrelated inputs. ENN optimizes the following objective, based on the MAML algorithm (Finn et al., 2017):
LENN(θ,Dbase,Dedit,Dloc) = Lbase(Dbase, θ) + cedit · Ledit(Dedit, θ′) + cloc · Lloc(Dloc, θ, θ′). (5)
The first term of Equation 5 is the base task loss; for a generative language model, we have Lbase(Dbase, θ) = − log pθ(Dbase) where Dbase is a batch of training sequences. Lbase is the edit reliability loss, encouraging the model to significantly change its output for the edit examples in Dedit. Finally, Lloc is the edit locality loss, which penalizes the edited model θ′ for deviating from the predictions of the pre-edit model θ on Dloc, data unrelated to Dedit and sampled from the same distribution as Dbase. See Sinitsin et al. (2020) for a more detailed explanation of ENN training and alternative objectives for Ledit and Lloc.
Comparing ENN and MEND. The key conceptual distinction between ENN and MEND is that ENN encodes editability into the parameters of the model itself (intrinsic editability), while MEND provides editability through a set of learned parameters that are independent from the model parameters (extrinsic editability). An advantage of ENN is that no new parameters are added in order to provide editability. However, this approach comes with several drawbacks. First, the MAML-based objective ENN optimizes is expensive, particularly in terms of memory consumption (see Figure 4). By further training the model parameters themselves, ENN cannot guarantee that the editable model it produces will make the same predictions as the original model. In order to approximately enforce this constraint during training, ENN must use an extra copy of the original base model to ensure that the editable model’s predictive distribution does not differ too much from it. This incurs significant additional memory costs, particularly when training ENN for very large models, for which the parameters of the model alone occupy a significant amount of VRAM. Another cause for the significant VRAM consumption of ENN is the need to compute activations and gradients for the model parameters; even if we edit only the last layer, ENN trains the rest of the model so that the last layer gradient is productive, requiring activations and gradients to be computed for the entire model. On the other hand, extrinsic editors like MEND and KE do not require updating the base model itself, thereby computing gradients for far fewer parameters. Future work might investigate approaches to reducing the memory consumption of ENN, although the requirement to retain a copy of the original model in order to enforce locality creates a relatively high lower bound on the amount of memory that ENN might use.
Regardless of memory consumption, extrinsic editors have the potential advantage of being able to edit more than one model; in theory, we might amortize the cost of training MEND over several base models at once. On the other hand, intrinsic editability must by definition be re-learned separately for each base model.

heading: B.2 KNOWLEDGEEDITOR (KE)

text:
Cao et al. (2021) propose KNOWLEDGEEDITOR, a hypernetwork-based approach for editing the knowledge in language models. KE is an RNN that conditions explicitly on the input, incorrect output, and new desired label and outputs a mask mi, offset bi, and a scalar scaling factor α to the gradient∇Wi for several of the weight matrices in a transformer model, where mi, bi,∇Wi ∈ Rd×d for a d × d weight matrix. The update to the model is θ′ = θ − α(mi ◦ ∇Wi) + bi. Because the weight matrices in state-of-the-art transformer models are very high-dimensional, the mask and offset output by KE are rank-1 to retain tractability.
Comparing KE and MEND. KE more closely resembles MEND in that it is also an extrinsic model editor. However, while MEND directly maps model gradients into model edits, the KE model editor uses the raw edit example as an input, outputting a single rank-1 mask and rank-1 offset over the fine-tuning gradient. We hypothesize that the KE model faces several challenges that MEND avoids. First, mapping the edit example itself into a model updates requires a translation from the high-level modality of data examples into the very low-level modality of model parameter updates. Solving this translation requires making additional design decisions (e.g., how to feed the edit input and label into the editor, what architecture to use for the editor), the optimal design for which may vary across problems. Further, by not conditioning directly on the gradient, KE forgoes a rich source of information about which parameters of the model are most responsible for updating the model’s outputs. In addition, by operating on the token-wise activations and gradients (i.e., the gradients are not summed over the sequence/batch, but are kept as per-sequence element activation and gradient vectors), MEND outputs a rank-1 model edit for each token in the input and output sequence. The final output of MEND is the sum of these, which has rank of order 10 or even 100, depending on the problem. In contrast, the KE editor outputs only a rank-1 gradient mask and rank-1 gradient offset, regardless of the information content of the edit example. This rank-1 constraint, irrespective of the size of the input, which we hypothesize causes KE’s failure to perform well for the Wikitext editing
task, which has significantly higher information content labels (10 tokens) than the FEVER or zsRE tasks.

heading: C EXPERIMENTAL DETAILS

text: For GPT and BERT-style models, all experiments edit the MLP weights in the last 3 transformer blocks (6 weight matrices total). For BART and T5-style models, all experiments edit the MLP weights in the last 2 transformer blocks in both the encoder and the decoder (8 weight matrices total). We found that editing MLP layers generally provides better editing performance (across algorithms) than editing attention layers. In line with past work (Cao et al., 2021), all reported performance numbers are on the validation set. For all algorithms, we use early stopping to end training early if the validation loss L = ceditLe +Lloc) does not decrease for 20000 steps on a subset of 500 validation examples, with a maximum number of training steps of 500,000. We use a batch size of 10 (with gradient accumulation) and the seed 0 for all experiments.

heading: C.1 HYPERPARAMETERS

text:
Fine-tuning. The fine-tuning baselines use model-dependent learning rates, which we found important in achieving good fine-tuning performance; using too large of a learning rate causes decreased locality (increased model degradation), while a learning rate too small causes slow edits. We use edit learning rates of 5e-6 for GPT-Neo and GPT-J and 1e-4 for T5 models, and 1e-6 for the smaller models, aiming to complete edits in less than 100 fine-tuning steps (as in Cao et al. (2021)). For the fine-tuning + KL-constraint baseline, we fine-tune on the loss ceditLe +Lloc, using a smaller cedit than for the learned algorithms (1e-2 for all models except GPT-J, which required 1e-3). Larger values of cedit provide little benefit from the locality loss. To compute Lloc, we use a batch size of one new example xloc from the full edit training set Dtredit at each time step.
ENN. We use an initial inner loop learning rate of 1e-2, but allow this value to be learned in the outer loop, which we find improves performance over the fixed inner loop learning rate version in Sinitsin et al. (2020). For all experiments, ENN fine-tunes all model parameters during training (even when we only edit the last few layers). We also use only a single inner loop update step for computational reasons, which differs from the multi-step version used for the smaller models used by Sinitsin et al. (2020). Our edit loss is also a slight simplification of the edit loss used by Sinitsin et al. (2020), which is
le(θ) = − log pθ(ye|xe, θ) + max yi log pθ(yi|xe, θ) (6)
The first term of this loss is the edit loss we use in our work; the second term is primarily intended to provide the property that le(θ) ≤ 0 when an edit is successful so that the iterative editing process can be stopped. However, in this work, because we use only a single gradient step of editing for ENN, this property is less important, and the second term simply amounts to an additional emphasis on pushing down specifically the largest incorrect logit (which the first term already does implicitly).
KE We use the implementation of KE provided by Cao et al. (2021), which can be found at https://github.com/nicola-decao/KnowledgeEditor, with minor changes to the computation of the KL constraint for consistency with other algorithms (see below). We use a learning rate of 1e-5.

heading: C.2 COMPUTING THE LOCALITY CONSTRAINT

text:
Computing the true KL-divergence between the pre- and post-edit model KL(pθ(·|xloc)∥pθ′(·|xloc)) quickly becomes computationally prohibitive for model outputs of more than a few tokens, requiring marginalization over possible answers. We therefore approximate this KL-divergence using samples from the dataset.6 For the seq2seq question-answering problem, we evaluate the KL divergence only at the tokens of the answer yloc, giving KLseq2seqapprox (θ, θ
′) = 1 |yloc| ∑|yloc| i=1 KL(pθ(·|xloc, y <i loc )∥pθ′(·|xloc, y <i loc )), where p(·|xloc, y <i loc ) is the distribution over next tokens yi given the locality input xloc and the label tokens for previous timesteps y<iloc . Similarly, for the Wikitext setting, we define KLautoapprox(θ, θ ′) = 1|xloc| ∑|xloc| i=1 KL(pθ(·|x <i loc)∥pθ′(·|x <i loc)). For FEVER fact-checking we compute the exact KL-divergence between Bernoulli distributions in closed form.

heading: C.3 ENVIRONMENT DETAILS

text: All runs are trained entirely on a single NVIDIA RTX Titan or A40 GPU. No gradient checkpointing or memory-reduction optimizations are used, although bfloat16 is used to fit the largest T5 model onto our GPU. In full precision, the parameters alone of the T5-11B model use all of the memory of our largest GPU. VRAM consumption for training MEND and KE on T5-11B (Figs. 3 and 4) is estimated by doubling the bfloat16 VRAM usage (Wang and Kanwar, 2019). While doubling half precision enabled estimating the memory consumption of ENN, we were unable to train ENN in half precision without numerical instability. All models are based on Huggingface Transformers implementations (Wolf et al., 2019) with some modifications in line with Cao et al. (2021). We use PyTorch (Paszke et al., 2019) for all experiments, specifically using the Higher library (Grefenstette et al., 2019) in order to implement the bi-level optimization in ENN as well as the inner loop of model editing for all algorithms.

heading: C.4 DATASET CONSTRUCTION & EXAMPLES

text: Datasets are constructed to provide pairs of edit input xe and plausible edit label ye. The edit label is not necessarily the ‘correct’ label; the goal is to provide realistic instances of the types of data we would expect to see during test. For example, our dataset might have a sample such as xe = Where was Ursula K. Le Guin born? and ye = Addis Ababa, Oromia, Ethiopia, even though Ursula K. Le Guin was born in Berkeley, California, USA. However, this fictitious example is still a useful assessment of our model’s ability to perform the general type of edit of ‘change a person’s birthplace’. For the zsRE question-answering dataset Cao et al. (2021) generate fictitious ye in this manner using the top predictions of a BART model fine-tuned on the task of question answering followed by manual human filtering. In practice, this produces alternate edit labels that are plausible and whose types match with the original label. For FEVER fact-checking, there are only two choices for labels, and we sample edit targets 1 and 0 with equal probability. For Wikitext generation, we use a distilGPT-2 model to generate plausible 10-token continuations for a given Wikitext prefix, with the similar motivation to zsRE of providing edit targets that share the structure of the types of edits that we will apply in practice, even if they are not always factual. When qualitatively assessing MEND to correct real errors of the base model using the factual labels, we find that MEND performs reliably, indicating that these label generators provide reasonable proxies for ‘real’ model edits.

heading: D RANK-1 GRADIENT FOR MLPS

text:
In the simplified case of an MLP and a batch size of 1, we describe the rank-1 gradient of the loss L with respect to the layer ℓ weight matrix Wℓ. We define the inputs to layer ℓ as uℓ and the preactivation inputs to layer ℓ+1 as zℓ+1 = Wℓuℓ. We define δℓ+1 as the gradient of L with respect to zℓ+1 (we assume that δℓ+1 is pre-computed, as a result of standard backpropagation). We will show that the gradient of the loss L with respect to Wℓ is equal to δℓ+1u⊤ℓ .
By the chain rule, the derivative of the loss with respect to weight W ijℓ is equal to
∂L
∂W ijℓ = ∑ k ∂L ∂zkℓ+1 ∂zkℓ+1 ∂W ijℓ = ∂L ∂ziℓ+1 ∂ziℓ+1 ∂W ijℓ (7)
the product of the derivative of L with respect to next-layer pre-activations ziℓ+1 and the derivative of next-layer pre-activations ziℓ+1 with respect to Wij . The second equality is due to the fact that ∂zkℓ+1 ∂W ijℓ = 0 for k ̸= i. Noting that ziℓ+1 = ∑ j u j ℓW ij ℓ , we can replace ∂ziℓ+1 ∂W ijℓ with simply ujℓ in Equation 7. Further, we defined δℓ+1 to be exactly ∂L∂ziℓ+1 . Making these two substitutions, we have
∂L
∂W ijℓ = δiℓ+1u j ℓ (8)
or, in vector notation,∇WℓL = δℓ+1u⊤ℓ , which is the original identity we set out to prove.

heading: E EDITING ATTENTION PARAMETERS

text:
Our experiments edit weights in the MLP layers of large transformers. Here, Table 9 shows the results of editing the attention layers, rather than MLP layers, observing that editing attention layers generally leads to reduced performance compared to editing MLP layers. For this comparison, we edit the same transformer blocks as for our main editing experiment in Table 3, but we edit the query/key/value/output matrices for each block instead of the two MLP matrices. The observation that editing MLP layers is more effective generally aligns with past work (Geva et al., 2021) suggesting that the MLP layers in Transformer architectures store human-interpretable, high-level concepts in the later layers of the model, motivating our choice of editing these layers in our original experiments. Further, we hypothesize that the improved effectiveness of editing MLP layers may
6We justify this choice by the fact that the model’s predictive distribution is similar to the locality sample distribution (as locality samples are drawn from the dataset the model was originally trained on). While this is not as principled as a true Monte Carlo estimate using samples from the model itself, it is reduces computational requirements of training and is easier to implement; the generally low drawdown for most models indicates that this approximation still provides a good locality constraint in practice.
simply be based on the fact that they make up a large majority of model parameters, as the MLP hidden state is often much higher-dimensional than the model’s hidden state.

heading: F ADDITIONAL QUALITATIVE EXAMPLES OF MEND

text: We provide additional qualitative examples of using MEND to edit a larger 770M parameter T5-large model (Roberts et al., 2020) in Table 10. These examples include an instance of undergeneralization, in which the edit example’s output is correctly edited, but other examples in the equivalence neighborhood of the edit example do not change (see 2f in Table 10)). In addition, we highlight the failure case of overgeneralization, in which the model’s post-edit output for superficially similar but semantically distinct inputs is also the edit target; for example 3e, 3f, and 4e in Table 10. Mitigating these failure cases for model editors (ensuring is an important priority for future work,

heading: G EDITING THROUGH MEMORIZATION OR CACHING

text:
Another simple approach to editing might be to cache the final layer hidden state ze (averaged over the sequence length) of the edit example xe and the tokens of the corresponding edit label ye. After an edit is performed, if the model receives a new input x whose final layer hidden state z is close to ze (i.e. ∥z − ze∥2 < ϵ), then the model outputs ye instead of its normal prediction. Here, we show that this approach is effective for editing problems with simpler inputs (zsRE questionanswering, FEVER fact-checking), where inputs are typically short, simple phrases with one subject, one relation, and one object, but fails completely on the Wikitext editing problem, where contexts are typically 10x as long, with diverse passages containing significant amounts of extraneous text and ‘distracting’ information. The results are presented in Table 11. We include the ‘optimal’ threshold ϵ∗ (the threshold that achieves similar drawdown to MEND), as well as the result of using 2ϵ∗ and 1 2ϵ
∗. We observe that the caching approach is fairly sensitive to the threshold hyperparameter, and a threshold that works well for one task may not work well for others.
For zsRE question answering, z is computed as the average hidden state of the question tokens; for FEVER fact-checking, z is the average hidden state of the fact statement tokens. For generative modeling, when predicting the token at time step t, we compute zt as the average hidden state for all previously seen tokens < t. In order to compute perplexity for the caching approach, we output onehot logits corresponding to ye. We experimented with scaling the one-hot logit by different factors, but found scaling by 1 to work well; scaling corresponds to changing the model’s confidence in its edit prediction but doesn’t change the prediction itself or the edit success.