Mitigating Heterogeneous Token Overfitting in LLM Knowledge Editing

Given a text $\boldsymbol{x}=(x_{1},\dots,x_{n})$, where each $x_{i}\in\mathcal{V}$ is a token from vocabulary $\mathcal{V}$, a large language model (LLM) parameterized by $\theta$ computes probability $\pi_{\theta}(\boldsymbol{x})$ based on chain rule (Bengio et al., 2000):

where $\pi_{\theta}(x_{i}\mid\boldsymbol{x}_{<i})$ is the predicted distribution of token $x_{i}$ given previous $\boldsymbol{x}_{<i}$. The LLM is usually trained with maximum likelihood estimation (Hochreiter, 1997; Sutskever, 2014; Cho et al., 2014). To generate a sentence $\boldsymbol{x}$, the LLM computes $\pi_{\theta}(x_{i}\mid\boldsymbol{x}_{<i})$ and draws $x_{i}$ from it; then $x_{i}$ is combined with $\boldsymbol{x}_{<i}$ as new inputs for future steps. This process completes if a special token that marks the end of the sentence is returned, or if the maximum length is reached.

Knowledge Editing (KE) aims to update specific knowledge in a pre-trained LLM while preserving unrelated others. A knowledge can be represented by natural language $(\boldsymbol{x},\boldsymbol{y})$, $\boldsymbol{x}$ describes the subject and relation, and $\boldsymbol{y}$ entails corresponding object. For instance, suppose $\boldsymbol{x}$ is The president of United States is, $\boldsymbol{y}$ can be Donald Trump. KE asks the LLM to respond given $\boldsymbol{x}$ with new $\boldsymbol{y}$, while satisfying the following criteria meanwhile (Zhang et al., 2024c): (1) Generality: the edited model should generalize to all equivalent inquires about the US president. (2) Portability: questions reasoned from the new knowledge such as the first lady of United States should be answered correctly. (3) Locality: unrelated knowledge such as the prime minister of Canada should be unchanged. These requirements of precisely updating specific knowledge proves non-trivial (Wang et al., 2023b; Zhang et al., 2024c).

In response to precise KE requirements, existing attempts restrict the updates to only a minimal amount of parameters. This design establishes remarkable progress in maintaining good locality (Zhang et al., 2024c; Wang et al., 2024b). However, it proves insufficient to maintain good generalizability (generalilty and portability) due to the so-called overfitting issue (Zhong et al., 2023; Zhang et al., 2024a).

Namely, many KE tasks involve one piece of new knowledge at a time, requiring to fine-tune an LLM on single training instance. In such challenging scenarios, the LLM often encounters severe overfitting even only a few parameters are updated. This greatly restricts its ability to generalize the edited knowledge. As shown in Zhong et al. (2023); Zhang et al. (2024a), edited LLMs usually pay excessive attention to the edited subject, but fail to address multi-hop reasoning questions involving the new knowledge. As a result, this limitation results in suboptimal portability.

As a direct evidence, Fig 1 shows the change of generality and portability loss¹¹1The perplexity loss of the ground truth answer to a question. at different iterations from fine-tuning LLaMA2 7B (Touvron et al., 2023) with LoRA, a representative KE baseline method (Zhang et al., 2024c). As the training goes on, the generality loss decreases. However, the portability loss decreases at the beginning of training, but starts to increase later. This confirms the existence of overfitting. More importantly, the ultimate portability loss is significantly larger than before editing, indicating that the reasoning ability is in fact undermined by the KE process,

Towards a deeper understanding of this overfitting phenomenon, we check the loss of each token, and find that different tokens tend to have distinct initial loss values. As depicted in Fig 2(a), before editing LLaMA2, only certain tokens (e.g., the beginning) have significant loss values. On the other hand, some tokens take small loss value and are initially-fitted by nature. As an intuitive explanation, consider the previous US president example. No matter a user wants to edit the answer to Donald Trump or Joe Biden, after seeing the first word Donald or Joe as a hint, the LLM is expected to be capable of infer the remaining part based on its pretrained knowledge.

