In-Context Learning with Iterative Demonstration Selection (2024)

Abstract: Spurred by advancements in scale, large language models (LLMs) have demonstrated strong few-shot learning ability via in-context learning (ICL). However, the performance of ICL has been shown to be highly sensitive to the selection of few-shot demonstrations. Selecting the most suitable examples as context remains an ongoing challenge and an open problem. Existing literature has highlighted the importance of selecting examples that are diverse or semantically similar to the test sample while ignoring the fact that the optimal selection dimension, i.e., diversity or similarity, is task-specific. Based on how the test sample is answered, we propose Iterative Demonstration Selection (IDS) to leverage the merits of both dimensions. Using zero-shot chain-of-thought reasoning (Zero-shot-CoT), IDS iteratively selects examples that are diverse but still strongly correlated with the test sample as ICL demonstrations. Specifically, IDS applies Zero-shot-CoT to the test sample before demonstration selection. The output reasoning path is then used to choose demonstrations that are prepended to the test sample for inference. The generated answer is followed by its corresponding reasoning path for extracting a new set of demonstrations in the next iteration. After several iterations, IDS adopts majority voting to obtain the final result. Through extensive experiments on tasks including reasoning, question answering, and topic classification, we demonstrate that IDS can consistently outperform existing ICL demonstration selection methods.

