Performance-Driven Demonstration Selection for In-Context Learning

1 Introduction

In-context learning (ICL) allows large language models (LLMs) to adapt to new tasks with strong performance using only demonstrations, without parameter updates. [Ding et al., 2026]; [Otero et al., 2026]; [Jiang et al., 2026] Despite its effectiveness, ICL is highly sensitive to the choice of demonstrations: different demonstration selections can lead to substantial performance variations, even when the demonstration number, the prompt, and the target LLM are fixed [Liu et al., 2025];[Peng et al., 2024] [Wang et al., 2023];[Shi et al., 2026];[Zuo et al., 2025].

Most existing approaches select demonstrations by optimizing heuristic or proxy criteria, such as semantic similarity between the input query and candidate demonstrations, diversity among selected demonstrations, uncertainty of demonstrations, or likelihood-based scores computed under the target LLM. [Li et al., 2025]; [Li and Qiu, 2023]; [Cai et al., 2025]; [Liu et al., 2024] While these criteria capture certain desirable properties of demonstrations, they are not optimized end-to-end with respect to the target LLM's downstream task performance. Consequently, demonstrations selected according to such proxy objectives do not necessarily translate into improved ICL performance and have been empirically shown to be suboptimal. This is corroborated by Table 1, where proxy-based selection methods (e.g., BM25, PPL [Gonen et al., 2023], and CD [Naik et al., 2023]) consistently underperform the best-performing approach.

Meanwhile, a growing line of work explores trainingbased supervised demonstration selection, where a retriever is trained using downstream task supervision [Liu et al., 2025; Ye et al., 2023; Wang et al., 2024]; [Chen et al., 2026]; [Chen et al., 2025] [Wang et al., 2025]. Despite being more principled than purely heuristic approaches, these methods still exhibit the following two limitations. On the one hand, many supervised selectors continue to rely on indirect surrogate signals, such as similarity- or uncertainty-derived labels or learned ranking scores. These signals do not explicitly model how the selected demonstrations align the LLM's performance. On the other hand, most existing training-based approaches define utility at the level of individual demonstrations and construct the in-context with k demonstrations via per-demonstration retrieval and ranking. They do not directly evaluate or optimize the set-level utility of a k-demonstration set. As a result, they fail to systematically capture interdemonstration interactions, including complementarity, redundancy, and other composition-induced effects.

Therefore, we propose that demonstration selection for ICL can be directly aligned with the target LLM performance and the ultimate objective: how well the target LLM performs on the target task given an input query-in-context pair. Meanwhile, we argue that constructing an in-context should account not only for the contribution of individual demonstrations, but also for the joint effect induced by their combination within the in-context. We therefore reformulate demonstration selection as a performance prediction problem by training a predictor (scoring model) fθ(x, C) that jointly encodes the query x and a candidate k-shot context C, and estimates the target LLM's downstream task performance for the pair (x, C). Supervised by the target LLM's actual task performance, fθ is sensitive to the set of demonstrations in C rather than any individual demonstration, thereby capturing composition effects (e.g., complementarity and redundancy) and enabling set-level optimization beyond perdemonstration ranking.

To realize this formulation, we propose PerformanceDriven Demonstration Selection (PDDS). PDDS learns a scoring model that predicts the target LLM's task performance for an input query-in-context pair. The scoring model is trained using supervision derived from the actual outputs of the target LLM, allowing it to capture how different demonstrations and in-contexts affect performance. Once trained, the scoring model can be used to efficiently evaluate candidate demonstrations and in-context at inference time, without requiring additional LLM calls or validation-time feedback. This design enables demonstration selection to be both performance-aligned and deployment-friendly.

A key challenge in PDDS stems from the combinatorial nature of in-context demonstration construction, which makes exhaustive evaluation over demonstration subsets computationally intractable. To address this challenge, we adopt a lightweight two-stage inference-time procedure guided by the PDDS scoring model. Specifically, we preselect the topwk most promising demonstrations based on their predicted scores and sort them in descending order. We then form a small set of candidate k-shot contexts from this sorted pool (e.g., via a sliding window), score each candidate in-context as a whole, and select the highest-scoring one as the final in-context. By leveraging the learned scorer, our procedure avoids combinatorial search while staying directly aligned with the target objective.

