In-Context Editing: Learning Knowledge from Self-Induced Distributions

TL;DR

Consistent In-Context Editing (ICE) tackles efficient knowledge updates in large language models by steering learning toward a contextual distribution induced by targeted knowledge, rather than optimizing to a one-hot target. The method combines a standard fine-tuning objective with a contextual consistency loss, enforcing $p_ heta(oldsymbol{x}|ig[oldsymbol{c},oldsymbol{q}ig]) \approx p_ heta(oldsymbol{x}|oldsymbol{q})$ while sampling context-conditioned sequences to guide updates. Empirical results on KnowEdit datasets show ICE achieves high edit accuracy and favorable locality and portability, with strong linguistic quality and robust continual editing capabilities, often outperforming baselines like ROME, MEMIT, FT-L, and FT-M. The work demonstrates that learning toward a contextual distribution yields more stable, generalizable updates, preserving the model’s existing knowledge while incorporating new information, albeit with higher computational costs due to in-context sampling and optional external-context generation. Overall, ICE offers a practical, gradient-based framework for continual knowledge editing with promising implications for personalized and updatable LLM deployments.

Abstract

In scenarios where language models must incorporate new information efficiently without extensive retraining, traditional fine-tuning methods are prone to overfitting, degraded generalization, and unnatural language generation. To address these limitations, we introduce Consistent In-Context Editing (ICE), a novel approach leveraging the model's in-context learning capability to optimize toward a contextual distribution rather than a one-hot target. ICE introduces a simple yet effective optimization framework for the model to internalize new knowledge by aligning its output distributions with and without additional context. This method enhances the robustness and effectiveness of gradient-based tuning methods, preventing overfitting and preserving the model's integrity. We analyze ICE across four critical aspects of knowledge editing: accuracy, locality, generalization, and linguistic quality, demonstrating its advantages. Experimental results confirm the effectiveness of ICE and demonstrate its potential for continual editing, ensuring that the integrity of the model is preserved while updating information.

Figures (6)

• Figure 1: Overview. (a) In-Context Learning: Utilizes context prompts without modifying model parameters, allowing dynamic adaptation but lacking parameter updates. (b) Traditional Fine-Tuning: Minimizes the distance between predictions and a one-hot target ($\delta_{x^*}$) using cross-entropy loss ($L_{ft}$), often leading to overfitting. (c) Consistent In-Context Editing (ICE): Adds a contextual loss ($L_{ice}$) to the traditional fine-tuning loss ($L_{ft}$). $L_{ice}$ minimizes the divergence between model outputs with and without a context prompt, aligning the model toward internalizing new knowledge. This helps ICE achieve effective knowledge incorporation while preserving general model stability.
• Figure 2: Continual editing with Llama2-7b-chat on WikiData$_{recent}$. Each edit builds on the previous model, risking deterioration over time. The model is assessed immediately after each edit without re-evaluating previous edits, testing its ability to update continuously. While most methods deteriorate, sometimes performing worse than the unedited version, our method, ICE, maintains integrity and achieves promising performance.
• Figure 3: Comparison of ICE with static and dynamic targets on an example, where the query is "The name of the country which Academy Award for Best Picture is associated with is?" and the target is "Wassoulou Empire". The line plots on the left illustrate the loss over optimization steps for static (top) and dynamic (bottom) targets under temperatures from 0.1 to 100. The figures on the right show how the probabilities of the top-6 predicted tokens for $x_2$, the second token following the target, change with iteration steps. The tokens are arranged from left to right in descending order of probability without context. At early steps, the token "Wass" appears due to its presence as the initial token in the target $\vb{x}^*$. At later steps, the probability of "Wass" in dynamic targets (top) significantly declines, indicating successful adaptation and suppression of repetitive token predictions. In contrast, for static targets (bottom), the probability of "Wass" remains relatively high throughout the optimization steps.
• Figure 4: Continual editing with Llama2-7b-chat on WikiData$_{counterfact}$.
• Figure 5: Continual editing with Llama2-7b-chat on ZsRE.

Theorems & Definitions (3)

• Lemma 1
• proof
• proof