Not All Demonstration Examples are Equally Beneficial: Reweighting Demonstration Examples for In-Context Learning (2023)

Abstract: Large Language Models (LLMs) have recently gained the In-Context Learning (ICL) ability with the models scaling up, allowing them to quickly adapt to downstream tasks with only a few demonstration examples prepended in the input sequence. Nonetheless, the current practice of ICL treats all demonstration examples equally, which still warrants improvement, as the quality of examples is usually uneven. In this paper, we investigate how to determine approximately optimal weights for demonstration examples and how to apply them during ICL. To assess the quality of weights in the absence of additional validation data, we design a masked self-prediction (MSP) score that exhibits a strong correlation with the final ICL performance. To expedite the weight-searching process, we discretize the continuous weight space and adopt beam search. With approximately optimal weights obtained, we further propose two strategies to apply them to demonstrations at different model positions. Experimental results on 8 text classification tasks show that our approach outperforms conventional ICL by a large margin. Our code are publicly available at https:github.com/Zhe-Young/WICL.

## introduction with the increase in model size and training corpus, large language models (llms) have demonstrated in-context learning (icl) capabilities @xcite @xcite . unlike traditional finetuning, which requires updating model parameters, icl allows llms to adapt to downstream tasks while keeping the parameters fixed. to enable icl, a small number of examples will be prepended before the query to form a prompt. the prompt is then fed into the llm for prediction. numerous studies have shown the effectiveness of icl, demonstrating that it can achieve, and sometimes even surpass, the performance of fine-tuned models. the performance of icl is closely related to the quality of demonstration examples @xcite . therefore, various methods have been proposed to select high-quality examples. for example, @xcite selects nearest examples based on semantic similarity, @xcite employs reinforcement learning for example selection, and @xcite maximizes mutual information on an unlabeled validation set for choosing templates. however, these methods assume the availability of a large number of examples for selection, which is not the case in real fewshot settings. from another aspect, previous icl methods treat all demonstration examples equally, which can be further improved, as the quality of demonstrations is not even. as shown in figure 1 , assigning different sets of weights to demonstration examples has a significant impact on icl performance. for a 13b gpt model, in terms of the accuracy on the sst2 dataset, the best weights surpass the worst by 40 points and surpass the non-13209 $ffxudf\ @xmath0.73 :hljkw9hfwru 5hjuhvvlrq/lqh weighting strategy by 15 points. this highlights the necessity of assigning appropriate weights to demonstration examples. in this paper, we propose the weighted incontext learning (wicl) method to enhance icl by determining and applying approximately optimal weights to demonstration examples. there are two key challenges for determining weights: how to evaluate the quality of a set of weights without additional validation data (weight evaluation), and how to efficiently search the best set of weights in the continuous space (weight search). for weight evaluation, we introduce a masked self prediction (msp) score as a proxy metric, which can be computed solely on demonstration examples but is strongly correlated with the final icl performance, as shown in figure 2 . for weight search, we discretize the weight space and employ beam search to efficiently find the weights that yield the highest msp score. with approximately optimal weights discovered, we propose two strategies to apply them to demonstrations. scaling key matrix (skm) applies the weights to the attention key matrices of demonstration examples, while scaling attention weights (saw) directly adjusts the attention weights. both strategies adjust the importance of demonstrations by influencing the attention weights according to a set of weights. we evaluate wicl on 8 text classification tasks and show that it achieves substantial accuracy improvement over the conventional icl method that follows a non-weighting strategy. in addition, our approach exhibits robustness to varying shot num-bers and templates. as for the approximately optimal weights discovered by our approach, we find that they demonstrate a close performance to the global optimal weights. furthermore, we discover that our example reweighting approach mainly works at middle layers, and reweighting only in a few middle layers even outperforms full layers. we summarize our contributions as follows: ## problem formulation for a text classification task, we are given a set of k demonstration examples denoted by @xmath1, y 1 ), (x 2 , y 2 ), ..., (x k , y k )}, where x i and y i denote the input text and its corresponding label, respectively. we define a transformation function t that maps each sample (x, y) to an icl example t (x, y) (e.g. t (x, y) = "sentence: x sentiment: y" for the sst2 dataset). the concatenation of these transformed examples is then defined as the final demonstration c: in the above equation, θ and d represent the parameters of llm and the candidate label set, softmax self-attention layer key×query attention weights × values weight vector new example key query attention weights weight vector value norm scaling key matrix scaling attention weights llm ffn self attention … … … option1: option2: or 1.2 1.0 0.8 1.0 example1 example2 example4 example3 att1 att2 att4 att3 weight vector 1.2 1.0 0.8 1.0 1.2 1.0 0.8 1.0 example1 example2 example4 example3 ⊙ icl example reweighting discovering optimal weights masked selfprediction score indicator ＆beam search ⊗ ⊙ ⊗ respectively. conventional icl does not take the weight of each example into account, and llm almost treats each example equally. to address this limitation, we introduce a weight vector @xmath2, w 2 , ..., w k ) for the demonstration, where w i denotes the weight of the i th example t (x i , y i ). a higher value of w i indicates that the example t (x i , y i ) is more important, and thus llm should pay more attention to it. when given a specific demonstration, the objective of wicl is to fully use demonstration examples and identify an appropriate weight vector w that enables the model to achieve performant performance. ## methodology demonstration example reweighting faces two challenges: (1) how to add weights to demonstration examples when given a weight vector. (2) how to find a performant weight vector in the absence of a held-out validation set. section 3.1 aims to solve (1) by reweighting at the self-attention layer, and section 3.2 aims to address (2) through weight quantization and beam search with an indicator. our approach is illustrated in figure 3 . ## experiments similar to previous work @xcite , we use 8 classical text classification datasets, respectively sst2 @xcite , cr @xcite , mr @xcite , subj @xcite , trec (voorhees and tice, 2000), dbpedia @xcite , agnews @xcite and rte @xcite , to evaluate the efficacy of our demonstration examples reweighting approach. in our experiments, we utilize 5 different gpt-like causal language models released by fairseq @xcite , and the number of parameters of these models is 355m, 1.3b, 2.7b, 6.7b, 13b, respectively. 1 previous works @xcite @xcite have 355m 1.3b 2.7b avg icl baseline 72.1 64.9 68.1 68.4 wicl-dual 80.8 76.1 77.7 78.2 wicl-skm 80.9 76.5 79.1 78.8 as the experiment results shown in ## analysis we also further analyze our experiments and investigate properties of our approach. since skm outperforms saw, we do our analysis experiments only with skm reweighting strategy. firstly, we compare different example masking/removing strategies when calculating msp (section 5.1), and compare msp with held-out validation set (section 5.2). nextly, we explore the robustness of our approach to different templates, shot numbers (section 5.3, 5.4). then, we find that our approach can obtain near-optimal weights (section 5.5). in section 5.6, we discover that example reweighting mainly works at middle layers, and reweighting on only a few middle layers can even outperform full-layer reweighting. section 5.7 shows some empirical results on relationship between example weight and example quality/position. ## conclusion in this paper, we find that treating demonstration examples equally can be a bottleneck for icl, and assigning proper weights to examples can significantly improve performance. based on this observation, we attempt to reweight demonstration examples and propose weighted in-context learning. to explicitly assign weights, we reweight examples in the self-attention layer by scaling key matrix or scaling attention weights; to determine the weight vector for examples, we adopt weight quantization beam search by maximizing masked self-prediction score, which can be obtained with any held-out validation set yet shows a strong correlation with the final performance. our approach can significantly improve icl performance in true few-shot setting, and experiments on 8 nlp tasks also demonstrate the efficacy of our approach.