Linear-Time Demonstration Selection for In-Context Learning via Gradient Estimation
This work addresses the challenge of prompt tuning and chain-of-thought reasoning for users of large language models by providing a faster and more accurate demonstration selection method, though it is incremental as it builds on prior similarity-based approaches.
The paper tackles the problem of efficiently selecting demonstration examples for in-context learning by introducing a gradient-based algorithm that estimates model outputs, achieving less than 1% error and scaling selection up to 37.7 times faster while outperforming existing methods by 11% on average.
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}\times$ on models with up to $34$ billion parameters, and outperform existing selection methods based on input embeddings by ${11}\%$ on average.