## introduction with the recent advancements in scaling up model parameters, large language models (llms) showcase promising results on a variety of few-shot tasks through in-context learning (icl), where the model is expected to directly generate the output of the test sample without updating parameters. this is achieved by conditioning on a manually designed prompt consisting of an optional task description and a few demonstration examples @xcite . fig. 1 shows an example describing how llms perform icl on the topic classification task. given a few text-topic pairs as demonstrations, icl combines them with the test sample as input, to the llm for inference. the output, i.e., 'technology', is generated by the model autoregressively without any parameter updates. despite the effectiveness, the performance of icl has been shown to be highly sensitive to the selection of demonstration examples @xcite . different sets of demonstrations can yield performance ranging from nearly random to comparable with state-of-the-art models @xcite . to alleviate the above issue, researchers in icl have proposed a number of methods to select a set of examples as few-shot demonstrations @xcite @xcite @xcite @xcite . however, for llms for which parameters or detailed output distributions are not available @xcite , it is still a common practice to randomly select examples or select examples that are semantically similar to the test sample as demonstrations, i.e., considering diversity or similarity. while several approaches investigate the combination of similarity and diversity when prompting with explanations, exploring compositional generalization, or choosing examples for annotation @xcite @xcite , it is not yet clear how to determine and leverage the optimal dimension for different tasks in icl and how the rationale for answering the query benefits the balance between these two dimensions. actually, the optimal dimension for selecting demonstration examples is task-specific. as we will show in §4, the diversity dimension is superior to the similarity dimension on commonsenseqa while the similarity dimension outperforms the diversity dimension on agnews and boolq. thus, it is unreasonable to claim that one dimension is consistently better than the other across different tasks. to fully leverage the merits of both dimensions, we propose iterative demonstration selection (ids) for icl (fig. 2 ) by utilizing how the test sample is answered. ids can iteratively select demonstration examples that are diverse but still have a strong correlation with the test sample through zero-shot chain-of-thought reasoning (zero-shot-cot) @xcite . specifically, zero-shot-cot, e.g., "let's think step by step.", is first applied to the test sample before selecting demonstrations to obtain a reasoning path. the training examples that are most semantically similar to the generated reasoning path are then selected as demonstrations. they are prepended to the test sample for inference. note that ids ensures that the generated answer is accompanied by the reasoning path through designed prompts. the new reasoning path is then used for extracting another set of demonstration examples by semantic similarity in the next iteration. after a few iterations, ids adopts majority voting to obtain the final result. empirical results on tasks spanning mathematical reasoning, commonsense reasoning, logical reasoning, question answering, and topic classification show that ids can consistently outperform previous icl demonstration selection baselines. in summary, our main contributions are: this work mainly explores how to select few-shot in-context learning demonstrations for llms by leveraging zero-shot-cot. in light of this, we review four lines of research that form the basis of this work: few-shot learning, in-context learning basics, demonstration selection for in-context learning, and chain-of-thought reasoning. ## problem formulation given the test set d test and the training set d train , the goal of icl demonstration selection is to find an optimal subset @xmath0, y 1 ), ..., (x k , y k )} (kshot) of d train as demonstration examples for each test sample (x i , ŷi ) to maximize the overall task performance on d test . more formally, the optimal selection method h is defined as: as demonstrated in previous work @xcite , the overall task performance is highly sensitive to the selection method h. different sets of demonstration examples can yield significantly different performance. for example, @xcite show that the minimum and maximum icl performance due to random sampling differs @xmath1, which emphasizes the importance of selecting good demonstrations for llms. a natural question is: what makes good incontext demonstrations? for llms, it is still a common practice to select a subset s consisting of examples that are diverse or semantically similar to the test sample as demonstrations, i.e., considering the diversity or similarity of s. to investigate whether one dimension is consistently better than the other one across different tasks, we conduct some pilot experiments on commonsenseqa @xcite , boolq @xcite and agnews @xcite . specifically, we randomly sample 100 examples from the original test set for experiments and conduct 4-shot learning using gpt-3.5 (gpt-3.5-turbo). following @xcite , we use sentence-bert @xcite to encode all samples. for each test sample, the similar-icl method selects the top-4 similar training data based on cosine similarity while the random-icl method randomly samples 4 training examples as few-shot demonstrations. inspired by @xcite , we apply self-consistency with 3 decoding paths (temperature 0.7) to similar-icl (named similar-icl-consistency) and run random-icl 3 times before majority voting (named random-icl-voting) to improve the robustness. the results of different methods on four datasets are reported in intuitively, semantically similar examples can help the model correctly answer the test query as they might share similar input-output patterns with the test sample which could unleash gpt-3.5's power of text generation. to further understand why the similarity dimension underperforms the diversity dimension on commonsenseqa, we present a case study in considering the strengths and weaknesses of both dimensions, we aim to design a method that can select demonstration examples that are diverse (minimizing misleading information) but still strongly correlated with the test sample, which is introduced in the next section. ## iterative demonstration selection based on the observations and considerations in §4, we introduce iterative demonstration selection (ids) for icl demonstration selection by leveraging how the test sample is answered (see fig. 2 for an illustration). intuitively, the demonstrations that are similar to the reason for answering a sample are strongly correlated with this sample. therefore, we propose to incorporate zero-shot chain-of-thought reasoning (zero-shot-cot) into ids to iteratively select demonstration examples that are diverse but still have a strong correlation with the test sample. specifically, for each test sample xi , ids mainly consists of four steps: similar-icl-consistency random-icl-voting which choice is the correct answer to the question? which choice is the correct answer to the question? examples: question: if you have cleaned off dust here it may be difficult to do your homework where? answer choices: (a) desktop (b) closet (c) most buildings (d) surface of earth (e) stove answer: a question: where is dust likely to be under? answer choices: (a) closet (b) ground (c) windowsill (d) attic (e) carpet answer: e question: where would you find a dustbin that is being used? answer choices: (a) utility closet (b) ground (c) cupboard (d) broom closet (e) kitchen answer: e question: dust accumulates where? answer choices: (a) ceiling (b) library (c) surface of earth (d) most buildings (e) desktop answer: d examples: question: she had a busy schedule, she had to run errands and pick up the kids the second she did what? answer choices: (a) make time for (b) take money (c) go outdoors (d) leave work (e) field answer: d question: what is the worst outcome of an injury? answer choices: (a) cause death (b) cause bleeding (c) falling down (d) become infected (e) claim insurance answer: a question: mom said that sarah should stay in bed until she was able to go to school again. what did mom say to sarah when she tried to get up? answer choices: (a) you're sick (b) were sick (c) more rest (d) rest more (e) get back under the covers answer: a question: john got a raise, but he lost rank. overall, it was a good what? answer choices: (a) demotion (b) push down (c) go off strike (d) lower (e) go off strike answer: a the response should follow the format: answer: {a, b, c, d or e} the response should follow the format: answer: {a, b, c, d or e} here is the test data. here is the test data. question: john wanted to clean all of the dust out of his place before settling down to watch his favorite shows. what might he hardest do dust? answer choices: (a) closet (b) under the bed (c) television (d) attic (e) most buildings question: john wanted to clean all of the dust out of his place before settling down to watch his favorite shows. what might he hardest do dust? answer choices: (a) closet (b) under the bed (c) television (d) attic (e) most buildings answer: e ✗ answer: d ✓ r and training examples to obtain the contextual representations and use cosine similarity to measure the similarity between representations. ## experiments in this section, we first describe the tasks and datasets, and then introduce methods compared in our work. finally, we present the experimental results.