Decoding by Contrasting Knowledge:
Enhancing LLMs’ Confidence
on Edited Facts

We explore why in-context editing (ICE) effectively edits knowledge in large language models (LLMs). Using dedicated experiments, we capture the logits output of knowledge influenced by the edit. Our results show that incorporating new knowledge through ICE significantly increases the predicted probability of generating edited facts during the decoding process.

Formally, given a sequence of tokens \(X^E = \{x^E_1, x^E_2, ..., x^E_{m-1}\}\) which includes guidance of an editing prompt, the probability of the next token \(x^E_m\) with editing guidance is computed as:

$$ P_E(x^E_m|x^E_{\lt m}) = \text{softmax}(\phi(h^E_m)) $$

This contrasts with the parametric probability distribution \(P_B(x^B_n|x^B_{\lt n})\), which solely represents the response of LLMs based on their unedited knowledge.

Our analysis shows that the introduction of new knowledge through ICE shifts the probability distribution of new knowledge significantly, while the logits for parametric knowledge remain largely unchanged. This indicates that ICE enhances the logits of new knowledge, thereby improving LLMs' confidence in editing facts.

Inspired by our observations, we propose Decoding by Contrasting Knowledge (DeCK) to enhance ICE's ability to overcome stubborn knowledge in LLMs. DeCK consists of two key components: an editing enhancement module and a contrastive decoding strategy.

First, we improve attention to new knowledge by enhancing the probability distribution \(P_E\). We introduce the Knowledge Enhancement Divergence (KED) to measure the difference between the enhanced distribution and a target distribution \(Q\):

$$ \text{KED}(P||Q) = \frac{1}{2} \sum_{i=1}^n w_i \left( P(v_i) \log \frac{P(v_i)}{M(v_i)} + Q(v_i) \log \frac{Q(v_i)}{M(v_i)} \right) $$

Here, \(M = \frac{1}{2}(P + Q)\) is the average distribution, and \(w_i\) is a weight based on semantic relevance. The semantic relevance function \(s\) measures the relevance of a token \(v_i\) to the edited knowledge \(E\), defined as:

$$ s(v_i, E) = \max_{e_j \in E} \text{sim}(v_i, e_j) \cdot \phi(v_i) $$

where \(\text{sim}(\cdot, \cdot)\) is a similarity function, and \(\phi\) is a frequency-based weighting function. We enhance the logits \(\phi(h^E_m)\) using:

$$ \text{Enh}(\phi(h^E_m), s) = \alpha \cdot \phi(h^E_m) + \beta \cdot s $$

where \(\alpha\) and \(\beta\) are scaling coefficients.

Next, we apply a contrastive decoding strategy by comparing the enhanced new knowledge distribution with the original parametric distribution to predict the next token. The probability of the next token is computed as:

$$ \hat{P}_E(x^E_m) = \text{softmax} \left( F \left( P_E(x^E_m), P_B(x^B_n) \right) \right) $$

where the function \(F\) is defined as:

$$ F \left( P_E(x^E_m), P_B(x^B_n) \right) = \log \frac{P_E(x^E_m)}{P_B(x^B_n)} $$

To ensure plausibility, we introduce an adaptive plausibility constraint (APC) strategy, filtering out tokens with low probabilities in \(P_E(x^E_m)\) and considering only high-score tokens. The subset \(V_{\text{head}}(x^E_m|x^E_{\lt m})\) is defined as:

$$ V_{\text{head}}(x^E_m|x^E_{\lt m}) = \{ x^E_m \in V : P_E(x^E_m) \geq \lambda \max_w P_E(w) \} $$

where \(\lambda\) is a threshold parameter.

By contrasting these distributions, DeCK effectively amplifies the signal of new knowledge while suppressing the influence of parametric knowledge. Our experiments demonstrate that DeCK significantly boosts the confidence of LLMs in edited facts, enhancing the performance of ICE in editing stubborn knowledge.

We conducted extensive experiments using the MQUAKE-3K dataset and its derivatives to evaluate the effectiveness of DeCK. Our results show that DeCK significantly enhances the performance of ICE across various models and datasets. Specifically, DeCK has increased ICE’s editing success rate in LLAMA2-13B-CHAT by 63% and in LLAMA3-8B-INSTRUCT by 219%.

Our contributions can be summarized by three points. First, as far as we know, we are the first to elucidate superior performance of ICE on the KE from a model interpretability perspective. Second, we find that stubborn knowledge significantly impacts the performance of ICE, and we propose DeCK to boost confidence in editing facts, enhancing ICE to overcome it. Third, extensive experiments on \textsc{MQuAKE} indicate that our DeCK can effectively enhance the performance of ICE without altering the internal model or modifying external prompts. DeCK can be easily integrated into any ICE method as a decoding component to enhance editing capabilities. Our work paves the way to develop the both effective and accountable KE methods for LLMs.

In this paper, we introduced Decoding by Contrasting Knowledge (DeCK), a novel decoding strategy aimed at enhancing in-context editing to overcome stubborn knowledge for LLMs. DeCK significantly improves editing accuracy but requires the reception of input from two different token sequences during the generation process, resulting in increased latency. Future work may focus on optimizing the transformers architecture or exploring alternative, more cost-effective versions of DeCK.

Please cite our paper if it's helpful to your work!

@misc{bi2024decoding,
  title={Decoding by Contrasting Knowledge: Enhancing LLMs' Confidence on Edited Facts}, 
  author={Baolong Bi and Shenghua Liu and Lingrui Mei and Yiwei Wang and Pengliang Ji and Xueqi Cheng},
  year={2024},
  eprint={2405.11613},
  archivePrefix={arXiv},
  primaryClass={cs.CL}
}