We conduct extensive experiments on 6 NLP tasks and 8 LLMs (ranging from 1B to 70B parameters) to evaluate the effectiveness of PDDS. Results show that performance-driven demonstration selection consistently improves ICL performance over strong heuristic- and learning-based baselines. Further analyses demonstrate that our method generalizes across different LLMs and datasets, highlighting the benefits of directly modeling LLM performance for in-context construction. Meanwhile, the PDDS scoring model can be applied as a plug-and-play component to other ICL methods and improve their performance.

Our contributions can be summarized as follows: ❶ We reformulate demonstration selection for in-context learning as a performance prediction problem, providing a new perspective that directly aligns selection with downstream task performance. ❷ We propose a performance-driven demonstration selection framework that learns to predict LLM task performance from input query-in-context pairs and uses this prediction to construct effective in-contexts. ❸ We demonstrate through extensive experiments that our approach outperforms existing demonstration selection methods in both effectiveness and robustness, while remaining efficient at inference time.

2 ProblemFormulation

We consider the standard ICL setting with a frozen target LLM. Let Dr = {(xi, yi)} denote a labeled retrieval dataset, where xi is a text and yi is its corresponding ground-truth label. Given a test input xt, the goal of ICL is to construct a context C = {(xj , yj)}, j = 1..k consisting of k demonstrations, such that the LLM achieves strong task performance when conditioned on (xt, C).

2.1 Performance-Driven Demonstration Selection

In contrast, we formulate demonstration selection as a performance prediction problem. Our objective is to directly model the task performance of the target LLM conditioned on an input query-in-context pair. Let Perf(LLM; xt,C) denote a task-specific performance metric (e.g., MSE [Hyndman and Koehler, 2006], BLEU [Post, 2018], ROUGE-1 [Ganesan, 2018]) achieved by the LLM when conditioned on input query xt and in-context C. We aim to learn a scoring model fθ (·):

where fθ predicts the downstream task performance of the LLM for a given input query-in-context pair. Under this formulation, demonstration selection is defined as

where Dr denotes the retrieval dataset. This formulation directly aligns the selection objective with downstream task performance.

3 Method

Figure 1 illustrates the overall framework of PDDS. Given an input query, PDDS first retrieves a pool of candidate demonstrations and constructs a set of candidate in-contexts. Unlike heuristic criteria, PDDS relies on a lightweight performance predictor trained with supervision from the frozen target LLM: for each query we construct candidate in-contexts and use the target LLM's task metric as labels to fit fθ(x,C) ≈ Perf(LLM; x,C). At inference time, PDDS performs a twostage in-context construction by (i) preselecting and ranking promising demonstrations and (ii) selecting the final k demonstrations at the context level, where all scores are produced by fθ without additional target-LLM calls. The target LLM is invoked only once with the selected in-context to produce the final prediction.

3.1 Training Data Construction

As the subfigure (a) Dataset construction in Figure 1 shows, training the scoring model for input query-in-context pairs requires supervision that reflects the downstream task performance of the target LLM. Therefore, we randomly split the retrieval dataset Dr into a query set Q and a demonstration pool D. For each query instance (x, y) ∈ Q, we first retrieve candidate demonstrations from D based on semantic similarity. Specifically, let ϕ(·) denote a sentence embedding function and s(·, ·) a similarity measure. Each candidate (x′ , y′ ) ∈ D is ranked by

Let πx denote the ranking permutation in descending order of sim(x, x′ ). Based on this ranking, we construct two types of in-contexts for each query. The single-demonstration context contains the most similar demonstration,

while the k-demonstration in-context consists of the top-k ranked demonstrations,

We then query the frozen target LLM with each input queryin-context pair (x,C) and compute a task-specific performance score. Formally, given a performance metric Perf(·), we obtain

These scores serve as supervision signals for learning the performance predictor (scoring model). This construction yields informative training signals without requiring exhaustive incontext enumeration. By contrasting in-contexts of different sizes and informational content, the scoring model is encouraged to learn how contextual composition influences downstream performance, rather than relying on superficial properties of individual demonstrations. As a result, the learned predictor (scoring model) can generalize beyond the specific contexts observed during training and be applied to arbitrary candidate in-contexts at inference time.

3.2 Performance ScoringModel

Given the constructed training set of input query-in-context pairs and their corresponding performance scores, we aim to learn a scoring model that predicts the downstream task performance of the target LLM. The scoring model takes an input query x together with a candidate in-context C as inputs and outputs a scalar score estimating the LLM's performance when conditioned on (x, C).

