In-Context Demonstration Selection with Cross Entropy Difference (2023)

Abstract: Large language models (LLMs) can use in-context demonstrations to improve performance on zero-shot tasks. However, selecting the best in-context examples is challenging because model performance can vary widely depending on the selected examples. We present a cross-entropy difference (CED) method for selecting in-context demonstrations. Our method is based on the observation that the effectiveness of in-context demonstrations negatively correlates with the perplexity of the test example by a language model that was finetuned on that demonstration. We utilize parameter efficient finetuning to train small models on training data that are used for computing the cross-entropy difference between a test example and every candidate in-context demonstration. This metric is used to rank and select in-context demonstrations independently for each test input. We evaluate our method on a mix-domain dataset that combines 8 benchmarks, representing 4 text generation tasks, showing that CED for in-context demonstration selection can improve performance for a variety of LLMs over baseline selection methods.

## introduction large language models (llms) have been widely successful across many nlp tasks @xcite @xcite . the primary method for llms to adapt to new tasks has been using in-context learning, where a few examples and labels are provided as input to the model @xcite . this simple approach has shown large improvements over zero shot settings, and even outperformed finetuning methods when the training dataset is small. however, the model's performance can be greatly influenced by which in-context demonstrations (icds) are selected into the prompt @xcite @xcite . selecting the best in-context demonstrations can be challenging. the variance in performance of similar demonstrations can be large, and the selected examples can introduce unfavorable prior biases on the output label space @xcite . the naive approach is to randomly sample demonstrations from the same source dataset. previous methods for selecting icds include simple methods such as selecting nearest neighbors by embedding distance @xcite and retrieval-based methods that require training a retriever model @xcite . this work presents a new method of selecting demonstrations that can be applied to any sized training data, requires training small peft models only and outperforms the nearest neighbor baseline on gpt-davinci-003 @xcite . we propose a cross entropy difference (ced) method for icd selection. ced has been used to select in-domain data from large mixed-domain datasets for domain adaptation @xcite @xcite . we borrow this idea to conduct icd selection. specifically, we utilize parameter efficient finetuning to train small models on training data that are used for computing the ced between a test example and every candidate in-context demonstration. the ced scores are used to rank and select in-context demonstrations. we present a theoretical explanation for the effectiveness of ced. ced approximates the gradient alignment between training and test examples. our analysis builds on previous findings that demonstrations operate as "meta-gradients" and shows that demonstrations with gradients similar to those of test inputs lead to improved performance in downstream tasks @xcite . we evaluate our proposed ced-icd selection method on a mixed-domain dataset composed of 8 datasets on 4 tasks: binary classification, multiple choice and extractive question answering and ab-stractive question answering. we show that downstream model performance using ced-icd outperforms nearest neighbor baselines and transfers across models allowing training small models for selection but evaluating test examples on much larger models including gpt-3.5. the contributions of this work are the following: 1) we present a method for selecting in-context demonstrations based on cross entropy difference. 2) we provide theoretical guidance for why selecting demonstrations based on their gradient alignment with test example is an effective heuristic. 3) we evaluate our method on a mixed-domain benchmark composed of 8 datasets from 4 tasks. 4) we evaluate our method on different sizes of gpt3 models, showing that this method transfers to larger models leading to performance improvements over baseline selection methods. ## related work our work combines ideas from three bodies of research: in-context learning (icl), data selection for domain adaptation and parameter efficient finetuning (peft). while in-context learning @xcite ) has shown very strong results for few-shot settings, recent work has shown that llms are very sensitive to the selected examples leading to large variance in performance @xcite , sensitivity to the order of examples @xcite and even lack of sensitivity to the actual labels @xcite . other work has attempted to mitigate these challenges by selecting in-context demonstrations by the nearest neighbor examples in embedding space @xcite , or training a retrieval mechanism @xcite . we build on this line of work by proposing a novel selection method that combines the observations from @xcite that domain similarity is a key characteristic of good in-context demonstrations, and observations from @xcite that using perplexity could be a good heuristic for prompt selection. previous work on domain adaptation has focused on finding in-domain examples from a large outof-domain dataset to train a model that achieves a better generalization on a target distribution @xcite @xcite . data selection is intended to maximize the distributional similarity between a training dataset and test dataset. however, cross entropy dif-ference has not been used previously at the example granularity to rank the "in-domainness" of training data in reference to just a single target example. we propose a natural extension of this framework for selecting demonstrations that are "in-domain" for a test input, which we demonstrate is an effective metric for selecting demonstrations for in-context learning. parameter efficient finetuning (peft) proposes a class of methods for augmenting model parameters with a small number of additional parameters that can be trained and stored efficiently @xcite @xcite . however, peft is usually used independently from in-context learning. @xcite report that combining in-context learning and peft has not been effective. @xcite does report some settings where peft and icl can be combined, but only under specific task conditions. we report similar findings, that in-context demonstrations do not improve peft models when selected randomly, however, we do see improvements in peft performance when combined with ced for selecting in-context demonstrations during both training and inference time. we also utilize the ability of a peft model, t-few @xcite , to train on very few examples to be able to effectively compute ced scores without overfitting to the target domain. ## methodology we propose a method for selecting in-context demonstrations (icds) by finding the training data that would minimize the perplexity of a test example, if a language model were finetuned on that training example. this approach stems from previous findings that in-context examples may act as a type of meta-gradient on the frozen llm @xcite and the assumption that models perform better on in-domain test data. as we show in the following sections, our method of using cross entropy difference finds the demonstrations that appear most likely to be from the same domain as the test example. ## experiments we evaluate ced-icd selection on both small models and the transfer of the selection method to larger models, including gpt-davinci-003. we evaluate the selection method in a mixed-domain setting where random demonstrations are not trivially indomain. we do not provide task or dataset labels as input to the selection model. we show in section 5.2, both ced and the nearest neighbors baseline do not exclusively select in-domain demonstrations in the mixed domain setting. in fact, we show in section 5.2 that out-of-domain examples may also be strong in-context demonstrations. in practical settings, a single llm may be used for multiple tasks and there may not be labels for the task type, especially with the more common chat interfaces. we find this mixed-domain setting to better reflect these realistic challenges. ## analysis we analyze the quality of cross entropy difference selection by computing the rank of selected demonstrations compared to an oracle. we also explore the presence of strong in-context demonstrations selected from out-of-domain data, compared to indomain data. ## conclusion this work shows that cross entropy difference can be used as a heuristic for selecting in-context demonstrations for in-context learning. we motivate our approach by linking previous observations that in-context demonstrations operate as metagradients on a frozen large language model to the data selection literature which has shown that ced is a method that selects examples that have a gradient most similar to the in-domain examples. we empirically show that we can use smaller models to compute ced scores and that this selection method effectively transfers to large models, such as the 175 billion parameter gpt3, and is able to improve performance over baseline selection methods.