PaperlayerView exact revision on arXiv

Source-linked AI summary

You Only Edit Once: Incentivizing In-Context Capability of LLMs via Local Demonstration Refinement

Jiarong Wen, Qi Wang, Yun Qu, Yixiu Mao, Heming Zou, Haoang Chi, Lizhou Cai, Yiqin Lv, Kaiyu Zhang, Yuhang Jiang, Xiangyang Ji

arXiv:2609.33609v2cs.LG

TL;DR

Because ICL depends on selecting useful demonstrations, exhaustive subset searches and repeated target-LLM queries can make selection costly. The paper trains Jev-LDE to make one structured edit to an initially retrieved context, reporting improved classification results and transfer without retraining. The authors note that consequential use requires task-specific validation.

  • Problem

    Selecting useful demonstrations is challenging because irrelevant or repetitive examples can hurt ICL, while evaluating combinatorial subsets may require repeated target-LLM queries.

    Source: §1 · pp. 1–2

  • Method

    The paper trains Jev-LDE with reinforcement learning to make one structured membership edit to an initially retrieved demonstration set before target-LLM inference.

    Source: §1 · p. 2

  • Results

    Average one-shot accuracy rises from 81.2% to 88.1%, and Jev-LDE improves Semantic TopK in 46 of 48 classification configurations while transferring to held-out benchmarks and models without retraining.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The authors conclude that learned local editing is a practical method for enhancing few-shot ICL with a lightweight, plug-and-play module.

    Source: §1 · p. 2, §6 · p. 10

    Jev-LDE may inherit biases and errors from target models and demonstration pools, and consequential applications require task-specific validation.

    Source: Ethics Statement · p. 10

Abstract

from arXiv · show

In-context learning (ICL) is crucial for boosting the inference performance of large language models (LLMs). However, the effectiveness of ICL in LLMs is greatly influenced by the choice of demonstration sets. Exhaustive searches over these sets are combinatorial, and existing selectors often rely on relevance or likelihood proxies to implicitly assess ICL quality. Making repeated queries to the target LLM with these strategies can incur substantial costs. This work simplifies selection by framing it as a constrained local search problem and presents local demonstration editing (LDE). Starting with an initially retrieved set of demonstrations, LDE employs a single structured edit to explore its surrounding neighborhood while balancing performance gains with search costs. Technically, LDE is reduced to a policy search problem, for which we train a small LLM, referred to as Jev-LDE. This model as the System-1 modifies the retrieved demonstration set by performing actions such as \texttt{Keep}, \texttt{Delete}, or \texttt{Replace} elements, all within a framework of reinforcement learning with verifiable rewards. At test time, Jev-LDE executes a single edit of the retrieved demonstration set, followed by one inference from the target LLM, avoiding the need for iterative context scoring or subset searches. Across standard classification benchmarks, various target LLMs with Jev-LDE as the plug-and-play module consistently improve ICL performance, and Jev-LDE shows transferability to held-out benchmarks and models without retraining. These findings indicate that the LDE approach offers an efficient and adaptable method for harnessing the ICL capabilities of target LLMs.

1 Introduction

In-context learning depends on the quality of its demonstrations, but finding an effective subset is costly and longer contexts do not always improve accuracy. The paper proposes one-edit refinement of retrieved contexts and reports improved classification accuracy and transfer to held-out benchmarks and target models.

Motivation

  • Demonstration selection is difficult because irrelevant or repetitive examples can hurt predictions, while exhaustive subset evaluation is combinatorial and can require repeated target-LLM queries.Context-window memory and computational constraints also limit how many demonstrations fit in a prompt.Source: §1 · pp. 1–2

Approach

  • Local demonstration editing refines an initially retrieved context with one structured membership edit, trading off performance gains against search efficiency.The approach is designed to work with existing preselection methods such as BM25.Source: §1 · p. 2
  • A small LLM, Jev-LDE, is fine-tuned with reinforcement learning as a plug-and-play System-1 editor that uses Keep, Delete, or Replace actions.At test time, it edits the retrieved set before the target LLM performs in-context inference.Source: §1 · p. 2

Experiments

  • Jev-LDE also improves few-shot performance on held-out benchmarks and diverse target models without retraining.The paper frames this as effective generalization across configurations.Source: §1 · pp. 2–3
  • Average one-shot classification accuracy rises from 81.2% to 88.1% with Jev-LDE, which improves Semantic TopK in 46 of 48 configurations.The reported accuracy is averaged across TREC, DBPedia, and Banking77 using four target LLMs.Source: Figure 1 · p. 1, §1 · pp. 2–3

2 From Combinatorial Shot Selection to Naive Local Search

The section formulates k-shot demonstration selection as an intractable combinatorial objective approximated by a selector, then examines whether one local membership change can improve the retrieved set. On TREC, oracle and learned edits show that substantial gains can lie within this one-edit neighborhood.