Performance Definition. Let Perf(LLM; x,C) ∈ R denote a task-specific performance score of the LLM under input x and context C. Without loss of generality, we define Perf(·) such that larger values indicate better performance. For evaluation metrics where lower values correspond to better performance (e.g., error or loss), we use the negated metric value as Perf(·). This convention ensures a consistent maximization objective across different tasks.

Model Input Representation. The scoring model operates on the full input query-in-context pair (x,C). The in-context C = {(xj , yj)}|C| j=1 is treated as an unordered set of demonstrations, where each demonstration consists of an input query-in-context pair. Rather than scoring individual demonstrations independently, the model evaluates the in-context as a whole, enabling it to capture interactions among demonstrations that jointly influence downstream performance.

Scoring Function. Let fθ(·) denote the performance scoring model parameterized by θ. The scoring model predicts a scalar score

which serves as an estimate of the true performance Perf(LLM; x, C).

Learning Objective. During training, supervision for the scoring model is obtained from the actual outputs of the target LLM. Given a training instance (x,C, r), where r = Perf(LLM; x,C), we optimize the model by minimizing a regression mean squared error (MSE) loss [Hyndman and Koehler, 2006]:

This objective trains the scoring model to accurately rank different in-contexts by their downstream task performance. Unlike prior demonstration selection methods that rely on proxy objectives or independently score individual demonstrations, our scoring model directly estimates the task performance of the target LLM under a complete input query-in-context pair. Therefore, the scoring model is trained in an end-toend manner aligned with the downstream performance. This alignment empirically leads to more effective demonstration selection than proxy-driven criteria.

3.3 Inference-Time In-Context Construction

As the subfigure (c) Testing process in Figure 1 shows, at inference time, given a test input query xt, we aim to construct a k-demonstration in-context that maximizes the predicted task performance. However, directly enumerating and evaluating all possible in-contexts is computationally infeasible due to the combinatorial nature of demonstration selection. Therefore, we use a two-stage procedure guided by the learned performance scoring model.

Stage 1: Preselect promising demonstrations. For each candidate demonstration (xi, yi) ∈ Dr (the retrieval dataset), we score the corresponding single-demonstration in-context:

We then keep the top-wk demonstrations with the largest ˆ ri , denoted as Dwk(xt).

Stage 2: Construct and score candidate contexts. We sort Dwk(xt) by ˆ ri in descending order and generate candidate in-contexts using a sliding window of size k:

where (x(j), y(j)) denotes the j-th demonstration in the sorted list. Each candidate in-context C ∈ C(xt) is evaluated by

and the final in-context is selected as

All inference-time scores are computed by the lightweight predictor (scoring model) fθ, without any additional calls to the target LLM. We feed the query and selected in-context to the target LLM for final prediction.

4 Experiment

4.1 Experimental Setup

(i)Tasks and Datasets. We cover six task types over eight datasets: Translation (WMT19 En-Zh dataset [Barrault et al., 2019]), text expansion (Gigatiny dataset [Rush et al., 2015]), summarization (Gigaword dataset), Semantic Textual Similarity (STS) (STS-B/STS14 datasets [Cer et al., 2017; Agirre et al., 2014]), classification (SST-5/Emotion datasets [Socher et al., 2013; Saravia et al., 2018]), and Translation quality evaluation (TQA) (EN-CS dataset [Ma et al., 2019; Barrault et al., 2019]), spanning classification/generation and discrete/continuous outputs.

(ii) Metrics and Baselines. We use BLEU [Post, 2018] (Translation), ROUGE-1 [Ganesan, 2018] (expansion/summarization), accuracy (classification), and report Mean Squared Error (MSE) [Hyndman and Koehler, 2006] for the STS and translation-quality evaluation tasks. Baselines include proxy-based demonstration selection methods DKNN [Li et al., 2025], CR [Li and Qiu, 2023], KATE [Liu et al., 2022], CD [Naik et al., 2023], ICCL [Liu et al., 2024], PPL [Gonen et al., 2023] (k=10), and training-based demonstration selection methods TTF [Liu et al., 2025], CEIL [Ye et al., 2023], and MoD [Wang et al., 2024].