Nonetheless, existing KE methods overlook this token-level difference. Consequently, they tend to overfit tokens that have varied losses at different speeds. For verification, we compute the pre-edited log-likelihood of tokens generated by the model with greedy decoding, and that of the editing instance during the KE process. We define underfitting degree (UD) as the difference between the pre-edited and running log-likelihood, negative UD indicates an overfitting. Fig 2(b) shows UD of different tokens when half of them are overfitted. Strong pattern of UD varies across different tokens confirms our concern. We dub this issue as heterogeneous token overfitting (HTO) of KE.

HTO’s direct cause lies in the training paradigm. Formally, given editing instance $(\boldsymbol{x},\boldsymbol{y}=[y_{1},\dots,y_{m}])$ where $\boldsymbol{y}$ contains $m$ tokens, many KE methods resort to a conventional LLM training objective²²2We restrict our study to the widely-used teacher-forcing mechanism (Lamb et al., 2016).. In particular, they seek to maximize likelihood of $\pi_{\theta}(\boldsymbol{y}\mid\boldsymbol{x})$ by minimizing an averaged cross-entropy (CE) loss with gradient descent on

Here ${\boldsymbol{c}}_{i}=\boldsymbol{x}\oplus\boldsymbol{y}_{<i}$ denotes the context for token $y_{i}$, $\delta_{y_{i}}(y)$ is the Kronecker delta function³³3$\delta_{y_{i}}(y)=1$ if $y=y_{i}$ else 0., and $\text{CE}[\cdot\|\cdot]$ computes CE between two distributions.

During training, gradient $\nabla_{\theta}\ell_{\text{CE}}(\theta)$ maximizes the probability of $y_{i}$ whiling minimizing the probabilities of all other candidates. When the model is repeatedly updated using gradient(s) from the single datapoint, as in KE, the probabilities of initially-fitted tokens become disproportionately large, while tokens with high initial loss values are gradually fitted. That is to say, HTO lies in indiscriminately optimizing CE loss of all tokens, without considering their difference. Existing attempts for mitigating overfitting such as early stopping (Yao et al., 2007) and label smoothing (Szegedy et al., 2016; Müller et al., 2019) also ignore this token-level difference, making them conceptually less suitable for HTO.

We present OVERTONE, a token-level strategy for HTO mitigation. Our method smooths $\boldsymbol{y}$’s distribution for fitting in an adaptive way. Specifically, we replace each delta distribution $\delta_{y_{i}}(y)$ with a unique smoothed target distribution $\pi_{\text{tar}}(y\mid{\boldsymbol{c}}_{i})$, and refine the cross entropy by a clipped forward KL divergence. Our complete loss is given by

where clipped $\max(\cdot,\epsilon)$ imposes a token-level early stopping when predicted $\pi_{\theta}$ is close enough to $\pi_{\text{tar}}$.

Principles of $\pi_{\text{tar}}$ design. We note that two principles should be met in order to make $\pi_{\text{tar}}$ a good distribution to target on. First, $\pi_{\text{tar}}$ should convey that ground truth token $y_{i}$ is most probable, otherwise, the objective may lead to incorrect knowledge. Second, compared to uniform prior that smooths all tokens equally, the model’s own pre-trained knowledge is a better prior to help mitigate forgetting problem (Zhang & Sabuncu, 2020; Lee et al., 2022).

