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

The rapid development of large language models (LLMs) like GPT, BERT, and their instruction-tuned variants has transformed NLP applications. Despite their success, these models rely on static parameters, which become outdated as new information emerges. Early methods such as MEND introduced gradient-based weight updates for targeted knowledge editing but faced high computational costs and limited applicability to black-box models. Recent advances shifted focus to prompt engineering and demonstration selection, exemplified by IKE and DR-IKE, which utilize different demonstration categories (COPY, UPDATE, RETAIN) to locally modify model behavior. However, these approaches often optimize a single objective, neglecting the complex trade-offs among reliability, generality, and specificity, leading to inconsistent performance and limited global coherence. Reinforcement learning was introduced to dynamically select demonstrations, but existing RL methods like DR-IKE primarily optimize for edit success, often sacrificing retention and generality. The challenge remains to develop a unified framework that systematically balances these conflicting goals while maintaining computational efficiency and robustness.

The core challenge in in-context knowledge editing is balancing multiple objectives: ensuring the edited fact is reliably updated (reliability), maintaining applicability across paraphrases and related queries (generality), and preventing unintended modifications to neighboring facts (specificity). Existing methods tend to optimize these objectives separately or focus on a single goal, resulting in trade-offs that compromise overall performance. Moreover, the demonstration selection process often ignores the global organization and sequential ordering of examples, which are critical for effective in-context learning. This leads to unstable updates, inconsistent results, and difficulty in scaling to diverse tasks and models. The problem is further compounded by the need for a unified framework that can adaptively balance these goals during inference, especially under limited context budgets.

The paper introduces MO-IKE, a multi-objective reinforcement learning approach that models prompt construction as a Constrained Markov Decision Process (Constrained MDP). Key innovations include:

• �� Formulating demonstration selection as a sequential decision-making process with explicit constraints on paraphrase consistency and retention.
• �� Designing a composite reward function that integrates edit success, paraphrase generality, and neighbor knowledge retention, balanced via fixed penalty coefficients.
• �� Employing Group Relative Policy Optimization (GRPO) for stable, multi-objective policy training.
• �� Enhancing the demonstration retrieval architecture to dynamically adjust the proportion of demonstration categories, ensuring global coherence and balanced optimization.
• �� Demonstrating superior performance over state-of-the-art baselines across multiple datasets and models, with robust zero-shot transfer capabilities.

These innovations collectively enable a principled, scalable, and effective multi-objective knowledge editing framework.

• �� Model the prompt construction as a Constrained MDP with states representing the current prompt and actions as selecting demonstration categories or stopping.
• �� Define the action space to include COPY, UPDATE, RETAIN demonstrations, and a stop token, allowing flexible ordering.
• �� Use a BERT-based retriever to sequentially select demonstrations, with probabilities computed via a softmax over dot products of learned embeddings.
• �� Construct a composite reward combining edit success, paraphrase consistency, and neighbor retention, with fixed penalty coefficients.
• �� During training, sample multiple prompt trajectories, evaluate them with the frozen LLM, and estimate advantages based on rewards.
• �� Optimize the retriever policy using GRPO, balancing exploration and exploitation, with regularization via KL divergence.
• �� In inference, select demonstrations greedily based on learned policy, constructing prompts that balance all objectives.
• �� The training process iteratively refines the demonstration selection policy, leading to globally coherent prompts that satisfy multiple constraints.

• �� Datasets include COUNTERFACT, ZSRE, and WIKI-COUNTERFACT, with a split into editable and demonstration pools.
• �� Models evaluated are Llama-3.2-3B, Mistral-7B, with frozen parameters.
• �� Baselines compared include FactPrompt, Edit-CoT, IKE, and DR-IKE.
• �� Metrics cover edit success, paraphrase consistency, retention rate, and overall harmonic score.
• �� Ablation studies test the impact of constraints and reward components, tuning penalty coefficients.
• �� Experiments demonstrate that MO-IKE outperforms baselines significantly, with improvements in reliability, generality, and retention across datasets and models.

• �� On Llama-3.2-3B, MO-IKE improves edit success from 85% to 92%, retention from 41% to 63.4%, and overall harmonic score from 61.7 to 75.7.
• �� On Mistral-7B, achieves over 80% edit success, with a 23% increase in retention, and maintains robustness in zero-shot transfer.
• �� Ablation results confirm that including retention constraints prevents knowledge degradation, and multi-objective reward shaping stabilizes training, leading to balanced performance.
• �� Across datasets, MO-IKE consistently surpasses prior methods, demonstrating its scalability and effectiveness.

• �� Suitable for dynamic knowledge maintenance in enterprise AI systems, enabling real-time updates.
• �� Applicable in knowledge bases, virtual assistants, and automated question-answering systems requiring accurate, up-to-date information.
• �� Future integration with auto-tuning mechanisms could further enhance adaptability and efficiency, broadening deployment in multi-task environments.

• �� Dependence on predefined demonstration categories and fixed penalty coefficients may limit flexibility across diverse tasks.
• �� High computational cost during training, especially for large datasets and multi-objective optimization.
• �� Manual tuning of demonstration proportions remains necessary; automatic adjustment mechanisms are needed for broader applicability.