2.1 Demonstration Selection Setup

  • The ideal k-shot selector chooses a size-k subset maximizing the target LLM’s likelihood of the gold answer, while a practical selector supplies an approximation.The feasible set consists of size-k subsets of the available demonstrations.Source: §2.1 · p. 3
  • The selection objective is intractable at test time, and evaluating the combinatorial set of candidates requires target-model queries despite unknown test labels.Traditional selectors therefore use proxy scores such as embedding similarity.Source: §2.1 · p. 3

2.2 One Local Edit Moves Accuracy

  • A one-edit neighborhood includes keeping the retrieved set, deleting one demonstration, or replacing one with a candidate outside the set.The replacement-distance constraint limits the search to sets at distance at most one from the initialization.Source: §2.2 · p. 3
  • On 4-shot TREC, oracle edits gain 16.8 and 26.0 percentage points over Semantic TopK for Qwen2.5-7B and Llama-3-8B, while Jev-LDE gains 7.4 and 8.2 points without test labels.Uniformly sampled valid actions hardly improve average performance; the figure shows the accuracy-change distribution under those actions.Source: §2.2 · p. 3, Figure 2 · p. 3

3 Learn to Edit Locally with Reinforcement Learning

The section formulates local demonstration editing as a query-specific search over a retrieved set and its one-edit neighborhood, then trains Jev-LDE with target-model feedback. Its analysis gives conditions for nearby successful edits and a finite-class feedback guarantee for learning an editing policy.

3.1 Local Editing Formulation

  • LDE restricts query-specific search to the retrieved initialization and its one-edit neighborhood, with Keep, Delete, and Replace actions defining the feasible edits.Replace swaps a selected demonstration for a candidate outside the retrieved set.Source: §3.1 · p. 4

3.2 Training with Target-Model Feedback

  • Jev-LDE learns a stochastic, query-aware editing policy with reinforcement learning, using frozen-target prediction correctness as verifiable reward and then editing once before target inference.The editor receives the task instruction, query, retrieved set, and replacement candidates; the target model remains frozen during policy optimization.Source: §3.1 · p. 4, §3.2 · pp. 4–5

3.3 Implementation and Analysis

  • The implementation trains Jev-LDE on static queries from source benchmarks with a fixed target LLM, then serves target LLMs across benchmarks without retraining.The pre-selector ranks the benchmark-level pool; its top-N demonstrations form the candidate pool and top-k initialize the edit.Source: §3.3 · p. 5
  • Under odd-k majority voting and the stated conditional-accuracy assumption, the optimal replacement distance has a geometric tail, with h=1 yielding single-replacement coverage of at least 2α−1.The result concerns random evaluation states and context-invariant binary query-label judgments.Source: §3.3 · pp. 5–6
  • Under fixed policy-class size, accuracy tolerance, and confidence, the sufficient feedback bound scales linearly with available actions, and the selected policy is within ε of the best policy in expected correctness reward.The guarantee is for the finite-class procedure in Proposition 6.Source: §3.3 · p. 6

4 Experiments

The experiments test Jev-LDE across classification benchmarks, shot budgets, target models, and selectors. They report broad gains from one edit, while also identifying less consistent reasoning-task results and modest inference-time overhead.

4 Experiments

  • The evaluation covers question, topic, and intent classification across training and held-out benchmarks, multiple target LLMs, and one- to four-shot budgets.Methods share data splits, prompts, and answer normalization; Semantic TopK initializes Jev-LDE’s shot set.Source: §4 · p. 7

4.1 Main Results

  • On held-out TACRED and TREC-50, DeepSeek-V4-Flash gains appear at every tested shot budget, ranging from 0.3–2.2 and 2.4–4.4 points, respectively.Across other settings, Jev-LDE improves 35 of 36 configurations, including 14 of 16 on held-out Banking77 and all displayed SST-2 settings.Source: Figure 5 · p. 7, §4.1 · pp. 7–8
  • Jev-LDE improves Semantic TopK in 46 of 48 classification configurations and is best or tied-best in 44; average one-shot accuracy rises from 81.2% to 88.1%.Gains remain positive on average at larger shot budgets.Source: §4.1 · p. 7, Table 1 · p. 8
  • Reasoning-task gains are less consistent: GSM8K differences stay within one point, while AQuA-RAT accuracy falls by 0.8 and 1.9 points for two target models.The AQuA-RAT declines are for Qwen2.5-7B-Instruct and Meta-Llama-3-8B-Instruct, respectively.Source: §4.1 · p. 8, Figure 6 · p. 9

