🤖 machine learning

Towards Reliable, Generalizable, and Specific In-Context Knowledge Editing via Multi-Objective Reinforcement Learning

This paper proposes Multi-Objective In-context Knowledge Editing (MO-IKE), a multi-objective reinforcement learning algorithm that formulates prompt construction as a Constrained Markov Decision Process to simultaneously optimize reliability, generality, and specificity, thereby achieving superior performance in updating LLM knowledge compared to prior methods.

Original authors: Xuzhong Wang, Maiqi Jiang, Tejal Nair, Girija Bhusal, Yanfu Zhang, Haipeng Chen

Published 2026-08-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Xuzhong Wang, Maiqi Jiang, Tejal Nair, Girija Bhusal, Yanfu Zhang, Haipeng Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Large language models are the engines behind a new generation of artificial intelligence, capable of writing stories, solving problems, and answering questions with a fluency that often mimics human thought. However, these systems have a fundamental limitation: their knowledge is frozen in time. Once a model finishes its initial training, the facts it holds become static. If a model was trained before the 2026 Winter Olympics, it will not know the final medal counts, no matter how many times you ask. Updating this knowledge usually requires a massive, expensive retraining process that is impractical for most users. To solve this, researchers have developed a technique called in-context knowledge editing. Instead of changing the model's internal code, this method provides the model with a set of examples in the conversation prompt, effectively teaching it a new fact for the duration of that specific interaction. The challenge lies in constructing the perfect set of examples: it must be strong enough to force the model to accept the new fact, flexible enough to work even if the question is phrased differently, and precise enough to ensure the model does not accidentally forget or alter unrelated facts it already knows.

A team of researchers has now introduced a new system called MO-IKE that significantly improves how these prompts are built. Previous attempts to automate this process relied on reinforcement learning, a method where a computer program learns by trial and error. While these earlier systems could successfully teach the model a new fact, they often did so at a high cost. In their rush to ensure the new fact was accepted, the systems would strip away the examples needed to protect the model's existing knowledge. This created a fragile balance where fixing one error caused others to appear. The new approach treats the construction of the prompt not as a simple list of items, but as a sequential decision-making process. The system learns to choose, one by one, whether to add an example that states the new fact, an example that rephrases the question, or an example that reinforces a fact that should remain unchanged. It stops adding examples only when it determines the prompt is complete.

The researchers trained this system using a multi-objective reward structure, which means the computer was rewarded for success in three areas simultaneously: getting the new fact right, handling different ways of asking the question, and keeping unrelated facts safe. By balancing these competing goals, the new method prevents the system from becoming too aggressive in its updates. When tested on a standard set of facts, the new system raised the success rate of editing from 87.1 percent to 91.1 percent. More importantly, it dramatically improved the preservation of unrelated knowledge, raising the retention rate from 41.0 percent to 63.4 percent. This means the model is now much less likely to "forget" facts it didn't need to change while learning the new ones. The system also proved to be highly adaptable; a version trained on one specific language model worked just as well when applied to a completely different model without any additional training.

The study demonstrates that the way examples are organized matters just as much as the examples themselves. Earlier methods often used a fixed template or focused only on a single type of example, which led to unbalanced prompts. The new system dynamically mixes different types of examples, creating a context that is robust and specific. In tests across five different frozen language models and four different datasets, including a massive benchmark with over 300,000 examples, the method consistently outperformed previous techniques. The researchers found that by carefully managing the trade-off between teaching new information and protecting old information, they could achieve a level of reliability and specificity that was previously out of reach. This work suggests that the future of updating artificial intelligence may not lie in massive retraining, but in smarter, more balanced ways of guiding the model through conversation.

Technical Summary: MO-IKE via Multi-Objective Reinforcement Learning

Problem Statement

Large Language Models (LLMs) rely on static parametric knowledge that becomes outdated after pretraining. While Knowledge Editing (KE) aims to update model behavior on specific facts without full retraining, In-Context Knowledge Editing (IKE) has emerged as a training-free alternative applicable to black-box models. IKE constructs a prompt by retrieving and concatenating demonstrations to induce the desired update.

However, constructing an effective IKE prompt involves balancing three often contradictory objectives:

Reliability: The model must successfully output the new fact for the target query.
Generality: The update must apply to semantically equivalent queries (paraphrases).
Specificity: The model must retain unrelated neighboring facts without alteration.

Recent Reinforcement Learning (RL) approaches, such as DR-IKE, attempt to dynamically select demonstrations to optimize reliability. However, the paper identifies two critical deficiencies in these state-of-the-art methods:

Neglect of Sequential Ordering: They often treat demonstration organization as a secondary concern or only optimize the ranking of specific types (e.g., RETAIN), ignoring the global inter-dependency of demonstration order.
Imbalance in Objective Optimization: By optimizing primarily for edit success (reliability), these methods inadvertently discard RETAIN demonstrations (which protect neighboring knowledge). This leads to a "reward hacking" phenomenon where the retriever filters out specificity-protecting examples, causing a collapse in Retention Rate (specificity).