In light of the two principles, we use $\delta_{y_{i}}$ and the LLM’s current knowledge from its predicted distribution $\pi_{\theta}$ to construct target $\pi_{\text{tar}}$. However, as will be verified later, directly use $\pi_{\theta}$ can be suboptimal due to the non-negligible noise it carries (Hewitt et al., 2022; Tang et al., 2024). Specifically, Tang et al. (2024) argued that $\pi_{\theta}$ mixes a distinct subset of informative tokens, and a subset of noisy tokens associating with small logits that fall outside $n\sigma$-distant away from the maximal value. By filtering out noisy tokens in $\pi_{\theta}$, the LLM performance can be boosted at inference time. We bring this insight to the training (editing) phase and mix the filtered distribution⁴⁴4For brevity $\pi_{\text{flt}}^{(i)}=\pi_{\text{flt}}(y\mid{\boldsymbol{c}}_{i})$, $\pi_{\text{tar}}^{(i)}$ is defined similarly. Plain $\pi_{\text{flt}}$ and $\pi_{\text{tar}}$ will be used when discussing the general idea. $\pi_{\text{flt}}^{(i)}$ with $\delta_{y_{i}}$ by

where $\lambda$ is a hyper-parameter. Namely, we adopt the candidate mixture $\pi_{\text{tar}}^{\text{can}}$ if it correctly assigns the maximal probability to $y_{i}$, otherwise, we skip the mixing and use $\delta_{y_{i}}$. This skip mechanism helps reduce potential knowledge conflicts by discarding $\pi_{\text{flt}}^{(i)}$ (from $\pi_{\theta}$) when it heavily relies on outdated knowledge, which often happens in the first few training steps, empirical benefit is shown in Sec 4.4. Algo 1 outlines the process of our solution.

This section provides theoretical analysis on key factors that merit OVERTONE for KE. All proofs and more in-depth technical background are deferred to App A.

Merit 1. OVERTONE is universal and efficient.

Merit 2. OVERTONE provides better updates.

OVERTONE leads to more effective parameter updates, as demonstrated through the lens of the influence function (Koh & Liang, 2017), outlined in the following informal theorem. Due to page limitations, the formal version and corresponding assumptions are deferred to Appendix A.3.

Merit 3. OVERTONE has close connection to DPO and other constrained optimizations.

One might question whether OVERTONE is conceptually superior to constrained optimization approaches, such as fine-tuning only a small set of specific parameters (Dong et al., 2022a; Dai et al., 2021), limiting update magnitudes (Zhu et al., 2020), or employing low-rank updates (Hu et al., 2021). We emphasize that OVERTONE introduces a new objective that can be solved with any optimization methods, regardless of whether constraints are imposed. In other words, OVERTONE can be seamlessly combined with existing constrained optimization-based solutions for KE.

Compared with conducting explicit DPO, OVERTONE does not require collecting preference data, and is more efficient thereof. Furthermore, as highlighted in Rozner et al. (2024), another challenge of applying DPO to KE is that determining win-loss data pairs can be unstraightforward in KE. In contrast, OVERTONE walks around this challenge by refraining from treating any token as unpreferred, and instead acts on a distribution level.

Base Models. We conduct experiments on two representative LMs, LLaMA 2-7b-Chat (Touvron et al., 2023) and LLaMA 3-8b-Instruct (Dubey et al., 2024), which have been widely studied in the literature (Zhang et al., 2024c; Wang et al., 2024b). From now on, we refer to the two LMs as LLaMA 2 and LLaMA 3 for brevity.

Tasks. Following Wang et al. (2023b); Zhang et al. (2024c), we edit different kinds of knowledge: WikiData${}_{\text{recent}}$, WikiData${}_{\text{counterfact}}$ (Cohen et al., 2024), WikiBio (Hartvigsen et al., 2024), and ZsRE (Yao et al., 2023). Besides the four popular benchmarks, we also explore more complex MQuAKE (Zhong et al., 2023; Wang et al., 2024d). Due to page limitation, we refer readers to Zhang et al. (2024c) for more benchmark details. When editing an LLM, we consider two scenarios: (1) Single Editing: one piece of knowledge is edited at a time. (2) Continual Editing: multiple pieces of knowledge are edited in a sequential way. This is more challenging due to forgetting and knowledge conflicting (Hartvigsen et al., 2024; Wang et al., 2024b).

