Kapu Nirmal Joshua

Author directory

2026

In-context learning (ICL) has emerged as a powerful paradigm for adapting large language models (LLMs) to new, data-scarce tasks using only a few carefully selected task-specific examples embedded in the prompt. However, given the limited context size of LLMs, a fundamental question arises: Which examples should be selected to maximize performance on a given user query? While nearest-neighbor-based methods have been widely adopted for this purpose, they suffer from well-known drawbacks in high-dimensional embedding spaces, including poor generalization and limited diversity. In this work, we study example selection in ICL from a principled, information-theoretic perspective. We frame the example selection task as a query-specific high-dimensional optimization problem: selecting a subset of exemplars from a larger example bank to minimize prediction error for a given query. To solve this, we first derive a principled surrogate objective that is approximately submodular, enabling the use of a greedy algorithm with an approximation guarantee. We further enhance our method by introducing an optimal design-based regularizer to encourage diversity in the selected examples. Empirically, we demonstrate significant improvements over standard retrieval methods across a suite of classification tasks, highlighting the benefits of structure-aware, diverse example selection for ICL in real-world, label-scarce scenarios.