Methodology: MO-IKE

The authors propose MO-IKE (Multi-Objective In-context Knowledge Editing), an algorithm that frames prompt construction as a Constrained Markov Decision Process (Constrained MDP) to jointly optimize reliability, generality, and specificity.

1. Formulation as a Constrained MDP

The retrieval process is modeled as a sequential decision-making problem $(S, A, T, R, C)$ :

State ( $S$ ): The currently constructed prompt context.
Action ( $A$ ): Unlike DR-IKE, which restricts actions to ranking RETAIN candidates, MO-IKE expands the action space to include selecting from COPY (direct restatement), UPDATE (paraphrased query), and RETAIN (unrelated fact) demonstrations, plus a STOP action. This allows the agent to learn the optimal global ordering and proportion of demonstration types.
Reward ( $R$ ): The primary reward is Edit Success (ES), indicating whether the model outputs the new fact.
Cost Constraints ( $C$ ): To enforce specificity and generality, the framework introduces cost constraints for Paraphrase Consistency (PC) and Retention Rate (RR). These are defined as the logical inverse of task success (i.e., a cost is incurred if the model fails to maintain the original fact or fails to generalize to a paraphrase).

2. Multi-Objective Reward Shaping

To solve the constrained optimization problem, MO-IKE employs Lagrangian relaxation to scalarize the multi-objective problem. The final reward signal $r(s, a)$ is computed as:
$r(s, a) = R(s, a) - \sum_{k \in \{PC, RR\}} \lambda_k C_k(s, a)$
Here, $\lambda_k$ are fixed hyperparameters (not learned dual variables) that penalize degradation in generality and specificity. This structure prevents the retriever from sacrificing specificity for reliability.

3. Policy Optimization

The retriever, implemented as a BERT-based encoder with trainable projection matrices, is trained using Group Relative Policy Optimization (GRPO).

Sampling: For each edit instance, the policy samples a group of $G$ prompt trajectories.
Evaluation: The frozen target LLM evaluates these trajectories to compute composite rewards.
Update: Advantages are estimated relative to the group mean, and the policy is updated to maximize the clipped surrogate objective while penalizing divergence from a reference policy via KL divergence.

Key Contributions

Identification of Deficiencies: The paper highlights that existing RL-based IKE methods (specifically DR-IKE) fail to balance the trade-off between edit success and knowledge retention, and they overlook the importance of the sequential ordering of diverse demonstration types.
MO-IKE Algorithm: The authors introduce a multi-objective RL algorithm that treats prompt construction as a Constrained MDP. By expanding the action space to include all demonstration categories and applying cost constraints via reward shaping, MO-IKE ensures a stable policy that balances reliability, generality, and specificity.
Architectural Improvements: The method introduces a dynamic retriever capable of learning the optimal "global" ordering of COPY, UPDATE, and RETAIN demonstrations, rather than treating them as static templates or optimizing only a subset.

Experimental Results

The authors evaluated MO-IKE on Llama-3.2-3B, Mistral-7B, and other models across COUNTERFACT, ZSRE, WIKI-COUNTERFACT, and the large-scale UniEdit benchmark (311K examples).

Performance Gains: On Llama-3.2-3B, MO-IKE improved Edit Success (Reliability) from 85.0% to 92.0% and Retention Rate (Specificity) from 34.7% to 57.7% against the DR-IKE baseline. The overall harmonic mean score increased from 54.9 to 73.5.
Specificity Preservation: Crucially, MO-IKE reversed the specificity degradation seen in DR-IKE. For Mistral-7B, it improved the Retention Rate by +22.7% over the baseline, shifting the Retention Metric from a negative -22.5 to a positive +32.1.
Zero-Shot Generalization: A retriever trained on Llama-3.2 generalized effectively to Mistral-7B in a zero-shot setting, maintaining performance comparable to a natively trained retriever.
Ablation Studies:
- Removing the retention constraint led to a significant drop in specificity, confirming the necessity of the multi-objective formulation.
- Expanding the action space to include all demonstration types (vs. only RETAIN) was shown to be critical for achieving balanced performance across all metrics.

Significance and Claims

The paper claims that MO-IKE provides a principled solution to the competing objectives of in-context knowledge editing. By framing the problem as a Constrained MDP and utilizing multi-objective reward shaping, the method achieves a "stable policy" that avoids the "reward hacking" of previous RL approaches.

The authors emphasize that their approach is training-free for the target LLM (the LLM remains frozen) and is applicable to black-box models. They note that while MO-IKE significantly improves upon existing retrieval strategies, the absolute retention rate remains a bottleneck inherent to in-context learning methods (hovering around 60%), suggesting that gradient-free approaches still have room for improvement. The work establishes that balancing the proportion and ordering of distinct demonstration types is essential for reliable, generalizable, and specific knowledge editing.