Correcting Language Model Outputs by Editing Salient Layers (2024)

Abstract: Large language models can accumulate incorrect or outdated knowledge as the real world evolves. Compared to typical solutions such as retraining, retrieval augmented generation, model editing offers an effective yet low cost solution to address this issue. However, existing model editing algorithms employ manual selection of edit layers, which requires prior domain knowledge or expensive architecture-specific empirical layer selection methods, such as causal tracing. In this work, we propose SaLEM (Salient Layers Editing Model), an efficient solution for data driven layer selection for the model editing task. Our solution utilizes layer-wise saliency maps for layer selection, and matches the accuracy of prior approaches but with only 1/3 of their edits, enabling efficient updates to the parametric knowledge in large language models.

## introduction large language models (llms) are well known for their capacity to store extensive factual knowledge, which enables them to perform well in tasks such as question answering @xcite . however, facts can change after model training, which can introduce inaccuracies in model predictions and degrade downstream task performance. updating such factual knowledge is typically done by fine-tuning the model with corrected answers but this is an expensive approach and is prone to model overfitting @xcite . model editing @xcite techniques, such as mend @xcite and rome @xcite offer a practical and an effective alternative approach to address this problem, where we selectively edit a small subset of model parameters to update the factual knowledge. an important prerequisite to do model editing is to identify the network layers most likely to store the corresponding facts to be edited. for instance, in mend @xcite , this selection is done manually. while grounded in intuition, this approach depends on the model developer's domain knowledge and faces the risk of introducing superfluous edits. on the other hand, in rome, @xcite , the authors employ causal tracing, which attempts to locate facts in an autoregressive neural network model by identifying hidden states which have the strongest causal effect on predictions of given facts. while effective, it is unclear if this technique generalizes beyond decoder model architectures. more importantly, layer selection through causal tracing is extremely costly, requiring full two autoregressive passes through the entire network for each layer and token in the input sequence. in this work, we develop a new approach for automated layer selection for model editing called salem @xcite . salem leverages gradient values for given dataset with respect to the parameters of the llm to be edited to create layer saliency profiles @xcite and outputs the most salient layer to be edited. the salient layer selection method is an inexpensive, effective, and architecture-neutral approach for this task. we then thread the salient layer selection approach with mend to apply edits using decomposed gradients with respect to the selected layers. extensive experimental analysis established the effectiveness of salem. our main contributions in this work are as follows: ## related work model editors: the need to update and adapt knowledge representations of language models has traditionally been served through fine tuning (kenton and toutanova, 2019). various model editing strategies have been explored, including modified fine-tuning methods that enforce locality of edits @xcite or minimize l2-norm parameter updates for reliable edits @xcite , updating model beliefs based on learned optimizers @xcite . however, parameter space constraints may not always translate effectively into function space for neural networks @xcite . to address this, fine-tuning can be incorporated with a kl-divergence @xcite constraint, but this may not yield generalizable edits. editable neural networks (enn) @xcite and knowledge-editor (ke) @xcite use meta-learning techniques @xcite to effectively edit base models, offering alternative paths for desirable edit capabilities. but @xcite requires costly specialized training of the original network, while @xcite lacks tractability. mend @xcite was proposed as a resource-efficient approach for training large language models by leveraging rank-1 gradients in a novel parameter update scheme. unlike traditional gradient-based meta-learning algorithms @xcite @xcite , mend introduces adaptability post-hoc to a pre-trained model, enabling effective model adaptation without high computational costs. mend, however, lacks a datadriven method for identifying most effective layers to edit, and instead applies edits to statically predetermined layers. to address this gap, @xcite employed causal mediation analysis @xcite to trace hidden state activations within gpt @xcite . this helped identify and update parameters within the forward mid-layers mlps that are decisive for last subject token in factual associations. causal tracing, however, is an expensive parameter discovery mechanism requiring two full autoregressive passes through the model for each token, and has been recently shown not to always offer insights on the optimal mlp layer to edit @xcite . ## datasets to evaluate our approach, we conducted experiments on a diverse set of datasets encompassing text classification with varying levels of accuracy, question-answering and generation tasks. we consider five text classification datasets: i) fever-factchecking @xcite -fact checking with respect to wikipedia information, ii) multinli ( williams et al., 2018 ## experimental results we conduct detailed experiments, comparing salem with four competitive baselines: i) ft (fine-tuning), ii) enn (editable neural networks), iii) ke (knowledge editing) and iv) mend, and report results on two key evaluation metrics: ea (edit accuracy) and dd (drawdown). ## conclusion facts stored in llms routinely get outdated, and model editing offers an elegant solution to selectively update these facts without compromising the integrity of the model. however, existing algorithms suffer from shortcomings such as relying on domain knowledge or using computationally expensive mechanism for layer selection. to address these shortcomings, here we propose salem, an effective and computationally efficient solution for layer selection which utilizes parameter saliency maps aggregated at various levels. our experimental results demonstrate that by identifying the salient layer, salem matches the edit success of mend and rome with considerably better computational efficiency. further, detailed evaluation of salem across various nlp tasks, including natural language inference, classification, questionanswering, and generation, demonstrate its robust performance.