4.2 Ablation Studies

  • Jev-LDE improves BM25, CASE, and TopK+ConE initializations in all 24 tested combinations; on one-shot TREC, doubling the replacement pool lowers accuracy from 80.6% to 79.4%.The pool-size comparison is for the default N=16 pool versus twice that size.Source: §4.2 · pp. 8–9, Figure 6 · p. 9, Figure 7 · p. 9
  • On one-shot TREC, Jev-LDE reaches 65.6% accuracy, beating Keep by 12.8 points and a prompted Meta-Llama-3-8B-Instruct editor by 7.8 points.Random editing and the untrained Qwen3-1.7B initialization score 23.2% and 31.2%, below Keep at 52.8%.Source: Figure 7 · p. 9, §4.2 · p. 9
  • On one-shot TREC, Jev-LDE raises accuracy from 65.4% to 80.6% with an 11.0% increase in end-to-end evaluation time over Semantic TopK.Runtime remains close to LMS3 and CASE and substantially below TopK+ConE.Source: Figure 8 · p. 9, §4.2 · p. 9

5 Related Work

Demonstration selection has moved from query-relevance ranking toward signals that estimate model-specific utility. LDE instead refines an existing selector’s output with one learned local edit.

  • Retrieval methods rank examples by query relevance, using embedding similarity in Semantic TopK or lexical matching in BM25; example order can also affect predictions.Semantic TopK; BM25; ordering (Lu et al., 2022).Source: §5 · p. 9
  • Later selectors use likelihood, entropy, model-specific retrieval, and stability signals, but these remain proxies for correctness of the complete downstream prompt.The cited approaches include Rubin et al. (2022), Iter et al. (2023), Peng et al. (2024), Wang et al. (2024), and Liu et al. (2025).Source: §5 · p. 10
  • LDE applies correctness feedback to one Keep, Delete, or Replace action on a retrieved set, unlike global subset search or iterative instruction-edit search.Its learned policy can commit to an edit without requiring a unique oracle action.Source: §5 · p. 10

6 Conclusion

The conclusion presents LDE as a single-edit alternative to pursuing state-of-the-art selection. It frames Jev-LDE as a lightweight module and reports gains and transfer without retraining.

  • The work asks whether a System-1 module can improve ICL by editing a retrieved demonstration set once, rather than pursuing a state-of-the-art selector.Source: §6 · p. 10
  • LDE frames this edit as constrained local search, with Jev-LDE trained through reinforcement learning as a lightweight, plug-and-play module.Source: §6 · p. 10
  • Classification experiments report improved ICL performance and generalization to new benchmarks and target models without retraining.The conclusion characterizes learned local editing as a practical method for enhancing few-shot ICL.Source: §6 · p. 10

Ethics Statement

The work uses public benchmarks and pretrained models, with no human participants or private personal data. It notes potential inherited biases and errors and calls for task-specific validation before consequential use.

  • The evaluation uses publicly available benchmarks and pretrained language models, without human participants or private personal data.Source: Ethics Statement · p. 10
  • Jev-LDE may inherit biases and errors from its target models and demonstration pools.Source: Ethics Statement · p. 10
  • The method is evaluated for research purposes; consequential applications require task-specific validation.Source: Ethics Statement · p. 10

Reproducibility Statement

The statement maps the paper’s methods, evaluation, and implementation material across sections and appendices. It also commits to releasing the inference code, post-trained Jev-LDE, and evaluated target LLMs in the final version.

  • Section 3 covers the LDE formulation, Jev-LDE’s action space and post-training, and the single-edit inference protocol.Source: Reproducibility Statement · p. 10
  • Section 4 describes datasets, baselines, and evaluation; Appendices A–C provide data processing, decoding, answer normalization, post-training configuration, and prompt templates.Source: Reproducibility Statement · p. 10
  • Appendices D and E provide complete shot-wise results for additional benchmark and pre-selector panels and for 8-shot and 16-shot comparisons.Source: Reproducibility Statement · p. 10
  • The authors state they will release the inference procedure code, post-trained Jev-LDE, and all evaluated target LLMs in the final version.Source: Reproducibility Statement · p. 11

A Dataset Construction and Evaluation Details

