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, such as similarity, diversity, or uncertainty, and select demonstrations independently, which may misalign with downstream performance and overlook set-level composition effects. Therefore, we propose Performance-Driven Demonstration Selection (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 entire in-context 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, 8 datasets, and 8 target LLMs ranging from 1B to 70B parameters, PDDS achieves state-of-the-art results and generalizes well across LLMs and datasets. PDDS remains effective with as few as 30 training instances and can be integrated as a plug-and-play component to enhance existing ICL methods.