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

Xuzhong Wang; Maiqi Jiang; Tejal Nair; Girija Bhusal; Yanfu Zhang; Haipeng Chen

Source-linked AI summary

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

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

arXiv:2608.25100v1Aug 25, 2026cs.AIcs.LG

TL;DR

Static parametric knowledge becomes outdated, motivating training-free in-context editing that must balance reliability, generality, and specificity. MO-IKE formulates prompt construction as a constrained sequential decision process with multi-objective reward shaping, and reports stronger performance across these objectives than prior retrieval strategies.

Problem
Static parametric knowledge becomes outdated, while prior in-context editing methods overlook the balance of reliability, generality, specificity, and global demonstration organization.
Method
MO-IKE models prompt construction as a Constrained MDP and trains a dynamic retriever with multi-objective reward shaping optimized by GRPO.
Results
MO-IKE reaches 92.0% edit success on Llama-3.2 and improves retention rate by +23.0% over DR-IKE while maintaining or exceeding peak reliability.
Takeaways & Limitations
MO-IKE consistently outperforms existing retrieval strategies across reliability, generality, specificity, and cross-model evaluation.
Takeaways & Limitations
MO-IKE is trained and evaluated only on the first 2,000 COUNTERFACT records, and its fixed constraint coefficients remain insufficiently explored.

Abstract

from arXiv · show

Large Language Models (LLMs) are powerful but limited by static parametric knowledge that becomes outdated once pretraining ends. Knowledge editing addresses this problem by updating model behavior on target facts without full retraining. In particular, in-context knowledge editing has gained attention because it is training-free and readily applicable to black-box LLMs. Recent reinforcement learning (RL)-based approaches improve over fixed retrieval strategies by adapting prompt construction to the quantity-quality trade-off. Despite initial success, they fail to model the prompt as a structured entity under the distinct and often competing objectives of reliability, generality, and specificity. Previous methods largely optimize a single objective and make decisions over only part of the prompt construction process, thereby overlooking both the balance of different objectives and the global organization of demonstrations. We propose Multi-Objective In-context Knowledge Editing (MO-IKE), a multi-objective RL algorithm that formulates prompt construction for in-context knowledge editing as a Constrained Markov Decision Process. MO-IKE trains a dynamic retriever to optimize competing objectives in knowledge editing, enabling more balanced and globally coherent prompt construction. On Llama-3.2, MO-IKE improves edit success (reliability) from 85.0% to 92.0%, paraphrase consistency (generality) from 77% to 79%, while increasing retention rate (specificity) by 23.0% compared to prior RL-based methods.

1 Introduction

In-context knowledge editing addresses stale parametric knowledge by inducing targeted updates through prompts, but prompt construction must balance reliability, generality, and specificity. MO-IKE treats this construction as multi-objective sequential optimization across demonstration categories.

Static parametric knowledge becomes outdated after pretraining, limiting LLMs on real-time or dynamic information.
In-context knowledge editing keeps the model frozen and induces target-fact updates through prompt context, making it training-free and applicable to black-box LLMs.
Prompt construction must support reliability, generality, and specificity, which respectively concern editing the target fact, handling equivalent queries, and preserving unrelated facts.
DR-IKE adapts the number of retrieved examples but only ranks RETAIN candidates and does not organize the full prompt sequence.
Reliability and generality require COPY and UPDATE demonstrations, whereas specificity requires RETAIN demonstrations, creating competing prompt objectives.
MO-IKE formulates prompt construction as a Constrained MDP and uses multi-objective RL to jointly optimize demonstration selection across reliability, generality, and specificity.

2 Related Work

Prior work spans gradient-based editing, in-context alternatives, retrieval-based prompt construction, and RL-based selection. The literature leaves balancing reliability, generality, and specificity as an open concern.

Gradient-based editors update model weights but incur substantial computational overhead, motivating gradient-free in-context knowledge editing.
Existing in-context editing approaches often improve edit reliability without adequately balancing reliability, generality, and specificity.
RAG supplements parametric knowledge with external documents, whereas IKE overwrites targeted parametric facts while preserving unrelated knowledge.
Demonstration selection is difficult because prompt performance depends on example interactions, ordering, and model biases within a combinatorial search space.

3 Problem Statement

The paper defines knowledge editing as changing a targeted fact in a frozen model while preserving behavior on equivalent and unrelated queries. IKE achieves this through retrieved demonstrations appended to a base prompt.

