Data Curation Alone Can Stabilize In-context Learning (2023)

Abstract: In-context learning (ICL) enables large language models (LLMs) to perform new tasks by prompting them with a sequence of training examples. However, it is known that ICL is very sensitive to the choice of training examples: randomly sampling examples from a training set leads to high variance in performance. In this paper, we show that carefully curating a subset of training data greatly stabilizes ICL performance without any other changes to the ICL algorithm (e.g., prompt retrieval or calibration). We introduce two methods to choose training subsets—both score training examples individually, then select the highest-scoring ones. CondAcc scores a training example by its average dev-set ICL accuracy when combined with random training examples, while Datamodels learns linear regressors that estimate how the presence of each training example influences LLM outputs. Across five tasks and two LLMs, sampling from stable subsets selected by CondAcc and Datamodels improves average accuracy over sampling from the entire training set by 7.7% and 6.3%, respectively. Surprisingly, the stable subset examples are not especially diverse in content or low in perplexity, in contrast with other work suggesting that diversity and perplexity are important when prompting LLMs.

## introduction in-context learning (icl) is a new paradigm for few-shot learning with pretrained large language models (llms) without any parameter updates. in icl, an llm can perform a new task simply by conditioning on a prompt foot_0 consisting of a sequence of labeled training examples. first introduced by gpt-3 @xcite , icl with llms has reached state-of-the-art few-shot performance across many tasks @xcite 2022; @xcite . compared with alternatives that use finetuning @xcite @xcite , icl does not require taskspecific training, which enables its use with very large language models, and it uses a unified model for all tasks, enabling easier deployment. despite its impressive few-shot performance, icl often exhibits unintuitive behavior @xcite . the standard icl approach is to randomly sample a few examples from a training set to construct a prompt @xcite ; however, prior work @xcite @xcite has found that icl is very sensitive to the choice of training examples and their order in the prompt. icl is also sensitive to small changes in prompt format @xcite . in this paper, we show that carefully curating a smaller training dataset from a larger pool can make icl much more stable. we define a training subset e to be stable if randomly sampling a sequence of examples as a prompt from e yields much higher average and worst-case accuracy than randomly sampling from the original training set. we propose two methods to identify such a sta-ble subset. our condacc method scores a training example by its average dev-set icl accuracy when combined with random training examples; these scores are closely related to data shapley values @xcite . our datamod-els method fits a linear regressor that predicts the llm's output based on which example is present at each index in the prompt; we score a training example highly if the associated weights from the linear model indicate that its presence improves accuracy. for both methods, we then select training examples with the highest scores to form the stable subset. while some prior work improves icl accuracy by retrieving a suitable prompt for each test example @xcite @xcite , we show that it is possible to achieve stably good accuracy with a randomly sampled prompt for all test examples, when given the "right" training (sub)set. our subset selection methods greatly improve performance across 5 classification datasets and 4 llms, with main experiments on gptj-6b @xcite and opt-13b @xcite . on average, condacc and data-models outperform the baseline that uses the entire training set without selection (named all) by 7.7% and 6.3%, respectively, when comparing the average accuracy over multiple sampled prompts. in contrast, baselines that choose examples found in high-performing prompts (topprompts) or examples that lead to high one-shot icl accuracy (oneshot) to form the subsets do not perform as well (see figure 1 ). our stable subset examples generalize to out-of-distribution test data, and we can even find stable subsets for binary classification tasks that only contain examples of one label; these findings suggest that the stable subset examples help llms understand the overall task definition. finally, we study what makes stable subset examples special by analyzing sequence length, perplexity, and diversity in both raw text and embedding spaces. we find that these examples do not have abnormally long sequence lengths or high perplexities. in contrast with prior work optimizing diversity for prompt selection @xcite , we find our stable subsets no more diverse than random subsets of the training data. in summary, we show that curating training data appropriately leads to more stable and accurate icl performance; we hope future work can develop new strategies for writing such helpful ex-amples. code and data are publicly available at https://github.com/terarachang/dataicl . ## problem setups we use the original icl formulation proposed by gpt-3, also known as the direct method, for all our experiments. specifically, an llm performs in-context learning on a new task based on a taskspecific prompt z formed by concatenating k labeled training examples, i.e., @xmath0, ..., z k ], where each z j is a training example (x, y) consisting of an input x and label y. the llm then makes predictions on a test input x test conditioned on the prompt z followed by x test , denoted by arg max @xmath1, x test ), where c is the set of possible labels. given a training set d tr , a dev set d dev , a heldout test set d test , and a predefined number of shots k, our goal is to select a stable training subset e ⊂ d tr , |e| > k, such that randomly sampling a sequence of k examples from e to form a prompt generally yields good performance on d test . we propose two setups, labeled and unlabeled. in labeled, our goal is to study which training examples consistently lead to high icl accuracy. we assume access to a large labeled d l @xmath2, y i )} ntr @xmath3, and a small labeled d dev . unlabeled is closer to the true fewshot learning setup @xcite , where we only have access to d dev and a large unlabeled training set {x i } ntr @xmath4. we pair each input x i with every possible label @xmath5, ỹ)|ỹ ∈ c}} ntr @xmath6. in both setups, our goal is to select a subset of e training examples from either d l tr or d u tr . note that we may select examples with incorrect labels under unlabeled. we only use the large labeled d test for evaluating our selection methods. ## methods we propose the following steps to identify a stable subset of e training examples: ## analysis we further analyze what makes the selected training examples special along different dimensions: sequence length, perplexity, and diversity. we compare good with bad training examples, where we use our condacc method to identify a bad (resp., good) subset by selecting e ′ training examples with the lowest (resp., highest) importance scores in each class ( §3.3). please refer to ## conclusion we propose two methods to identify stable training subsets for in-context learning, achieving substantially higher average and worst-case accuracy across different setups. our condacc method is intuitive and easy-to-implement, while our data-models method provides informative weights that enable further analyses. the success of our methods implies that when provided with the "right" training set (in our case, a subset of 20 examples to randomly sample prompts from), icl could be far less sensitive to the choice of training examples and their orders in a prompt. our analyses on stable subsets find that they do not contain outliers with especially long sequence lengths or high perplexities, and are also no more diverse than random subsets of the training data. we hope our work is a step towards developing guidelines for finding or writing better training examples.