(iii) Target LLMs and Other Setups. Our main LLMs are GLM-4 9B, LLaMA 3.1 8B, and Qwen2.5 7B; we further test 1B-70B LLMs (LLaMA 3.2 1B, DeepSeek 32B [Guo et al., 2025], LLaMA 3.3 70B and Qwen3 4B/8B [Yang et al., 2025]). The number of demonstrations for all methods is 10. We randomly split the retrieval dataset Dr into a query set Q and a demonstration pool D, where Q and D account for 20% and 80% of the data, respectively. The wk in Eq. 10 is 20.

4.2 Main Results

Table 1 reports the overall performance across six tasks and eight datasets. PDDS consistently achieves strong results and attains the best performance in the majority of settings. Specifically, it improves BLEU by 0.032 on machine translation, increases ROUGE-1 by 0.009 and 0.016 on summarization and expansion, respectively, reduces MSE by 0.060 on STS14, and exceeds the strongest competing baseline by 5.8% accuracy on text classification.

Overall, these results show that directly optimizing predicted LLM performance leads to more effective and robust demonstration selection than proxy-driven heuristics or contrastive-learning-based methods.

4.3 Ablation Study

We conduct ablation studies to quantify the contribution of each design choice to overall performance, including the scoring model, the two-stage in-context construction strategy, and the preselection size wk. Ablating performance-driven scoring model (fθ vs. proxy criteria). We ablate fθ by constructing in-contexts using only Eq. 5 (KATE-style retrieval) without performance scoring (Figure 2). Without the scoring model, the performance of PDDS degrades noticeably. Specifically, the average accuracy drops by 9.47%, the average MSE worsens by 0.12, the average ROUGE-1 decreases by 0.02, and the average BLEU drops by 0.09. These consistent drops underscore the scoring model's critical role in PDDS.

Ablating two-stage in-context construction. In PDDS, in-context construction is performed in two stages. In Stage 1, PDDS retrieves wk candidate demonstrations and ranks them in descending order according to the predicted performance scores. In Stage 2, the scoring model is further applied to select the final k demonstrations from the ranked candidate set to form the final in-context. We now analyze the contribution of each stage separately.

Ablating Stage 1. To evaluate the role of Stage 1, we replace performance-driven preselection with the similarity criterion in Eq. 3 to retrieve the top-wk most similar candidates, and then apply Stage 2 to construct the final in-context. As shown in Figure 2, removing Stage 1 consistently hurts performance: the average accuracy drops by 4.52%, the average MSE worsens by 0.04, the average ROUGE-1 decreases by 0.01, and the average BLEU drops by 0.01. This underscores the importance of Stage 1 in PDDS.

Ablating Stage 2. To ablate Stage 2, we form the incontext by directly selecting the top-k demonstrations from the ranked wk candidates produced in Stage 1. As shown in Figure 2, excluding Stage 2 consistently degrades performance: average accuracy drops by 5.02%, MSE increases by 0.26, and ROUGE-1/BLEU each decrease by 0.01. These results underscore the importance of Stage 2 in PDDS.

Sensitivity to preselection size wk. We set wk = 20 by default and evaluate wk ∈ {10, 15, 20, 25}. Performance improves as wk increases but quickly saturates; e.g., on Emotion with Llama, accuracy rises from 51.5% (wk = 10) to 62.1% (wk = 25) and then stays around 62.1%. A moderate candidate pool captures most gains, while larger wk mainly adds computation.

5 Generalization

Cross-retrieval generalization (plug-and-play). Given a fixed external retriever, we only score the retrieved candidate pool (with wk candidates) using our performancedriven scoring model and then select the top-k demonstrations accordingly, without retraining or modifying the retriever. As shown in Table 2, replacing the original ranking/selection strategy with our plug-and-play scoring consistently improves performance. Specifically, it increases the average accuracy from 48.77% to 55.32% and the average ROUGE-1 from 0.059 to 0.067. These results indicate that the scoring model captures performance-relevant signals beyond retriever-specific heuristics and can be integrated in a plug-and-play manner into existing ICL pipelines.

Cross-dataset generalization. We examine whether the scoring model generalizes to unseen datasets beyond those used for training. Specifically, we train the scoring model on SST5 and evaluate it on Emotion under the same classification setting. As shown in Table 3, the transferred scoring model achieves an average accuracy of 65.47% and consistently outperforms the baselines reported in Table 1. These results indicate that performance-driven selection learns dataset-agnostic signals that remain predictive across distribution shifts, rather than overfitting to the training corpus.