Editing Methods. We apply OVERTONE to four representative KE methods from different families that have achieved state-of-the-art performance (Zhang et al., 2024c; Wang et al., 2024c). FT-M (Zhang et al., 2024c) fine-tunes a special layer identified by causal-tracing analysis wherein the knowledge is stored. LoRA (Hu et al., 2021) learns additive low-rank updates for model parameters on the new knowledge. MELO (Yu et al., 2024) and WISE (Wang et al., 2024b) incorporates additional parameter copies to learn new knowledge, along with some gating mechanism to determine whether original or new knowledge should be used at inference time. Despite incorporating certain explicit or implicit constraints on the learnable parameters, these methods are all trained to minimize the CE loss. For better benchmarking, we also report results from two widely-studied methods ROME (Meng et al., 2022a) and MEMIT (Meng et al., 2022b). ROME applies a causal-tracing analysis to identify the layer wherein the knowledge is stored and then solves an analytic rank-one update, and MEMIT extends ROME by identifying a series of layers to edit and finding the updates as least-squares solutions. To reflect the challenging nature of KE under data scarcity regime, we focus on KE methods that do not require a larges-scale hard-to-access training data, or training additional models. No data augmentation were applied during the editing.

Evaluation Criteria. We evaluate the performance from four aspects as discussed in Sec 2: reliability (Rel.), generality (Gen.), portability (Por.), and locality (Loc.). Due to page limits we refer readers to Zhang et al. (2024c); Wang et al. (2024b) for their formulations. We report the average of different metrics for more complete comparisons.

Implementation Details. All of our experiments are implemented in EasyEdit (Wang et al., 2024c). More details and hyper-parameters can be found in App B.

We evaluate the effectiveness of OVERTONE in conducting Single Editing on ZsRE, WikiData${}_{\text{recent}}$, WikiData${}_{\text{counterfact}}$, and WikiBio with different KE methods. WISE was tested on ZsRE, the only benchmark that contains additional irrelevant data during the editing time that is required by WISE.

Single Editing results are reported in Tab 4.2. From the table, all KE methods gained significant improvement from the proposed OVERTONE paradigm. Specifically, The four methods hardly performed comparable to baselines ROME and MEMIT from normal training, but were capable of exceeding them when trained with OVERTONE. For instance, without OVERTONE, ROME achieved the highest and the second-highest average performance for editing LLaMA 2 and LLaMA 3 respectively on Wiki${}_{\text{recent}}$. However, when equipped with OVERTONE, FT-M, LoRA, and MELO outperformed ROME on both tasks.

We next check where the improvement was made. from the table, the first gain was from improved portability. To see this, note that when editing LLaMA 2 on ZsRE, LoRA reached a portability that was nearly three times of the base version. Similarly, MELO also reached an almost doubled portability. More evidence can be found from editing LLaMA 3 as well. In addition, all methods, especially those initially fall short in maintaining good locality, achieved excellent performance in this regard. As an evidence, LoRA’s reached a nearly five times locality improvements when editing both LLaMA 2 and LLaMA 3 on Wiki${}_{\text{counterfact}}$. We want to highlight that, all these improvements were made without compromising editing reliability. That is to say, all the four methods achieved better trade-offs between reliability and reasoning (and locality) from the proposed OVERTONE. More importantly, this success was established in a model-agnostic manner, in the sense that OVERTONE is not specialized for any particular KE method studied here. Instead, it offers a highly flexible and generic paradigm that can be combined with existing solutions in a plug-and-play manner.

More Complex Editing task. To further evaluate how OVERTONE performs on complex benchmark in the filed of KE, we test FT-M and LoRA with editing the two LLMs on MQuAKE-2002 (Wang et al., 2024d)⁵⁵5This is a cleaned version of MQuAKE by fixing knowledge conflicts (Wang et al., 2024d)., following Zhong et al. (2023). This task requires the edited LLM to answer one- and two-hops reasoning questions about the edited knowledge. Experiment results are reported in Table 4.2. As before, OVERTONE was capable of achieving better portability without hurting the editing performance.

These empirical results echo well with our theoretical analysis, and confirm the superiority of OVERTONE.