Auto-regressive In-context Demonstration Selection

Effective demonstration selection is crucial for maximizing large language model (LLM) performance in few-shot in-context learning. Because of effects such as recency bias, the effectiveness of demonstrations depends heavily on their contextual relationship to the specific query and on the ordering in which they are presented, making demonstration selection a complex combinatorial problem. To address these two challenges, we introduce AutoSelect, a novel framework that formulates demonstration selection as an auto-regressive sequential decision process. At each step, AutoSelect embeds the query and previously selected demonstrations into matrix representations to preserve structural information, and a trainable policy model sequentially selects the next best exemplar. To navigate the factorial space of demonstration permutations, our framework formulates a Kullback-Leibler (KL)-regularized optimization problem, from which an optimal policy induces an optimal Plackett-Luce (PL) ranking over all possible demonstration sequences. We prove that minimizing a tractable policy-level cross-entropy (CE) loss provably bounds the worst-case discrepancy between our policy's induced PL ranking and the optimal one, enabling tractable prioritization of high-quality sequences. Empirically, AutoSelect outperforms existing heuristic and learning-based methods across nine diverse datasets, achieving up to an 11\% improvement over the strongest baseline. Analytical studies and a case study further highlight AutoSelect's key properties, as well as its transferability and generalizability.

Large language models can answer new questions more accurately when their prompts include a few examples as context, but picking those examples is surprisingly tricky. An example that helps one question may mislead another, and simply rearranging the same examples can change the answer. Trying every possible set and order is far too slow for practical use. AutoSelect addresses this by building each prompt one example at a time. For the current question, it looks at the examples already chosen, decides which one should come next, and can stop when adding more is unlikely to help. This lets AutoSelect account for how examples support or interfere with each other, instead of judging each example alone. Across nine tasks and datasets, AutoSelect outperformed existing selection methods, improving results by up to 11% over the strongest comparison method. By choosing the right number of examples in a better order, AutoSelect can make language models more reliable and efficient when they are guided by in-context examples.