Cross-LLM generalization. Finally, we evaluate whether the scoring model transfers across target LLMs. As shown in Table 5, a scoring model trained with one target LLM can be directly applied to a different target LLM. Although cross-LLM transfer is slightly worse than training on the target LLM itself (by Llama3.1 8B), it still yields strong performance: when trained on GLM4 9B, the PDDS achieves 53.2 accuracy on Llama3.1 8B. We conjecture that, since both GLM4 9B and Llama3.1 8B are trained to approximate the same underlying text-to-label mapping, they may induce similar decision boundaries (or decision regions) over the input query-in-context space. As a result, they tend to exhibit similar behaviors on the same input query-in-context pair, which helps explain why the scoring model can generalize across LLM backbones. Overall, these results suggest that the scoring model captures relatively model-agnostic signals correlated with in-context performance.

6 Discussion

Evaluation across target LLM scales and families. We further evaluate PDDS on five additional target LLMs (1B70B). As shown in Table 6, PDDS remains competitive across scales and families, achieving up to 79.2% accuracy. Overall, PDDS transfers well across LLM families and scales, suggesting the scorer captures performance-relevant signals.

Cost considerations. We evaluate PDDS overhead on SST5 in terms of training time, storage, token cost, and latency. Training the scoring model takes 1.2 hours and 472 MB. At inference, PDDS uses 189.5 tokens and 0.31 seconds per query on average (Table 4). Overall, the one-time training cost is amortized over repeated use, making PDDS practically acceptable.

Training data efficiency (few-shot supervision). We study whether the PDDS scoring model can be effectively trained with only a small amount of labeled data. Specifically, we limit the number of training queries to 30. To enrich supervision in this setting, we construct additional single- and k-demonstration in-contexts, C (1) (x) and C(k) (x) (Equ. 4 and 5). These in-contexts are generated not only via similarity-based retrieval, but also by leveraging retrieval strategies from several ICL methods, including BM25, CD, and PPL. This yields a diverse set of query-in-contextperformance triples despite the limited supervision, enabling stable training of the scoring model. As shown in Table 7, PDDS still achieves up to 52.0% accuracy with 30-shot supervision.

7 Conclusion

We revisit demonstration selection for ICL from a performance-driven perspective. Specifically, we formulate it as predicting downstream LLM performance for an input query-in-context pair, rather than optimizing heuristic proxies. Based on this formulation, we propose a lightweight and end-to-end framework that learns a scoring model to directly score and construct effective in-contexts.

Extensive experiments across six task categories, eight datasets, and eight target LLMs show that our method consistently improves ICL performance by aligning context construction with predicted LLM performance. Further analyses demonstrate that the scoring model generalizes across datasets, tasks, and target LLMs, and can be effectively trained with few-shot data, making the approach practical and scalable. Meanwhile, our method can be applied as a plug-and-play approach to improve the performance of other ICL methods. Overall, our results suggest that aligning context construction with predicted LLM performance offers a principled alternative to proxy-based heuristics, and reframes context selection as an explicit, performance-oriented decision problem in ICL.

Acknowledgements

Supported by the Shenzhen Science and Technology Program (CJGJZD20240729141505007), the National Natural Science Foundation of China (Nos. U2541229, 62441619, and 62411540034), and the Ningbo Science and Technology Innovation 2025 Major Project (2025Z027), and the National Natural Science Foundation of China under Grant 62502550, Shenzhen Science and Technology Program (KJZD20240903095700001).

Contribution Statement

† Wenqiang Wang and Mingbo Yang contributed equally to this work as co-first authors. ∗ Xiaochun Cao is the corresponding author. Wenqiang Wang made the primary contribution to this work, including developing the main idea, formulating the performance-driven demonstration selection framework, designing the scoring model and the two-stage incontext construction strategy, conducting the main analysis, and writing the manuscript. Mingbo Yang contributed to code implementation, experiment execution, and result organization. Aiping Zhang contributed to methodological discussions, experimental analysis, and manuscript revision. Yan Xiao and Peng Chen contributed to dataset preparation, baseline comparison, and experimental evaluation. Jianjie Huang contributed to technical discussions and manuscript refinement. Xiaochun Cao supervised the project, provided overall research guidance, and contributed to the final manuscript revision.