Knowledge editing maps a frozen model and original factual triple to an edited model that reflects the revised fact.
Reliability requires the edited model to answer the exact target query with the revised fact.
Generality requires the revised fact to hold for semantically equivalent or logically dependent queries.
Specificity requires unchanged outputs for queries depending on unrelated knowledge.
IKE concatenates retrieved natural-language demonstrations to a base prompt while keeping the language model parameters frozen.
Retrieved demonstrations are categorized as COPY, UPDATE, or RETAIN, representing direct restatements, paraphrased updates, and neighboring facts to preserve.

4 Methodology

MO-IKE constructs prompts through sequential selection over COPY, UPDATE, and RETAIN demonstrations, evaluates reliability under generality and specificity constraints, and optimizes a shaped reward with GRPO.

Constrained MDP: MO-IKE expands the action space beyond RETAIN to all COPY, UPDATE, and RETAIN demonstrations, enabling global ordering across categories.
Constrained MDP: At each step, the retriever selects an unused demonstration or a stop token, deterministically appending selections until the prompt is finalized.
Constrained MDP: Edit Success is the primary reward representing reliability, while Paraphrase Consistency and Retention Rate constrain generality and specificity.
Multi-Objective Reward Shaping: The composite reward uses fixed-multiplier Lagrangian relaxation to penalize degradation in paraphrase consistency and retention rate.
MO-IKE Pipeline: The pipeline retrieves candidates, sequentially selects demonstrations, evaluates prompts with the target LLM, and optimizes the retriever policy.
MO-IKE Pipeline: During training, MO-IKE samples prompt trajectories, computes composite rewards and group-relative advantages, then updates the policy with GRPO.

5 Experiments

MO-IKE is evaluated across multiple knowledge-editing benchmarks, models, and objectives, with ablations examining its constrained reward and demonstration-selection design. It generally improves reliability and specificity while supporting cross-dataset and cross-model performance.

Evaluation Metrics: The evaluation measures reliability, generality, and specificity through ES and EM, PC and PM, and RR and RM, respectively.Overall Score is the harmonic mean of ES, PC, and RR.
Main Results: 92.0 ES is achieved by MO-IKE on Llama-3.2, increasing from DR-IKE’s 85.0, while RR improves by +23.0 over DR-IKE.For Mistral-v0.3, RR improves by +22.7 over DR-IKE.
Main Results: 42.0 on ZSRE and 59.0 on Wiki are MO-IKE’s highest overall Scores, alongside the strongest RR on both datasets.The reported cross-dataset results indicate performance beyond COUNTERFACT.
Cross-Model Generalization: A retriever trained on Llama-3.2 transfers zero-shot to Mistral-7B without performance drop and performs practically identically to a natively trained version.The transferred retriever outperforms the other baselines in the reported comparison.
Ablation Studies: Removing the retention penalty lowers RR, while the full multi-objective reward achieves better Score than using only a retention constraint.The authors attribute the latter difference to more informative group-relative advantage signals under multiple constraints.
Demonstration Structure: Optimizing the full COPY, UPDATE, and RETAIN sequence outperforms optimization restricted to RETAIN demonstrations across all three metrics.MO-IKE expands the action space to all demonstration categories and captures their global interactions.
Demonstration Structure: MO-IKE’s soft stopping mechanism preserves contextual heterogeneity, avoiding the aggressive removal of RETAIN candidates associated with DR-IKE’s greedy optimization.The analysis links this structural balance to improved editing specificity.

Conclusion

MO-IKE formulates prompt construction as sequential, constrained multi-objective reinforcement learning for dynamic demonstration retrieval. Across multiple models, it outperforms existing retrieval strategies on reliability, generality, specificity, and cross-model generalization.

Conclusion: MO-IKE trains a dynamic retriever through a Constrained MDP to optimize reliability, generality, and specificity jointly.Prompt construction is treated as sequential decision-making rather than fixed retrieval.
Conclusion: Across multiple LLMs, MO-IKE consistently outperforms existing retrieval strategies across the three knowledge-editing objectives.The conclusion also reports cross-model generalization and ablations supporting the constraint, embedding, and ranking designs.

Limitations

The paper identifies limitations in dataset scale, reward-constraint hyperparameter analysis, retention performance, and evaluation scope.

