Precise Localization of Memories: A Fine-grained Neuron-level Knowledge Editing Technique for LLMs

Memory-based For memory-based editors, some specific modules store the edit knowledge are used for post-edit response. SERAC (Mitchell et al., 2022b) stores edits in an explicit memory and learns to reason over them to modulate the base model’s predictions as needed. Zheng et al. (2023) explores in-context knowledge editing (IKE), a method without any gradient and parameter updating. GRACE (Hartvigsen et al., 2024) is a lifelong model editing method that implements spot-fixes on streaming errors of a deployed model, ensuring minimal impact on unrelated inputs. Recently, Yu et al. (2024) proposes MELO, a novel method that alters the behavior of LLMs by dynamically activating certain LoRA blocks according to the index built in an inner vector database.

Meta-learning Based on hypernetwork, several meta-learning methods have been proposed to edit models. De Cao et al. (2021) presents KnowledgeEditor (KE), a method which can be used to edit knowledge and, fix bugs or unexpected predictions without the need for expensive re-training or fine-tuning. MEND (Mitchell et al., 2022a) introduces 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. Hase et al. (2023b) proposes SLAG, a training objective for sequential, local, and generalizing updates with a better performance. Han et al. (2023) proposes a novel divide-and-conquer framework, drawing on dynamic inference to break the zero-sum phenomenon in multiple edits.

Locate-then-edit Although prior research has explored knowledge storage mechanisms, the precise methods by which LLMs retain knowledge remain unclear. Studies have indicated that knowledge is often embedded within FFNs (Geva et al., 2021; 2022; Dai et al., 2022). Building on these, locate-then-edit methods have gained traction by first locating specific regions of knowledge storage and then executing targeted editing. A leading example is ROME, which innovatively employs causal tracing to pinpoint parameters intended for edits and directly updates them (Meng et al., 2022). This foundational work has paved the way for additional methods such as MEMIT (Meng et al., 2023), PMET (Li et al., 2024), and EMMET (Gupta et al., 2024), enhancing the capacity to incorporate and modify larger quantities of knowledge. The advantages of locate-then-edit methods include increased precision in knowledge modification and the ability to selectively edit specific information, making them a vital advancement in the ongoing development of more effective and reliable LLMs.

Transformer (Vaswani et al., 2017) is one of the most successful architectures designed for different tasks (Song et al., 2024b; Li et al., 2023) and there has been increasing interest in interpreting and analyzing the internal mechanisms of transformer-based models. Previous research has aimed to characterize the types of information encoded in individual neurons. Dai et al. (2022) explores the identification of “knowledge neurons”, which encode specific commonsense knowledge acquired during pre-training. Additionally, Wang et al. (2022) presents a technique for identifying “skill neurons” in pre-trained transformer-based language models, which are crucial for specific tasks. Schwettmann et al. (2023) explains how LLMs convert visual representations into corresponding texts by introducing a procedure for identifying “multimodal neurons”. More recently, Pan et al. (2024) proposes a novel method for finding “multi-modal neurons”, which elucidates how multi-modal LLMs (Zeng et al., 2024) bridge visual and textual concepts for captioning (Song et al., 2024a).

Following previous work (Dai et al., 2022; Wang et al., 2022; Schwettmann et al., 2023; Pan et al., 2024) on selecting neurons in Transformer-based models, we present a neuron localization method for knowledge editing. Specifically, we hypothesize that a knowledge $(s_{j},r_{j},o_{j})$ is stored in specific neurons, which are activated when LLMs receive input $(s_{j},r_{j})$, exhibiting a tendency to produce the output $o_{j}$. Therefore, our objective is to quantify contribution of each neuron to the current output and locate those neurons with higher impact. Following Pan et al. (2024), who calculates contribution scores in multi-modal LLMs, we similarly compute contribution scores within LLMs.

For each token $t$ in the output $o_{j}$, we compute contribution score for each neuron $u_{i}$ at layer $l$ as:

where ${\bm{q}}^{l}_{i,-1}$ is the activation output at the last token for neuron $u_{i}$ at layer $l$, $(\cdot)_{t,i}$ represents the $t$-th row and $i$-th column of the input matrix, and $\mathbf{W}_{u}$ is the unembedding matrix.

Here we regard $\mathbf{W}_{u}\mathbf{W}_{\text{out}}^{l}\in\mathbb{R}^{v\times d_{m}}$ as a projection function projecting from activations of the neurons to distribution of the vocabulary, where $d_{m}$ is the intermediate size and $v$ is the vocabulary size and regard ${\bm{q}}^{l}_{i,-1}$ as a coefficient of the projection, respectively. This projection explicitly demonstrates the varying levels of focus that different neurons pay to different tokens, enabling us to calculate the contribution score. We provide detailed derivation in Appendix A.

