PaperlayerView exact revision on arXiv

Source-linked AI summary

Data Efficient Sample Selection for In-Context Learning

V Venktesh, Cem Levi, Avishek Anand

arXiv:2609.06670v2cs.LGcs.CL

TL;DR

Selecting demonstration subsets for ICL is challenging because existing methods may not capture complex subset–performance relationships or generalize from static task-level choices to unseen queries. DEARICL trains a nonlinear, instance-level subset ranker using gap-index bandit sampling of borderline subsets. It reports 8.08–15.9% performance gains on smaller LLMs.

  • Problem

    Existing ICL selection methods may use static task-level subsets that fail to generalize to unseen queries, and linear surrogates do not capture complex subset–performance dependencies.

    Source: §1 · pp. 1–2

  • Method

    DEARICL uses gap-index top-m arm identification to sample borderline subsets as an auxiliary objective for training a nonlinear surrogate for instance-level ranking.

    Source: §1 · p. 2

  • Results

    8.08–15.9% performance improvement is reported for DEARICL on smaller LLMs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The trained surrogate can rank example subsets for instance-level selection, and DEARICL adds negligible inference overhead in that setting.

    Source: §3.2 · p. 4, §6 · p. 9

    The work does not focus on theoretical regret bounds, which the authors describe as hard to derive for gap-index frameworks.

    Source: §8 · p. 9

Abstract

from arXiv · show

The In-context learning (ICL) paradigm aids large language models (LLMs) to adapt to new tasks without need for fine-tuning. However, selecting an optimal combination of demonstration examples from a large pool of example subsets is a challenging problem. Existing approaches for selection do not model the complex relationship between ICL samples and downstream LLM performance. They typically perform static task-level selection, choosing subsets once offline, which can fail to generalize to unseen queries. We introduce DearICL (Data Efficient Algorithm for Ranking) ICL samples, a new framework that models demonstration example selection as a subset ranking problem. DearICL employs a non-linear surrogate employing a differentiable sorting objective within a gap-index bandit algorithm. The gap-index based approach enables fine-grained separation of good arms and borderline arms, which is used as an auxiliary objective to train the non-linear surrogate through sufficient sampling of borderline arms, supporting instance-level subset ranking. On exemplar selection benchmarks with open-source LLMs, DearICL achieves 8.08-15.9% accuracy gains over strong linear bandit baselines, with low sample complexity. Code and data: https://github.com/VenkteshV/DearICL.

1 Introduction

Selecting ICL demonstrations is difficult because subsets interact, while static selection can miss the needs of unseen queries. DEARICL uses gap-index sampling to train a nonlinear, instance-level subset ranker from LLM feedback.

1 Introduction

  • Subset interactions create a large search space, and linear reward surrogates cannot represent complex subset–performance dependencies.Source: §1 · p. 1
  • Static task-level selection may fail on unseen queries, while running a bandit algorithm for each query is computationally expensive.Source: §1 · pp. 1–2
  • DEARICL uses top-m selection as an auxiliary objective to sample borderline subsets and train a nonlinear surrogate for instance-level ranking.The framework reports 8.08–15.9% performance improvements on smaller LLMs.Source: §1 · p. 2
  • 8.08–15.9% performance improvement is reported for DEARICL on smaller LLMs.Source: §1 · p. 2

2 Related Works

Prior work on exemplar selection spans automated selection, static task-level methods, and instance-level methods; the paper also situates its ranker and bandit objective within learning-to-rank and top-m identification research.

2 Related Works

  • Automated ICL exemplar selection includes reinforcement learning, trained retrievers, determinantal point processes, constrained optimization, and learning-free similarity, complexity, or MMR methods.Source: §2 · p. 2
  • Static bandit- or Bayesian-optimization methods preselect task-level examples, whereas instance-level methods offer query-specific selection at higher inference-time computational cost.Source: §2 · p. 2
  • DEARICL learns rankings from LLM feedback and uses fixed-confidence top-m identification, unlike click-feedback rankers and methods focused mainly on best-arm identification.GIFA provides a unified top-m framework but requires many gap-index computations and comparisons, leading to high sample complexity.Source: §2 · p. 2

3 Methods

The methods frame demonstration selection as top-m subset identification and use it to train a nonlinear ranker for query-specific inference. DEARICL’s analysis supplies pairwise-gap and sample-complexity bounds for the auxiliary selection process.

3.1 Problem Definition

  • For ICL, a model receives input–output demonstrations followed by a test input, and predicts its output; demonstrations can include rationales for some tasks.Source: §3.1.1 · p. 3