Limitations: MO-IKE is trained and evaluated on only the first 2,000 COUNTERFACT records, limiting assessment of overall capabilities.The authors state that this dataset scale may be insufficient for evaluating the model’s broader capabilities.
Limitations: Fixed constraint coefficients leave the precise impact of reward hyperparameters on RL optimization unexplored.The limitation arises from incorporating all constraints directly into the reward function with static coefficients such as lambda.
Limitations: Retention rate remains around 60% despite overall improvement, reflecting an inherent in-context-learning bottleneck.The authors identify room for improvement in this gradient-free method.
Limitations: Evaluation is limited to one knowledge-editing task, leaving generalization to other in-context learning applications unclear.The paper specifically mentions model reasoning as an untested application.

A Training Procedure for Multi-objective GRPO

MO-IKE trains a retriever to construct prompts sequentially under reliability, generality, and specificity constraints. It uses a constrained-reward formulation and GRPO to optimize groups of candidate prompts.

GRPO training: Training generates groups of candidate trajectories, computes rewards and group-relative advantages, then applies clipped policy optimization with a KL penalty.The algorithm reconstructs states from sampled demonstration sequences before updating retriever parameters by gradient descent.
Sequential construction: The retriever sequentially samples demonstrations without replacement, updates the prompt state, and can terminate through a learnable stop action.The action space combines available demonstrations with a stop token, making prompt length and organization jointly selectable.
Problem formulation: MO-IKE models prompt construction as a Constrained MDP whose state is the current prompt and whose actions select demonstrations or stop construction.The prompt is evaluated for edit reliability while paraphrase consistency and retention rate act as constraints.
Reward optimization: MO-IKE scalarizes editing success and constraint costs through a fixed-multiplier Lagrangian composite reward optimized with GRPO.The fixed coefficients λPC and λRR apply static penalty weightings to balance editing performance against paraphrase-consistency and retention constraints.
Additional baseline: RAG-RL achieves moderate edit success but collapses on retention, reaching 8.0% RR on CounterFact with Qwen2.5-7B-Instruct.This evaluation tests an RL retriever trained for open-domain QA when transferred directly to in-context knowledge editing.

E Training Dynamics

MO-IKE’s retrieval policy converges rapidly during training. After a sharp first-epoch improvement, later epochs mainly refine ES and PC while RR remains stable.

Training dynamics: By the first epoch, MO-IKE establishes a strong, balanced demonstration ordering, marked by an initial peak in ES and RR.The training-dynamics analysis attributes this pattern to rapid policy convergence.
Training dynamics: During epochs 2 and 3, performance largely plateaus while ES and PC receive stable refinements and RR does not degrade.The plateau indicates limited additional change after the initial policy improvement.

F Hyperparameter Sensitivity.

MO-IKE remains comparatively stable across a 100× sweep of constraint weights. The observed directional changes are modest and consistent with the intended trade-offs between PC and RR.

Sensitivity results: Across 100× sweeps of λPC and λRR, ESR varies by ≤2.0 points, PC by ≤3.3, and RR remains unchanged.MO-IKE outperforms all baselines at every tested setting without monotonic degradation in any metric.
Sensitivity results: Upweighting PC improves PC by +1.0, whereas upweighting RR degrades PC by -2.3 without further improving RR.The results place the default (1, 1) weighting in a flat region balancing both objectives.

G Additional Experiments

Additional experiments test MO-IKE across larger datasets, model scales, heuristic baselines, newer LLMs, and case-study prompt structures. The results support balanced performance and retention beyond the primary setting.

Larger dataset: On UNIEDIT’s 311K examples, MO-IKE achieves the best Score and best RR.UNIEDIT is roughly 15× larger than COUNTERFACT and mixes relation types and knowledge domains.
Model scaling: Score improves as model size increases, with the highest overall rates achieved by models with at least 7B parameters.This experiment evaluates MO-IKE across multiple frozen LLMs.
Natural baselines: Simple duplication or shuffling produces moderate performance but does not match MO-IKE’s balance of ES, PC, and RR.The comparison tests whether prompt selection and ordering alone can reproduce the multi-objective RL framework’s results.
2025-era LLMs: On Qwen3-4B-Instruct, MO-IKE exceeds DR-IKE by approximately 23% on RR and achieves the highest overall score of 70.2%.MO-IKE preserves the newer model’s high ESR and PC while differentiating methods through retention.
Case study: The case study shows that IKE fails to execute the target edit, whereas DR-IKE executes it but discards RETAIN candidates and alters a related fact.The comparison links prompt composition with the reliability–specificity trade-off.