After quantifying the contribution of each neuron, we rank all scores of neurons across all layers by the descending order and pick out top-$k$ neurons, denoted as $u^{k}$. These neurons are regarded as carriers of knowledge $(s_{j},r_{j},o_{j})$ and make significant contributions to the output $o_{j}$. We follow the same procedure to locate neurons for each token $t$ in $o_{j}$, and use the set $\mathcal{U}_{j}=\left\{u^{k}_{1},u^{k}_{2},\cdots,u^{k}_{|o_{j}|}\right\}$ to represent neurons of $o_{j}$, where $|o_{j}|$ means the token length of $o_{j}$. Algorithm 1 summarizes the entire process of neuron localization.

We locate key neurons $\mathcal{U}_{j}$ of $o_{j}$ as described above, and then modify the model weights corresponding to locations of selected neurons to update the knowledge. For each neuron $u\in\mathcal{U}_{j}$, we assume that $u$ is the $i$-th neuron at layer $l$. Then we compute a vector ${\bm{z}}\in\mathbb{R}^{d_{h}}$ and add it to the $i$-th row of matrix $\mathbf{W}_{\text{out}}^{l}$ for updating, where $d_{h}$ is the hidden size. If we stack vector ${\bm{z}}$ for each neuron $u$ as ${\bm{Z}}_{j}=[{\bm{z}}_{1}~{}|~{}{\bm{z}}_{2}~{}|~{}\cdots~{}|~{}{\bm{z}}_{|\mathcal{U}_{j}|}]$, our objective can be succinctly represented as learning an optimized ${\bm{Z}}_{j}$ based on $\mathcal{U}_{j}$, which is then applied to the model $\mathcal{M}$, resulting in a post-edited model $\mathcal{M}^{\prime}$. Following Meng et al. (2022), the objective $\mathcal{L}({\bm{Z}}_{j})$ consists of editing loss $\mathcal{L}_{\text{edit}}({\bm{Z}}_{j})$, KL divergence $\mathcal{L}_{\text{KL}}({\bm{Z}}_{j})$ and repetition penalty loss $\mathcal{L}_{\text{pen}}({\bm{Z}}_{j})$. The editing loss utilizes negative log-likelihood to maximize the probability of the target $o^{*}_{j}$:

During the editing process, we aim to avoid altering unrelated knowledge or impacting the model’s language capabilities. To this end, we add a KL divergence constraint of prompt that contains the subject and relation to the model, which is calculated by:

where $\mathbb{P}^{\prime}\left[\cdot\right]$ represents the probability distribution of output from position 1 to position $\ell_{p}-1$, assuming the length of the input prompt is $\ell_{p}$, which is different from $\mathbb{P}\left[\cdot\right]$. Except KL divergence, to prevent the post-edited model from generating the editing target $o^{*}_{j}$ repeatedly, we also introduce a repetition penalty constraint. At the last position of the complete prompt $p(s_{j},r_{j},o^{*}_{j})$, we use negative log-likelihood to maximize the probability of not generating $o^{*}_{j}$:

Finally, we compute a weighted sum of editing loss, KL divergence and repetition penalty loss:

where $\alpha$ and $\beta$ are hyperparameters.

In language models, the later layers are closely tied to the model’s language capabilities (Geva et al., 2021; Dai et al., 2022; Wang et al., 2022; Pan et al., 2024). Arbitrary modifications to these later layers may impair model’s linguistic abilities and result in responses with lower quality. To ensure the stability of LLMs, we implement layer freezing (LF) in our method. Specifically, for a LLM with $L$ layers, when locating neurons, we exclude the last $l_{f}$ layers, focusing only on the first $L-l_{f}$ layers. This ensures that no modifications are made to the last $l_{f}$ layers during the editing process.

Models and datasets. We conduct experiments on the KnowEdit (Zhang et al., 2024) benchmark with GPT-J-6B (28 layers) (Wang & Komatsuzaki, 2021), LLaMA-2-7B (32 layers) (Touvron et al., 2023b) and LLaMA-3-8B (32 layers) (Dubey et al., 2024). KnowEdit is an integrated benchmark for evaluating various knowledge editing methods, which contains six datasets for different evaluation types. We select three datasets including knowledge insertion and knowledge modification in our experiments: $\text{WikiData}_{counterfact}$ (Cohen et al., 2024), $\text{WikiData}_{recent}$ (Cohen et al., 2024) and ZsRE (Levy et al., 2017). Notably, the locality evaluation in KnowEdit primarily focuses on changing the relation. The proportion of prompts where the subject changes in datasets $\text{WikiData}_{counterfact}$, $\text{WikiData}_{recent}$ and ZsRE is only 0.9%, 0.1%, and 0.0%, respectively.

