Dynamic $k$-shot In-Context Learning

Wenqiang Wang, Yan XIAO, Xiaojun Jia, Yangshijie Zhang, Peng Chen, Bin Zeng, Xiaochun Cao

Keywords: In-context learning

Abstract: In-context learning (ICL) allows large language models (LLMs) to learn new tasks from demonstrations and to predict unseen inputs without parameter updates. Existing studies typically fix the number of demonstrations as a static hyperparameter (e.g., 5 or 10), overlooking the variability across models and inputs. We empirically find that the same query text may yield different outcomes depending on the number of demonstrations used. Motivated by this observation, we propose Dynamic-$k$ In-Context Learning (D-$k$-ICL), a novel method that adaptively determines the most suitable number of demonstrations for each query text. The core component is a performance predictor—a neural network that jointly encodes the query text and candidate in-contexts (constructed with varying demonstration counts) to estimate expected inference quality. At inference time, we retrieve the top-$k$ semantically similar demonstrations and progressively vary $k$ to generate candidate in-contexts. The predictor then selects the candidate most likely to yield the best output, thereby dynamically adapting both the number and composition of demonstrations. Across three LLMs and eight datasets, D-$k$-ICL achieves considerable results, with up to 77.8\% accuracy, 0.641 MSE, 0.271 ROUGE-1, and 0.295 BLEU. Furthermore, even when trained under few-shot, weakly supervised, or self-supervised settings, the predictor remains effective. Finally, D-$k$-ICL consistently improves performance on commercial LLMs such as GPT-4o, demonstrating its robustness and broad applicability.

Primary Area: applications to computer vision, audio, language, and other modalities

Submission Number: 25480