3.1.1 Subset Selection for ICL - Preliminaries and Limitations of Existing Formulation

  • The formulation selects k demonstrations from a candidate pool to form the test context, while offline static subsets may not generalize to unseen queries and per-query bandit convergence is costly.Source: §3.1.1 · p. 3

3.1.2 Top-m Arm Selection as Auxiliary objective

  • DEARICL treats candidate subsets as bandit arms and uses top-m identification with a nonlinear surrogate to separate near-optimal from suboptimal subsets.The framework provides theoretical bounds on pairwise-gap error and the arm sampling required to fit the surrogate.Source: §3.1.2 · pp. 3–4

3.2 DEARICL: A sample-efficient algorithm fitting of non-linear Ranking surrogate

  • A differentiable sorting objective trains the nonlinear surrogate on LLM-feedback rewards, while gap-index sampling targets ambiguous borderline arms for query-specific ranking at inference.The algorithm updates its shortlist and challengers, samples arms, acquires feedback, and retrains the ranker.Source: §3.2 · pp. 4–5

3.3 Sample Complexity Theoretical Bounds

  • The analysis gives a high-probability sample-complexity upper bound for DEARICL’s auxiliary objective, representing the expected LLM calls needed for convergence.The stopping-time bound applies on the good-gap event; the analysis also establishes a pairwise-gap estimation bound.Source: §3.3 · pp. 5–6

4 Experimental Setup

The experiments address whether DEARICL converges faster, improves downstream performance through instance-level selection, and retains efficiency. They evaluate reasoning and translation tasks using open-source LLMs and compare against selection and ranking baselines.

4 Experimental Setup

  • The evaluation asks whether DEARICL speeds top-m subset identification, improves downstream performance with instance-level selection, and does so efficiently.Source: §4 · pp. 6–7

4.1 Experimental Setup

  • Experiments use GSM8K and AquaRAT for mathematical reasoning and WMT 2019 for translation, measured with Exact Match or BertScore, respectively.BertScore between generated and ground-truth rationales and answers also supplies reward feedback for GSM8K and AquaRAT.Source: §4.1 · p. 7
  • The primary evaluation uses the open-source Llama3.2-3b model, with additional reported results on gpt-4o-mini and Deepseek-R1:7B.All approaches use a 1,000-token generation limit and temperature 0.25.Source: §4.1 · p. 7
  • Comparisons include CASE, BRIDGE, EXPLORA, PiRank, UDR, and other learning-to-rank baselines, with DEARICL's bandit settings shared for fair comparison.The setup uses m = 10, |Ct| = 5, δ = 0.05, ε = 0.1, and example subsets sampled with replacement across five training clusters.Source: §4.1 · p. 7

5 Results

DEARICL's instance-level ranker outperforms static, retrieval-based, dynamic bandit, and learning-to-rank alternatives. It also converges faster and reduces selection calls, with only negligible inference-time overhead.

5.1 Empirical Verification of Convergence in DEARICL

  • DEARICL's gap-index decreases toward zero and its top-m simple regret declines across rounds; it converges earlier than CASE, which stagnates after round 70.The passage reports corresponding gap-index and simple-regret results for other datasets in the appendix.Source: §5.1 · p. 7

5.2 Performance Comparison for Example Subsets Selection

  • DEARICLdynamic outperforms static selection, KNN/MMR, dynamic CASE, and learning-to-rank approaches, reaching up to 15.99% over dynamic CASE on AquaRAT.The reported gains are attributed to modeling subset reward structure and separating borderline from top-m arms; gap-index exploration also outperforms random exploration.Source: §5.2 · p. 8

5.3 Efficiency Comparison

  • DEARICL converges in 79 versus 238 rounds on AquaRAT and 130 versus 510 on GSM8K, and reduces top-m selection LLM calls by approximately 63–75% versus prior approaches.Table 3 presents the LLM-call comparison.Source: §5.3 · p. 8, Table 3 · p. 8
  • DEARICLdynamic adds only a few milliseconds of latency over static CASE and is more efficient than CASEdynamic at test time.Its lightweight ranker avoids the extra inference overhead associated with NeuralNDCG.Source: §5.3 · p. 9

6 Conclusion

DEARICL uses a sample-efficient gap-index bandit objective to fit a nonlinear surrogate for instance-level ICL subset selection, with faster convergence and negligible inference overhead.

  • DEARICL uses a sample-efficient gap-index MAB framework as an auxiliary objective to fit a nonlinear surrogate that scores example subsets for instance-level ICL selection.The conclusion reports faster convergence than existing MAB algorithms and negligible inference overhead for instance-level use.Source: §6 · p. 9

8 Limitations