Baselines. We categorize baseline methods into three types of knowledge editing. The first category consists of methods that directly modify model parameters, such as Fine-Tuning (FT) and LoRA (Wu et al., 2023). The second category includes memory-based methods, for which we select In-context Knowledge Editing (IKE) (Zheng et al., 2023), which retrieves the most pertinent demonstrations. The third category focuses on locate-then-edit methods, which are central to our study. Although Knowledge Neurons (KN) is also a neuron-level knowledge localization method, it employs a significantly different technique than ours, selecting neurons via gradient-based attributions and modifying the corresponding FFN weights by adding scaled embedding vectors. Importantly, ROME (Meng et al., 2022), as a pioneer of causal tracing localization, has further advanced the locate-then-edit methods and significantly influenced the field, while MEMIT (Meng et al., 2023) has built upon this foundation with notable enhancements. PMET (Li et al., 2024) serves as an improvement over MEMIT. Both ROME and MEMIT not only represent critical developments but have also achieved substantial popularity, making them essential comparisons in our work.

Evaluation metrics. As described in § 3, we adopt four evaluation metrics in our experiments: Edit Success, Portability, Locality, and Fluency (Zhang et al., 2024). Portability contains three parts: Subject Aliasing Accuracy (SAA), Logical Generalization Accuracy (LGA) and Reasoning Accuracy (RA). Subject aliasing replaces the question’s subject with an alias or synonym to evaluate performance on other descriptions of the subject. Logical generalizations are changes that are semantically related to the modified fact and expected to change by the edit. Reasoning examines the reasoning ability with changed facts. Locality consists of two parts: Forgetfulness Accuracy (FA) and Relation Specificity Accuracy (RSA). Forgetfulness evaluates whether the post-edited model retains the original objects in one-to-many relationships, whereas relation specificity evaluates whether any other attributes of the subject, which have been previously updated remain unaltered.

In Table 5.2, we show quantitative editing results on $\text{WikiData}_{counterfact}$. Our approach demonstrates the best Edit Success, Portability and Locality among various locate-then-edit methods. We observe that previous locate-then-edit methods with causal tracing localization perform poorly when handling similar but unrelated knowledge, exhibiting generally low Locality. To achieve better editing results, we sacrifice some Fluency but without compromising the original model’s language capabilities. Editing results on $\text{WikiData}_{recent}$ and ZsRE can be found in Appendix D.

In this section, we present ablation study to assess the impact of various components on the overall performance of our method. Specifically, we first test the impact of removing neuron localization and layer freezing on performance. Next, we investigate effects of restricting neuron localization to a single layer and explore how varying number of selected neurons affects editing. Finally, we examine results of removing KL divergence and repetition penalty constraints in the editing process.

Removing neuron localization and layer freezing. Since our approach also employs a locate-then-edit methodology, it is essential to verify the effectiveness of the initial localization step. To this end, we maintain the editing process unchanged and conduct experiments by replacing carefully selected neurons with randomly selected ones. Table 2 lists ablation results. When neurons are selected at random, both Edit Success and Portability demonstrate varying degrees of decline, particularly evident in SAA metric, suggesting that our chosen neurons are sensitive to the knowledge being edited. In contrast, ROME experiences only a slight decrease in performance without localization, supporting the hypothesis that causal tracing is not essential. On the other hand, we assess the effectiveness of layer freezing. As shown in Table 2, without layer freezing, the model’s language capabilities are compromised, leading to a significant drop in fluency. It is speculated that even minor modifications, when applied to the last few layers, can result in catastrophic consequences.

Restricting neuron localization to a single layer. In the previous experiments, we do not restrict modifications to a specific layer as in ROME. Instead, we determine which layers to modify solely based on the scores (note that, due to layer freezing, the last few layers are not altered). We now manually restrict neuron localization to a single layer and modify only the model weights within that layer to determine whether this manual intervention yields better results. For the specific layer $l\in\{5,10,15,20,25\}$, only neurons in layer $l$ will be selected. Results with LLaMA-2 and LLaMA-3 are listed in Table 3 and Table 10, respectively. We observe that although specifying a particular layer sometimes performs better on some metrics (e.g., RSA and FA of layer 25), overall performance across all metrics is still optimal when no layer restrictions are applied. While manually restricting neuron localization to a single layer can be effective based on experience, relying on our algorithm to automatically locate neurons may be a more appropriate option in the absence of prior information.

