Selecting Demonstrations for Many-Shot In-Context Learning via Gradient Matching (2025)

Abstract: In-Context Learning (ICL) empowers Large Language Models (LLMs) for rapid task adaptation without Fine-Tuning (FT), but its reliance on demonstration selection remains a critical challenge. While many-shot ICL shows promising performance through scaled demonstrations, the selection method for many-shot demonstrations remains limited to random selection in existing work. Since the conventional instance-level retrieval is not suitable for many-shot scenarios, we hypothesize that the data requirements for in-context learning and fine-tuning are analogous. To this end, we introduce a novel gradient matching approach that selects demonstrations by aligning fine-tuning gradients between the entire training set of the target task and the selected examples, so as to approach the learning effect on the entire training set within the selected examples. Through gradient matching on relatively small models, e.g., Qwen2.5-3B or Llama3-8B, our method consistently outperforms random selection on larger LLMs from 4-shot to 128-shot scenarios across 9 diverse datasets. For instance, it surpasses random selection by 4% on Qwen2.5-72B and Llama3-70B, and by around 2% on 5 closed-source LLMs. This work unlocks more reliable and effective many-shot ICL, paving the way for its broader application.

## introduction in-context learning (icl) enables pre-trained large language models (llms) to perform tasks by learning from input-output examples (or "demonstrations") provided during inference @xcite . this allows llms to adapt to new tasks through forward propagation, without weight updates via back-propagation, offering a flexible alternative to traditional fine-tuning. since the performance of icl is often sensitive to the choice of demonstrations @xcite , significant efforts have been made to im-figure 1 : instance-level retrieved demonstrations are not suitable for many-shot icl, since the variable context cannot utilize prefix-caching and can lead to theoretically o(n 2 ) inference complexity @xcite when using n-shot demonstrations. here we illustrate the runtime of llama3-70b (dubey et al., 2024) on cm-sqa @xcite using vllm with 4 a100 gpus in tensor-parallel @xcite . prove demonstration selection. most studies concentrate on instance-level retrieval, aiming to identify suitable demonstrations for each test query independently @xcite . this mainly involves considering similarity @xcite between demonstrations and the query, as well as auxiliary factors such as complexity @xcite , perplexity @xcite , difficulty @xcite , and diversity @xcite . another line of work focuses on task-level selection, which seeks to find a fixed set of demonstrations that achieves the best average performance across all test queries from a target task @xcite . current approaches primarily explore reinforcement learning on the selection policy @xcite and bayesian inference to explain the demonstration effectiveness @xcite as potential solutions. however, for the recently emerged many-shot in-context learning @xcite paradigm, demonstrations are selected simply by random in existing work @xcite @xcite , which may lead to suboptimal performance. nevertheless, existing methods for fewshot demonstration selection are not well-suited to many-shot scenarios. on one hand, applying instance-level demonstrations conflicts with caching and reusing the hidden states of the same long context in language models @xcite , leading to theoretically o(n 2 ) inference complexity @xcite for n-shot demonstrations, as illustrated in fig. 1 . on the other hand, existing task-level selection methods are designed for very few-shot scenarios (e.g., 4-shot in their implementations), facing challenges from exploration complexity @xcite and performance early saturation @xcite for more shots. to further release the potential of many-shot icl, it now requires a selection method with scalable effectiveness to many-shot scenarios. to address such research gap, we revisit icl demonstration selection from a "learning" perspective-we hypothesize that the data requirements for in-context learning and fine-tuning are analogous. building on this premise, we conduct latent concept learning @xcite on a small language model (lm) to learn the optimal in-context task embeddings. then we compute the latent gradient that each example provides to the learning process. finally, we select n-shot demonstrations from the whole fine-tuning set with the minimized l 2 distance between their average latent gradients, so as to approach the learning effect on all examples within the selected n-shot demonstrations. we validate our proposed method on 9 datasets from 5 distinct nlp tasks, each covering scenarios from 4-shot to 128-shot. our method selects demonstrations through gradient matching on relatively small models, e.g., qwen2.5-3b or llama3-8b, and surpasses the widely-adopted random selection by an average of 4% on llama3-70b and qwen2.5-72b. furthermore, the selected demonstration set exhibits transferability to closed-source llms, consistently outperforming random selection by around 2% on qwen-turbo, glm-4-flash, doubao-pro-32k, gpt-4o-mini, and deepseek-v3. our source code is available at https://github .com/zhangjf-nlp/manyshoticl-clg.git . ## methodology previous research has identified several key attributes that contribute to the effectiveness of incontext demonstrations, such as similarity @xcite , diversity @xcite , and coverage @xcite . while intuitively beneficial, these attributes lack a direct connection to the learning dynamics induced by demonstrations within the language model. we posit that the ideal demonstration set for icl should not only present static attributes of high-quality fine-tuning data but, more fundamentally, should actively facilitate and guide the language model's learning process on the target task. drawing inspiration from the principles of effective training data selection in supervised learning, particularly from the field of dataset condensation @xcite , we introduce curriculum latent gradient (clg), a novel approach for task-level in-context demonstration selection. clg leverages the concept of latent task embeddings and analyses their learning trajectory to identify demonstrations that can effectively guide the model's learning process. we begin by reviewing latent concept learning (sec. 3.1), highlighting its strengths and limitations as a foundation. we then articulate the core principles behind curriculum latent gradient (sec. 3.2), detailing how we capture and leverage the learning dynamics of latent task embeddings to guide demonstration selection. finally, we describe the demonstration selection process based on these learning dynamics (sec. 3. 3 ). an overview of our proposed methodology is depicted in fig. 2 . ## experiments we conduct experiments on 9 datasets from 5 distinct tasks, and show that demonstrations selected by our proposed method can effectively improve the many-shot icl performances over random selection as well as some straightforward methods. ## conclusion many-shot icl has emerged as a promising way to utilize llms on downstream tasks, but the random selection applied in existing work may produce sub-optimal demonstrations and learning results. in this work, we hypothesize that data requirements for in-context learning and fine-tuning are analogous, and propose curriculum latent gradient (clg) to select a demonstration set that aligns with the entire training set in learning dynamics on lms. we validate our method across various datasets and llms, showing its significant improvements compared to random selection, e.g., improving by 4% on open-source llms and by 2% on closed-source llms. this work unlocks more reliable and effective many-shot icl, paving the way for its broader adoption and application in real-world scenarios.