Adaptive Token Biaser: Knowledge Editing via Biasing Key Entities (2024)

Abstract: The parametric knowledge memorized by large language models (LLMs) becomes outdated quickly. In-context editing (ICE) is currently the most effective method for updating the knowledge of LLMs. Recent advancements involve enhancing ICE by modifying the decoding strategy, obviating the need for altering internal model structures or adjusting external prompts.However, this enhancement operates across the entire sequence generation, encompassing a plethora of non-critical tokens.In this work, we introduce **A**daptive **T**oken **Bias**er (ATBias), a new decoding technique designed to enhance ICE.It focuses on the tokens that are mostly related to knowledge during decoding, biasing their logits by matching key entities related to new and parametric knowledge.Experimental results show that ATBias significantly enhances ICE performance, achieving up to a 32.3% improvement over state-of-the-art ICE methods while incurring only half the latency.ATBias not only improves the knowledge editing capabilities of ICE but can also be widely applied to LLMs with negligible cost.

## introduction large language models (llms) @xcite @xcite accumulate a substantial volume of factual knowledge during pretraining. however, some of this knowledge may quickly become outdated, resulting in decreased reliability of llms @xcite @xcite . due to the substantial cost associated with retraining, knowledge editing (ke) @xcite @xcite has been proposed to update the knowledge in llms by injecting new knowledge or modifying parametric knowledge. as currently the most effective ke method, incontext editing (ice) @xcite @xcite indicate that editing stubborn knowledge solely through external context prompts is challenging, as this knowledge has been established in llms with strong confidence during pre-training, as illustrated in figure 1 . recent stateof-the-art ice method deck @xcite enhances the editing of stubborn knowledge by modifying entire generating sequence during decoding. however, this approach carries potential risks, not only introducing the possibility of inference errors but also incurring higher latency costs. in this work, we explore enhancing ice for editing stubborn knowledge during the decoding stage of llms, without altering internal llms' parameters or modifying external prompts. we propose adaptive token biaser (atbias), a new ke framework for llms that enhances ice by matching key entities and biasing the logits of specific tokens. the framework of atbias is shown in figure 2 . unlike previous decoding @xcite @xcite , atbias focuses more on the matched tokens rather than the entire generated sequence. we argue that modifications on other tokens are unnecessary, leading to redundant computational costs and even mistakes. for example, in generating text the author richard dawkins wrote "misery", the key terms "richard" and "dawkins" merit attention over other words in the text. indiscriminate adjustments to other words (such as "the", "author", etc.) can pose a potential risk of introducing fundamental errors in the logical coherence of the entire inference statement. the main goal of atbias is to increase the generation probability of tokens related to new knowledge while decreasing that of parametric knowledge. capturing key textual entities is a prerequisite for matching crucial tokens. atbias provides a parametric induction and entity extraction module, which can efficiently extract key entities from both new facts and parametric facts induced from llms. we also introduce knowledge caching, enabling the aforementioned process to be completed offline. this ensures our atbias performs efficient editing with only a single inference. we design a specialized filtering mechanism that ensures our approach only considers topranked and high-probability predicted tokens. the probabilistic-ranking filter not only significantly reduces the likelihood of implausible tokens having their logits erroneously amplified but also greatly improves the time efficiency of atbias. tokens related to key entities cannot be precisely located due to the tokenization rules. therefore, we developed an n-gram and jaccard-based simi-larity comparison algorithm to match tokens with entities. we introduce bias to the logits of both new and old knowledge entities based on the similarity computed between the filtered tokens and these entities. the tokens related to new knowledge are more likely to be output than parametric knowledge during the generation of llms, thus significantly enhancing the editing capabilities of ice. experimental results indicate that our atbias significantly enhances ice performance, achieving up to a 32.3% improvement over state-of-the-art decoding methods while incurring only half the latency. this means that atbias not only further improves editing capabilities but can also be widely applied to llms with negligible cost. furthermore, we suggest that research into decoding methods should focus more on key tokens rather than the entire sequence in generation. ## preliminary llms decoding. the primary goal of llms during decoding is to predict the succeeding word within a provided context sequence. formally, given a sequence of tokens {x 1 , x 2 , . . . , x t-1 } of length t -1, we can calculate the probability distribution of next token over the vocabulary set v: multi-hop editing. multi-hop editing is a highly challenging task in ke, aimed at verifying whether a fact has been thoroughly edited in llms. it not only edits the specific knowledge but also all related knowledge within the multi-hop relations impacted by this edit. for example, consider the two-hop question in figure 2 . the original answer would be "united states" with the facts stephen king wrote "misery", stephen king is a u.s. citizen. with an edit richard dawkins wrote "misery" and existing knowledge richard dawkins is british, the edited output answer should be "united kingdom". ## methods the framework of atbias is shown in figure 2 . first, we induce llms to output parametric knowledge by clozing the retrieved new knowledge, and then we extract the knowledge entities from them (section 3.1). this process can be optimized through knowledge caching (section 3.5). next, we refine the tokens using a probability and rank-based token filter (section 3.2), and match key entities with an n-gram and jaccard similarity calculation algorithm (section 3.3). finally, we adaptively bias the logits of the crucial tokens (section 3.4) to predict the next tokens. ## related work factual hallucinations. factual hallucinations have garnered widespread attention due to their significant side effects, as llms generate content that deviates from established world knowledge @xcite @xcite @xcite . these hallucinations can arise from various sources and at different stages of the llm life cycle @xcite . outdated knowledge is a major factor contributing to factual hallucinations. atbias enhances ke during the inference stage in llms to mitigate these hallucinations. knowledge editing. ke @xcite has been proposed to update information in llms, enabling accurate responses to current questions. in general, there are three lines of works for ke. model editing @xcite @xcite involves adding or altering the model parameters responsible for the undesirable output. meta-learning methods (de cao et al., 2021; mitchell et al., 2021) use a hypernetwork to learn the necessary adjustments for editing llms. in-context editing methods (ice) (madaan et al., 2022; zhong et al., 2023; zheng et al., 2023; bi et al., 2024c) demonstrate significant potential, enabling the editing of llms by prompting them with edited facts and retrieving editing demonstrations from the edit memory. decoding strategy. recent work modifies various decoding strategies to enhance different alignments by altering the logits of the original tokens during generation. cd (li et al., 2023) compares powerful expert language models with weaker amateur language models to enhance fluency and coherence. dola (chuang et al., 2023) contrasts mature layers with premature layers, while icd (zhang et al., 2023a) compares with models injected with hallucinations, aiming to enhance the factual accuracy of the model. deck (bi et al., 2024a) enhances ice by highlighting the output probability increment of new knowledge in contrast to the parametric knowledge. unlike the aforementioned decoding methods, atbias proposed in this paper only needs to adjust key tokens to enhance ke and mitigate factual hallucinations in llms. ## conclusion in this work, we propose a new ke framework, atbias, to enhance ice. atbias focuses on the crucial tokens that are mostly related to knowledge during the generation, biasing their logits by matching the knowledge entities. this design effectively reduces the potential risk of introducing fundamental errors in the logical coherence of the entire inference statement. experimental results show that atbias significantly improves the editing success rate of ice and outperforms the current best decoding methods. furthermore, the latency of atbias is at most 1.01 times that of the base-line, meaning atbias not only enhances ice but can also be widely applied with negligible cost.