Easier to Judge than to Find: Predicting In-Context Learning Success for Demonstration Selection

TL;DR

This paper introduces DiSP, a framework for predicting the success of demonstration-context pairs in in-context learning, improving selection efficiency and speed.

Contribution

DiSP offers a novel sample-and-judge approach that predicts demonstration success, enabling faster and more effective demonstration selection in ICL.

Findings

DiSP improves average accuracy by up to 3.4% over baselines.

Achieves up to 23x speedup in end-to-end inference.

Effective across multiple classification datasets with large language models.

Abstract

In-context learning (ICL) is highly sensitive to which demonstrations appear in the prompt, but selecting them is expensive because the space of possible demonstration contexts and combinations is enormous. We argue that demonstration selection is \emph{easier to judge than to find}: predicting whether a specific query--context pair $(q,D)$ will succeed is cheaper and more general than searching for an optimal $D^\star$. Based on this insight, we propose DiSP, a sample-and-judge framework that stratifies queries by difficulty. DiSP runs random demonstration trials to estimate success rate of each training query, trains a lightweight router to predict difficulty from the query, and trains level-specific judges for sampled demonstrations. At inference, DiSP performs stop-on-acceptance judging under an explicit budget, emitting diagnostic risk tags when no suitable context is found. Across five classification datasets with Llama~3--8B and Qwen~2.5--7B, DiSP achieves the best average accuracy, improving over strong learned selection baselines by up to 3.4\%, while achieving up to $23\times$ end-to-end wall-clock speedup.