We study the problem of demonstration selection, which involves selecting a subset of examples for pre-pending to a query to a language model. This problem is closely related to in-context learning and model inference. Since the inference cost of a transformer model scales quadratically with sequence length, the selection problem becomes especially challenging in a long-context scenario. In this paper, we tackle this problem by building on state space models (SSMs), which has linear inference time given the input embeddings. Our approach works in two stages. First, we learn a small set of SSMs through a distillation procedure — given a (trained) transformer model, we partition all the layers into consecutive (layer) groups — for each group, we learn one SSM to mimic the input-output behavior, thereby capturing structural similarity in adjacent layers. Second, we map the SSM outputs to a small set of tokens, and apply them for downstream demonstration selection. We run extensive experiments in both synthetic and real-world datasets to validate this approach. We show that the SSMs trained from distillation incur a relative approximation error of less than 0.7%. As for downstream demonstration selection, we find that on several text classification and reasoning tasks, our approach reduces FLOPs by 14.2× and improves accuracy by 6.48% relative to baseline selection methods.

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 𝟏% 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.