The approach currently relies on supervised question–answer pairs, and the paper does not derive theoretical regret bounds. The authors identify broader settings and ranking losses as future extensions.

  • The current approach relies on supervised question–answer pairs, which may be unavailable in some domains.Selecting questions and augmenting them with answers using LLMs or other data-augmentation mechanisms is proposed but outside this work's scope.Source: §8 · p. 9
  • The work provides sample-complexity bounds but does not focus on theoretical regret bounds for its nonlinear stochastic gap-index framework.The authors note that deriving regret bounds for gap-index frameworks is difficult and report empirical convergence of regret and gap-index instead.Source: §8 · p. 9
  • Future directions include extending the framework beyond ICL and adapting it to other relaxed learning-to-rank losses.The authors also identify regret bounds as a potential extension.Source: §8 · p. 9

9 Ethical Considerations and Risks

The paper reports using public math and machine-translation datasets without private or harmful information, and describes the datasets and LTR hyperparameter table.

9 Ethical Considerations and Risks

  • The experiments use publicly available math and machine-translation datasets without private or harmful information.The authors also state they do not prompt LLMs in a format intended to elicit harmful or biased information.Source: §9 · p. 9

B Dataset Description

  • AquaRAT contains 100,000 algebraic word problems, with 254 problems each in its dev and test sets.Problems include answers and step-by-step rationales.Source: B Dataset Description · p. 13
  • GSM8K contains 8.5K multi-step reasoning problems, and evaluation uses 1,319 test questions.Source: B Dataset Description · p. 13
  • WMT 19 uses the English–Chinese translation split and has 173k training sentence pairs.The passage gives a 3,981-sentence test-set example for WMT18 zh-en, not WMT19.Source: B Dataset Description · p. 13

C LTR Hyperparameters

D Gap and simple regret analysis for GSm8K and WMT19

For GSM8K and WMT19, DEARICL’s gap index approaches the stopping threshold as simple regret minimizes across rounds. The ablation reports that some baselines improve with DEARICL-collected data but remain behind DEARICLdynamic.

F Data ablation and extension to other LTR loss functions

G Dataset prompts

  • Table 4 categorizes LTR-model hyperparameters by loss function and framework.Source: Table 4 · p. 14

D Gap and simple regret analysis for GSm8K and WMT19

F Data ablation and extension to other LTR loss functions

H Qualitative Analysis

A qualitative example compares the demonstrations selected for an algebra problem by DEARICLdynamic and NeuralNDCG.

H Qualitative Analysis

I Proof for Theorem 3.2

The proof establishes a high-probability pairwise-gap confidence event for DEARICL by bounding dropout noise, SGD stability error, and model bias. It uses this event to support the top-m sample-complexity analysis.

I Proof for Theorem 3.2

  • For an AquaRat algebra test instance, DEARICLdynamic selected three similar questions, while NeuralNDCG selected none.The exemplars’ rationales are omitted in the presentation; the passage says AquaRat exemplars include rationales.Source: I Proof for Theorem 3.2 · p. 15
  • Uniform stability controls prediction changes when one sampled arm-and-reward example is added to the surrogate’s training set through one-epoch SGD.The stated stability conditions include bounded gradients and a Lipschitz loss.Source: I Proof for Theorem 3.2 · p. 18
  • The proof combines bounds on dropout noise, SGD stability, and model bias to establish the pairwise confidence event used in DEARICL’s sample-complexity analysis.The confidence width is defined for pairwise gaps between arms.Source: I Proof for Theorem 3.2 · pp. 16–19
  • The differentiable NDCG relaxation has an L-Lipschitz loss when the reward, discount, and soft-permutation operator satisfy the stated bounds.The soft-permutation Lipschitz constant is specified as O(1/τ).Source: I Proof for Theorem 3.2 · p. 19

J.2 Proof Blueprint for Theorem 3.3

The proof turns the stopping rule and pairwise confidence widths into per-arm and total sample bounds. The resulting high-probability upper bound is stated in Theorem 3.3.

J.2 Proof Blueprint for Theorem 3.3

  • The proof inverts the stopping condition to bound how many samples an arm needs for its confidence width to become small relative to its gap.The stated intuition is that nonzero-gap arms must be sampled sufficiently before the stopping rule triggers.Source: §J.2 · p. 21
  • For clarity in deriving the sample bound, the proof uses the leading variance term and sets aside stability and bias terms, which it describes as negligible in this derivation.The predictive variance is stated to decrease as an arm receives more samples.Source: §J.2 · p. 21
  • Theorem 3.3 states an upper bound on sample complexity; the cited proof gives per-arm bounds on event E and sums them over arms to obtain a total-sample upper bound.Source: §J.2 · p. 22
Loading 2609.06670v2…