Dynamic Retriever for In-Context Knowledge Editing via Policy Optimization

TL;DR

DR-IKE tackles the problem of stale LLM knowledge by introducing a dynamic, reward-guided retrieval framework that selects and prunes in-context demonstrations for knowledge editing without modifying model weights. A light-weight BERT retriever is trained via policy gradients to surface only the most informative Retain demonstrations, while a learnable budget controller adaptively shortens prompts for easy edits and expands support for difficult ones. Across CounterFact and multiple LLMs, DR-IKE achieves up to 17.1% higher edit success with up to 41.6% latency reductions and better preservation of unrelated knowledge, demonstrating strong adaptivity to task difficulty and efficiency in black-box settings. The approach offers a scalable, model-agnostic alternative to gradient-based editors, with potential applicability to temporal, numerical, and paraphrase-rich updates in real-world deployments.

Abstract

Large language models (LLMs) excel at factual recall yet still propagate stale or incorrect knowledge. In-context knowledge editing offers a gradient-free remedy suitable for black-box APIs, but current editors rely on static demonstration sets chosen by surface-level similarity, leading to two persistent obstacles: (i) a quantity-quality trade-off, and (ii) lack of adaptivity to task difficulty. We address these issues by dynamically selecting supporting demonstrations according to their utility for the edit. We propose Dynamic Retriever for In-Context Knowledge Editing (DR-IKE), a lightweight framework that (1) trains a BERT retriever with REINFORCE to rank demonstrations by editing reward, and (2) employs a learnable threshold to prune low-value examples, shortening the prompt when the edit is easy and expanding it when the task is hard. DR-IKE performs editing without modifying model weights, relying solely on forward passes for compatibility with black-box LLMs. On the COUNTERFACT benchmark, it improves edit success by up to 17.1%, reduces latency by 41.6%, and preserves accuracy on unrelated queries, demonstrating scalable and adaptive knowledge editing. The code is available at https://github.com/mwnafee/DR-IKE .

Figures (4)

• Figure 1: Motivating example for adaptive in‑context editing. (a) A fixed‑shot editor appends the same large set of demonstrations to every query; the oversized prompt diffuses model attention, hindering both a straightforward location swap (Winter Paralympics: Vancouver $\rightarrow$ Berlin) and a more concept‑heavy update (Hannah Point: Antarctica $\rightarrow$ Europe), whose corpus evidence is dominated by Antarctic descriptions. (b) Our dynamic editor adjusts prompt length to task complexity, retaining a minimal context for the simple edit while retrieving additional, targeted facts for the harder one.
• Figure 2: Overall architecture of the proposed framework. Training stage (left). Given an edit query, the retriever scores a pool of retained candidate demonstrations. A learnable threshold $\sigma$ converts these scores into an integer $k$, the number of prompt variants to generate. The frozen LLM processes each variant; its outputs yield the reinforcement signal used to update the retriever via policy gradients. Whenever the aggregated outcome flips between success and failure, $\sigma$ is adjusted to capture the revised difficulty estimate. Inference stage (right). At test time, the learned $\sigma$ truncates the ranked list once, yielding a single concise prompt to the frozen LLM.
• Figure 3: Case study illustrating prompt construction and example selection for challenging queries under IKE versus our method.
• Figure 4: Retain budget distribution for LLaMA 3.1-8B under DR-IKE. Each bar shows the frequency of a given Retain count per prompt. Mean = $\mathbf{3.02}$, Std = $\mathbf{0.96}$.

Theorems & Definitions (3)

• Definition 3.1: Knowledge Editing
• Definition 3.2: In‑context Knowledge Editing
• Definition 3.3: Demonstration Categories zheng2023ike