# Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimation

arXiv: https://arxiv.org/abs/2508.19999 · HTML: https://mlsanity.com/paper/2508.19999/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 39 references that are in the map.
- First posted on arXiv 27 Aug 2025.

## Builds on (most cited first, 30 of 31)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Understanding Black-box Predictions via Influence Functions](https://mlsanity.com/paper/1703.04730/) (2017)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Learning to Prompt for Continual Learning](https://mlsanity.com/paper/2112.08654/) (2021)
- [A Survey on In-context Learning](https://mlsanity.com/paper/2301.00234/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [A Survey on Data Selection for Language Models](https://mlsanity.com/paper/2402.16827/) (2024)
- [Efficiently Identifying Task Groupings for Multi-Task Learning](https://mlsanity.com/paper/2109.04617/) (2021)
- [Rethinking with Retrieval: Faithful Large Language Model Inference](https://mlsanity.com/paper/2301.00303/) (2022)
- [In-Context Learning with Long-Context Models: An In-Depth Exploration](https://mlsanity.com/paper/2405.00200/) (2024)
- [Understanding and Improving Information Transfer in Multi-Task Learning](https://mlsanity.com/paper/2005.00944/) (2020)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [Revisiting Demonstration Selection Strategies in In-Context Learning](https://mlsanity.com/paper/2401.12087/) (2024)
- [Robust Fine-Tuning of Deep Neural Networks with Hessian-based Generalization Guarantees](https://mlsanity.com/paper/2206.02659/) (2022)
- [Parallel Structures in Pre-training Data Yield In-Context Learning](https://mlsanity.com/paper/2402.12530/) (2024)
- [From Few to Many: Self-Improving Many-Shot Reasoners Through Iterative Optimization and Generation](https://mlsanity.com/paper/2502.00330/) (2025)
- [Scalable Multitask Learning Using Gradient-based Estimation of Task Affinity](https://mlsanity.com/paper/2409.06091/) (2024)
- [Boosting Multitask Learning on Graphs through Higher-Order Task Affinities](https://mlsanity.com/paper/2306.14009/) (2023)
- [Noise Stability Optimization for Finding Flat Minima: A Hessian-based Regularization Approach](https://mlsanity.com/paper/2306.08553/) (2023)
- [Scalable Fine-tuning from Multiple Data Sources: A First-Order Approximation Approach](https://mlsanity.com/paper/2409.19458/) (2024)
- [Toward In-Context Teaching: Adapting Examples to Students' Misconceptions](https://mlsanity.com/paper/2405.04495/) (2024)
- [Efficient Ensemble for Fine-tuning Language Models on Multiple Datasets](https://mlsanity.com/paper/2505.21930/) (2025)
- [Divide, Reweight, and Conquer: A Logit Arithmetic Approach for In-Context Learning](https://mlsanity.com/paper/2410.10074/) (2024)
- [Exploring Demonstration Ensembling for In-context Learning](https://mlsanity.com/paper/2308.08780/) (2023)
- [Compute-Optimal LLMs Provably Generalize Better With Scale](https://mlsanity.com/paper/2504.15208/) (2025)

## Built on by (most cited first, 8 of 8)

- [Scalable Multi-Objective and Meta Reinforcement Learning via Gradient Estimation](https://mlsanity.com/paper/2511.12779/) (2025)
- [Efficient Estimation of Kernel Surrogate Models for Task Attribution](https://mlsanity.com/paper/2602.03783/) (2026)
- [WinQ: Accelerating Quantization-Aware Training of Language Models Around Saddle Points](https://mlsanity.com/paper/2605.17471/) (2026)
- [Efficiently Learning Branching Networks for Multitask Algorithmic Reasoning](https://mlsanity.com/paper/2512.01113/) (2025)
- [Learning Multimodal Embeddings for Traffic Accident Prediction and Causal Estimation](https://mlsanity.com/paper/2512.02920/) (2025)
- [MoPLEx: Estimating Plackett-Luce Mixture Models for Multi-Objective Alignment](https://mlsanity.com/paper/2608.25200/) (2026)
- [Long-Context Demonstration Selection Using State Space Models](https://mlsanity.com/paper/2609.17888/) (2026)
- [Efficient Task Adaptation in Large Language Models: A Survey of Weight-Based, Prompt-Based, and Embedding-Based Adaptations](https://mlsanity.com/paper/2610.00928/) (2026)

## Abstract (arXiv)

This paper introduces an algorithm to select demonstration examples for in-context learning of a query set. Given a set of n examples, how can we quickly select k out of n to best serve as the conditioning for downstream inference? This problem has broad applications in prompt tuning and chain-of-thought reasoning. Since model weights remain fixed during in-context learning, previous work has sought to design methods based on the similarity of token embeddings. This work proposes a new approach based on gradients of the output taken in the input embedding space. Our approach estimates model outputs through a first-order approximation using the gradients. Then, we apply this estimation to multiple randomly sampled subsets. Finally, we aggregate the sampled subset outcomes to form an influence score for each demonstration, and select k most relevant examples. This procedure only requires pre-computing model outputs and gradients once, resulting in a linear-time algorithm relative to model and training set sizes. Extensive experiments across various models and datasets validate the efficiency of our approach. We show that the gradient estimation procedure yields approximations of full inference with less than 1% error across six datasets. This allows us to scale up subset selection that would otherwise run full inference by up to 37.7× on models with up to 34 billion parameters, and outperform existing selection methods based on input embeddings by 11% on average.