Performance-Driven Demonstration Selection for In-Context Learning

Abstract

In-context learning (ICL) enables large language models (LLMs) to adapt to new tasks with considerable performance gains, yet its effectiveness is highly sensitive to the choice of demonstrations. Most existing selection methods rely on heuristic or proxy signals (e.g., similarity, diversity, or uncertainty) and select demonstration independently, which may misalign with downstream performance and overlook set-level composition effects. Therefore, we propose \textbf{P}erformance-\textbf{D}riven \textbf{D}emonstration \textbf{S}election (PDDS), which directly aligns demonstration selection with ICL performance. PDDS formulates selection as predicting the target LLM’s downstream task performance for a given query–in-context pair, replacing proxy heuristics with a performance-aware objective. Unlike prior approaches that rank demonstrations independently before composing the prompt, PDDS evaluates the \textit{entire} in-context as a set, capturing inter-demonstration interactions and composition effects. PDDS trains an end-to-end scorer using supervision from the target LLM’s actual task outcomes. At inference time, it selects high-scoring in-context sets without additional target LLM calls or validation-time feedback. Across 6 NLP tasks, 10 datasets, and 8 LLMs (1B--70B, including GPT-4o), PDDS achieves state-of-the-art results and generalizes well across LLMs, datasets, and tasks. PDDS remains effective with as few as 30 training instances, scales to candidate pools of up to 100 demonstrations, and can be integrated as a plug-and-play component to enhance existing ICL methods.