Varying number of selected neurons. During the editing process, the number of selected neurons likely influences the extent of modifications to the model — the more neurons selected, the greater the number of model parameters altered, and vice versa. Therefore, we vary the number of selected neurons to observe the changes in the metrics. For each number of neuron $k\in\{1,3,5,10,15,20\}$, top-$k$ neurons are selected for each token. Figure 2 plots metric curves on LLaMA-2. We can observe that as the number of neurons increases, Portability (i.e., SAA, LGA and RA) generally improves while Locality (i.e., RSA and FA) tends to slightly decrease. This suggests that selecting a greater number of neurons may provide a more comprehensive localization and further enhance the model’s ability to update the targeted knowledge, but it also increases the risk of unintentionally altering unrelated memories. Results on LLaMA-3 could be found in Appendix D.

Removing KL divergence and repetition penalty constraints. To minimize the impact on the model’s inherent language capabilities, we adopt KL divergence and repetition penalty constraints during the editing process. Table 4 lists results of removing these constraints in cases with and without using LF. When using LF, the effects of KL divergence and repetition penalty constraints are not significant; however, when LF is not applied, we observe that (1) KL divergence constraint is important for the locality of model editing, and removing it leads to a significant decline in the RSA metric. (2) Repetition penalty constraint has minimal impact on portability and locality but significantly affects fluency. Without it, the post-edited model is more likely to produce repetitive text (e.g., “The next host city of the Olympic Games is Los Angeles Los Angeles Los Angeles …”).

Efficiency evaluation. A key advantage of FiNE, due to its fine-grained approach, is its notable efficiency. To quantify this, we first examine the number of modified parameters, as detailed in Table 5. Both ROME and MEMIT modify the weights of the second layer in FFNs, resulting in a substantial number of parameter modifications, ranging from $10^{7}$ to $10^{8}$. In contrast, FiNE only edits a subset of neurons, reducing the number of modified parameters to approximately $10^{4}$, which allows for a more fine-grained and precise editing in LLMs. Additionally, we assess editing time and memory usage at Float32 and Float16 precision in Figure 5 and Figure 9, respectively. FiNE exhibits a significant time advantage over ROME and MEMIT, particularly at Float32 precision, being approximately 4$\times$ to 6$\times$ faster. In terms of memory usage, FiNE also offers a slight benefit for LLaMA-2 and LLaMA-3.

Localization analyses. We analyze our neuron-level localization from two perspectives: distributions and textual meanings. (1) We plot the distribution of unique neurons located by FiNE (see Figure 5). We aggregate statistics for all located neurons across the entire $\text{WikiData}_{counterfact}$ dataset. We observe that these key neurons widely occur in higher layers, which is consistent with previous work (Wang et al., 2022; Pan et al., 2024), but different from layers that ROME and MEMIT edit. (2) For insight into neurons filtered by Eqn. 3, we follow the Logit Lens (nostalgebraist, 2020; Zhong et al., 2022; Geva et al., 2022), which converts hidden states into a set of logits for each vocabulary token. Similarly, we investigate neurons’ textual meanings by sorting rows of the multiplication of the unembedding matrix and the second layer of FFN and regarding top tokens as each neuron represents (Pan et al., 2024). Table 5 shows an example, which indicates that neurons selected by FiNE are highly related to the source knowledge. We list more examples in Appendix D.

Editing method scaling. In knowledge editing, the ability to simultaneously edit multiple knowledge facts is a crucial objective that enhances the practical application of various methods. Several approaches (Meng et al., 2023; Li et al., 2024; Gupta et al., 2024) have been specifically developed to achieve this goal. Although our method does not incorporate a specialized design for this purpose, we posit that our more precise localization may reduce the inter-dependencies among different knowledge facts during the editing process, thereby intuitively contributing to improved editing scalability. To verify our hypothesis, we progressively increased the scale of editing targets from 100 to 800. Figure 6 plots experimental results with GPT-J. ROME struggles significantly when handling 100 edits, and ceases to function effectively as the number of edits increases further, with all metrics approaching zero. We observe that our method continues to operate effectively even when handling a larger number of edits, although performance is lower compared to single-instance editing. Additionally, we unexpectedly find that our method closely matches MEMIT on metrics FA, SAA, LGA, and RA. We attribute this to our fine-grained neuron-level localization approach, which only modifies a small number of neurons, and results in subtle but crucial changes to LLMs.

Case study. Table 6 provides an example of the editing and locality testing results across different methods. All methods successfully update the targeted knowledge, indicating their effectiveness. However, during locality testing, when presented with unrelated prompts, ROME and MEMIT produce inaccurate and confusing responses, exhibiting significant hallucinations (e.g., Czech Republic and Salzburg). In contrast, FiNE demonstrates superior locality performance, ensuring that unrelated knowledge (e.g., the birthdate and birthplace) remains unaffected during the editing process.