The evaluation standardizes key data and inference conditions across selectors while retaining method-specific candidate scopes. It reports broad paired gains for Jev-LDE and documents benchmark, ablation, and wall-time protocols.

  • All selectors share evaluation examples, demonstration pools, shot budgets, downstream prompts, decoding parameters, and answer normalization, while candidate scope remains method-specific.Source: A Dataset Construction and Evaluation… · p. 14
  • Demonstrations come from training splits and evaluation uses fixed test splits, except SST-2 uses validation because public test labels are unavailable.CASE validation examples come from the benchmark demonstration pool and are disjoint from the test split; Jev-LDE edits within a smaller semantic neighborhood than global selectors.Source: A Dataset Construction and Evaluation… · p. 14
  • The evaluation covers four target models at 1–4 shots, with deterministic decoding; DeepSeek-V4-Flash is additionally tested on TREC-50 and TACRED.TREC-50 uses its complete 5,452-example training split and 500-example test split; TACRED uses 10,000 stratified training-pool examples and 1,000 sampled test examples.Source: A Dataset Construction and Evaluation… · p. 15
  • The one-shot TREC policy ablation varies the action-selection policy while holding initialization, candidate neighborhood, editing state, action space, and evaluation protocol fixed.Compared policies include no edit, random action selection, un-post-trained Qwen3-1.7B, prompted Meta-Llama-3-8B-Instruct, and post-trained Jev-LDE.Source: A Dataset Construction and Evaluation… · p. 16
  • Pre-selector tests apply one edit within each selector’s own top-16 neighborhood, while candidate-pool tests vary replacement candidates across 7, 15, 23, and 31.The end-to-end wall-time comparison includes retrieval, candidate construction, editing, and target inference, measured over five runs.Source: A Dataset Construction and Evaluation… · pp. 16–17
  • Jev-LDE improves 46 of 48 configurations, all 24 pre-selector combinations, 14 of 16 held-out Banking77 settings, and 35 of 36 Qwen-target settings without retraining.Paired comparisons use the same evaluation examples; reported gaps include 15.2 points on one-shot TREC and 6.9 points on average one-shot accuracy.Source: A Dataset Construction and Evaluation… · p. 17

B Jev-LDE Post-Training and Implementation Details

Jev-LDE is post-trained with target-model correctness feedback on locally edited demonstrations, using filtered training states and a fixed target. The appendix reports training dynamics, prompt templates, and supplementary result tables.

B Jev-LDE Post-Training and Implementation Details

  • The editor is initialized from Qwen3-1.7B and trained with online correctness feedback from frozen Meta-Llama-3-8B-Instruct; multiple rollouts are used only during post-training.The appendix identifies GRPO as the group-relative optimization objective and keeps the target fixed.Source: B Jev-LDE Post-Training and… · pp. 17–18
  • Training states use AGNews, TREC, DBPedia, and GSM8K training-pool examples, with retrieved initial shots and remaining top-16 candidates forming local edit choices.Evaluation-split examples and labels are excluded; states from held-out transfer benchmarks are not included.Source: B Jev-LDE Post-Training and… · p. 17
  • States are retained only when probed local actions yield both binary reward outcomes, identifying consequential decisions without assigning optimal actions or labels.The authors note sparse-reward challenges but do not investigate policy-optimization issues.Source: B Jev-LDE Post-Training and… · pp. 17–18
  • Mean reward rises from 0.475 in the first 100 updates to 0.713 in the final 100, while actor entropy loss falls from 0.168 to 0.119.Evaluations use the final checkpoint at update 752; Figure 9 plots raw and smoothed reward and entropy-loss trajectories.Source: B Jev-LDE Post-Training and… · p. 18, Figure 9 · p. 18

C Prompt Templates and Examples

D Completed Results for Heatmap

  • Tables 5 and 6 provide underlying accuracy values for the generalization heatmap and TREC pre-selector ablation, respectively.Table 6 compares each pre-selector’s unedited result with its result after one Jev-LDE edit.Source: Table 5 · p. 22, Table 6 · p. 22

E Classification Results at Larger Shot Budgets

At larger shot budgets, Jev-LDE often improves Semantic TopK accuracy; the accompanying analyses bound local-edit learning under stated assumptions and identify limits on task and edit scope.

E Classification Results at Larger Shot Budgets

F.1 Proof of Proposition 4

  • The finite-shot analysis permits dependent judgments under conditional reliability, but its repair result is specific to fixed binary judgments and majority aggregation.Marginal accuracy alone is insufficient for the stated guarantee.Source: §F.1 · pp. 23–24

F.2 Proof of Proposition 6

  • The policy-selection analysis bounds the expected reward gap by ε, while local editing uses 53 actions versus 2,380 global sets when N=16 and k=4.The reward guarantee does not require the policy class to contain an oracle; the action-count reduction motivates a smaller sufficient feedback bound under fixed policy complexity and exploration coverage.Source: §F.2 · pp. 24–25

G Limitations

  • Improvements on GSM8K are small and inconsistent, and average accuracy decreases on held-out AQuA-RAT, limiting transfer evidence beyond classification.The authors state that observed classification transfer does not establish comparable benefits for mathematical reasoning or broader generation tasks.Source: G Limitations · p. 25
  • LDE cannot recover useful demonstrations outside the available candidate pool or handle cases requiring multiple coordinated changes.Its effectiveness therefore depends on the initial retrieval and available replacement candidates.Source: G Limitations · p. 25

Related papers

Loading 2609.33609v2…