Effective In-Context Example Selection through Data Compression (2024)

Abstract: In-context learning has been extensively validated in large language models. However, the mechanism and selection strategy for incontext example selection, which is a crucial ingredient in this approach, lacks systematic and in-depth research. In this paper, we propose a data compression approach to the selection of in-context examples. We introduce a two-stage method that can effectively choose relevant examples and retain sufficient information about the training dataset within the in-context examples. Our method shows a significant improvement of an average of 5.90% across five different real-world datasets using four language models.

heading: 1 Introduction
text: Drawing inspiration from recent research that regards Large Language Models (LLMs) as an efficient means of compressing pre-training datasets, and the notion that In-Context Learning (ICL) can be seen as fine-tuning on example datasets (Dai et al., 2022), we assume that LLMs can achieve data re-compression through ICL. In other words, an effective training dataset compression method can aid in the selection of in-context examples. Looking at the matter from another perspective, it is evident that fine-tuning the entire dataset would yield the best results. However, in the case of ICL, we typically choose only a few examples as LLM prompts due to the limitations of input window length. By employing data compression techniques, we can ensure that the majority of data information is preserved in the in-context examples, which is also the aim of dataset pruning.
Based on the aforementioned analysis, we propose to utilize the influence function (Koh and
⋆ Equal contribution. †Corresponding author. Work partially done at Engineering Research Center of Next- Generation Intelligent Search and Recommendation, Ministry of Education.
Liang, 2017), which has exhibited efficacy in dataset pruning, to select examples for ICL. However, recent studies on ICL have revealed that the relevance between the query source and in-context examples is critical for ICL. Furthermore, the influence function requires the gradient of parameters, which is computationally expensive and inefficient. To tackle the aforementioned issues, we suggest a two-stage method (I2CL). Firstly, relevant examples for the query input are recalled, which ensures the correlation between the examples and the query source. Secondly, our meta-gradient-based influence function is utilized to calculate the influence score for each recalled example. Finally, based on the influence score, in-context examples are selected from the recalled examples. Notably, our framework compresses important information from the training set into the in-context examples, thereby enhancing the performance of ICL. Additionally, our framework is data-independent, relies solely on a small number of model parameters, and does not require the training of any additional models. Numerous experiments indicate that I2CL shows a significant improvement of an average of 5.90% on five different real-world datasets using multiple language models.

heading: 2 Background

heading: 2.1 In-Context Learning
text: The ICL scenario of LLMs can be regarded as a conditional text generation problem. Concretely, the probability of generating a target y is conditioned on the context C, which includes k examples and the source x. Therefore, the probability can be expressed as:
pLLM(y | C, x) = T∏
t=1
p (yt | C, x, y<t)
where LLM denotes the parameters of the large language model, and C =
1 871
{x1, y1, x2, y2, . . . , xk, yk} is a context string concatenating k training instances. For example, (xk, yk) is concatenated with the special character, e.g., “\n” or “Sentence: x; Sentiment y.” which is denoted as pk. In this paper, we have different example sets C at different stages, C1 in the first stage, and C2 in the second stage, where C2 is a subset of C1.
Dai et al. (2022) explains language models as meta-optimizers and understands ICL as a kind of implicit finetuning:
F̃ICL(q) = WZSLq+ ∑
i
( WV x ′ i ⊗ ( WKx ′ i )T) q
=WZSLq+∆WICLq =(WZSL +∆WICL)q,
(1)
where ZSL denotes the zero-shot learning, which only contains the source x; x ∈ Rd is the input representation of a query token t, and q = WQx ∈ Rd′ is the attention query vector; x′ ∈ Rd denotes the input representations of the example’s token; WQ,WK ,WV ∈ Rd
′×d are the projection matrices for computing the attention queries, keys, and values, respectively. Dai et al. (2022) regards WV X ′ as some meta-gradients, which are used to compute the updated matrix ∆WICL.

heading: 2.2 Dataset Pruning
text: Investigating the data redundant problem not only helps to improve the training efficiency but also helps us understand the representation ability of small data and how many training samples are required and sufficient for a learning system. (Yang et al., 2022) proposed to use the Influence Function to accurately and fast estimate the parameter change caused by weighting an example p for the training dataset. The influence of weighting p on the parameters is given by:
Iparam(p) = dθ̂δ,p dδ ∣∣∣∣ δ=0 = −H−1 θ̂ ∇θL(p, θ̂), (2)
where Hθ̂ = 1 n ∑ pi∈D ∇2θL(pi, θ̂) is the Hessian and positive definite by assumption, Iparam(p) ∈ RN , N is the number of network parameters, D is the original dataset. After getting the weighting of each example p, (Yang et al., 2022) propose generalization-guaranteed pruning or cardinalityguaranteed pruning to get the final compressed dataset D̂.

heading: 3 Method

heading: 3.1 Recall
text: Given the training dataset D and the query source x, we use BM25 (Robertson et al., 2009) to retrieve a set of relevant examples C1 for x. For each example pj in D, the BM25 score, denoted as R(pj , x), is computed. This score reflects the relevance of example pj to the query x. Specifically:
Rj = BM25(pj , x), (3)
where Rj is the relevance score of example pj with respect to query x.
Subsequently, we form the set C1 which consists of the top-N examples with the highest relevance scores:
C1 = {pj |j = 1, 2, . . . , N}, (4) where N is the number of examples we wish to recall for the given query x.

heading: 3.2 Influence-Awared Rerank
text: For each p in C1, we calculate the input representation of tokens p as P and the meta-gradient Gp = WV P . To compute Iparam(p) in Eq. (2), we require the Hessian of p for the parameter WV , which necessitates the computation of second-order derivatives. However, we only have access to firstorder derivatives approximations of the parameters. Considering that LLMs typically employ crossentropy loss and maximum likelihood estimation (MLE) for fine-tuning, we have opted to employ the Fisher matrix as an approximation of the Hessian (Barshan et al., 2020). The key to the approximation process is as follows:
∇2f(x) ≈ ∇f(x)∇f(x)⊤ (5) Then, combining the Eq. (2) with Eq. (5), the
expression of the influence function for p is:
Iparam(p) = −Ĥ−1 θ̂ Gp, (6)
where Hθ̂ = 1 n ∑ pi∈D GpG ⊤ p .
The score of C1 is determined by a combination of the influence score and the relevance score, represented as: S = { ∥Iparam(p1)∥2F +R1, . . . , ∥Iparam(pN )∥2F +RN } .
Finally, the K in-context learning examples in C2 are chosen by:
C2 = {pi|i ∈ I} ,where I = arg max I⊆{1,2,...,|S|}
|I|=K
S.
2

heading: 4 Experiments
text: In this section, we empirically verify the efficiency of our approach. The source code and all experiments have been shared at https://github.com/ ke-01/ICL.

heading: 4.1 Experiments setup
text: This section introduces the detailed information about our experiments.
Models. We use the open source GPT2 model family (Radford et al., 2019) (i.e., GPT2-Small, GPT2-Medium, GPT2-Large, GPT2-XL) as a representative of large models to verify the effectiveness of our method.
Datasets. We use five datasets spanning four tasks: linguistic analysis, hate speech detection, tweet classification, and semantic similarity. Specifically, we employ Linguistic Acceptability dataset (Cola) (Warstadt et al., 2019), online hate speech detection dataset (Ethos and Ethos-disability) (Mollas et al., 2022), Tweet_eval-stance_feminist and Tweet_eval_stance_hillary (Barbieri et al., 2020) from Twitter, and Sentences Involving Compositional Knowledge dataset (Sick) (Marelli et al., 2014). We use Accuracy and F1 score as evaluation metrics. Detailed dataset statistics and the prompt templates used can be found in Appendix A.2 and Appendix A.1.
Implementation Details. In the study, we chose K = 3 and K = 4 demonstrations to contrast ex-
ample selection methods from training data. We set N = 100 for all models. Sentences were either truncated or supplemented to have a uniform length at 50% of the average sentence length. Although using multiple transformer layers’ meta-gradient might be beneficial, considering the time efficiency, we used the first layer and obtained higher accuracy than baseline models.
Baselines. Considering the model proposed in this paper is unsupervised and requires no training, it possesses a higher generalizability and operational efficiency compared to models that undergo supervised training. To ensure a fair comparison, our primary baseline is the unsupervised BM25based In-Context Example Selection. Previous work (Wang et al., 2023; Gupta et al., 2023) has demonstrated that BM25 constitutes a robust baseline for demonstration selection, hence we juxtapose our methodology against BM25. The demonstrations selected by the BM25 are utilized across all GPT2 models.

heading: 4.2 Overall Performance
text: Tables 1 and 2 display results for three and four ICL examples, respectively. Observing the last two rows, our method consistently outperforms across all models and datasets. Using three and four examples, we surpass BM25 by averages of 5.17% and 6.64% in all metrics. Specifically, accuracy sees improvements of 6.33% and 7.80% over BM25. This underscores our approach’s superiority. We found
3
some higher model performance with three ICL examples compared to four, which can be explained by overfitting and example quality. Overfitting in few-shot learning means too many examples leads to adaptation to specific instances rather than general patterns, reducing accuracy on unseen data. Furthermore, if the additional fourth example is of lower quality or relevance, it can degrade model performance.

heading: 5 Related Work

heading: 5.1 In-context Learning
text: In-context learning (ICL) has emerged as a fresh approach in natural language processing (NLP), where large models predict based solely on contexts supplemented by several examples (Dong et al., 2022; Shin et al., 2022; Zhang et al., 2023; Bai et al., 2023). Numerous studies have sought to modify, improve, and comprehend ICL, encompassing topics like prompt tuning (Kim et al., 2022; Wang et al., 2022a; Mishra et al., 2022), intrinsic mechanism analysis (Chan et al., 2022; Li et al., 2023; Garg et al., 2022), evaluations (Srivastava et al., 2023; Wang et al., 2022b), and its use across various fields (Sun, 2023; Sun et al., 2024a,b), among others.

heading: 5.2 Demonstration Selection
text: The goal of demonstration selection is to identify optimal examples for ICL. (Liu et al., 2022) demonstrated that choosing the nearest neighbors as in-
context examples is an effective approach. The used distance measures include the pre-set L2 distance or the cosine similarity based on sentence embeddings. They introduced KATE, an unsupervised kNN retriever for in-context example selection. (Rubin et al., 2022) suggested a two-phase retrieval process for demonstration selection. For a given input, it initially employs an unsupervised retriever (like BM25) to retrieve similar candidate examples and then uses a supervised retriever, EPR, to pick demonstrations from these candidates. Recent studies indicate that LLMs exhibit strong sensitivity to the examples chosen, resulting in significant performance variations (Nie et al., 2022), dependency on example sequence (Lu et al., 2022), and at times, an insensitivity to the actual labels (Min et al., 2022). Our research focuses on reducing training overhead and condensing crucial data from the training set into in-context examples, which in turn amplifies the ICL’s effectiveness.

heading: 6 Conclusion
text: In summary, inspired by LLMs and ICL’s potential, we devised a two-stage method using the influence function for optimal in-context example selection. Our approach ensures relevance with the query source and efficiently determines influence scores. The result is an enhancement in ICL performance, with our experiments validating our model’s effectiveness. Our framework stands out due to its data-independent nature and minimal re-
4
liance on model parameters.

heading: 7 Limitation and Future Work
text: Given resource constraints and page limitations, we provide limited validation in this paper. Our model is a model-agnostic and free-training approach that can be applied to various in-context learning selection models. In the future, we will validate the effectiveness of our model on more large-scale language models, baselines, and datasets.

heading: 8 Acknowledgements
text: This work was funded by the National Key R&D Program of China (2023YFA1008704), the National Natural Science Foundation of China (No. 62376275, 62377044), Beijing Key Laboratory of Big Data Management and Analysis Methods, Major Innovation & Planning Interdisciplinary Platform for the “Double-First Class” Initiative, PCC@RUC, funds for building world-class universities (disciplines) of Renmin University of China. Supported by the Outstanding Innovative Talents Cultivation Funded Programs 2024 of Renmin University of China.

heading: A APPENDIX FOR REPRODUCIBILITY
text: A.1 Dataset In our experiments, we employ Linguistic Acceptability dataset (Cola) (Warstadt et al., 2019), online hate speech detection dataset (Ethos and Ethos-disability) (Mollas
6
Dataset Training Size Valid Size Label class
Sick 4439 495 3 Cola 8551 1043 2 Ethos-disability 346 87 2 Tweet_eval-stance_feminist 597 67 3 Tweet_eval-stance_hillaryt 620 69 3
Table 4: Dataset statistics.
et al., 2022), Tweet_eval-stance_feminist and Tweet_eval_stance_hillary (Barbieri et al., 2020) from Twitter, and Sentences Involving Compositional Knowledge dataset (Sick) (Marelli et al., 2014). Table 4 shows more detailed statics of the datasets.
A.2 Templates. Based on (Min et al., 2022; Yoo et al., 2022), we employ minimal templates. For the GPT-2 series model, we distinguish between input and label using a space, and this is also applied between demonstration examples. Table 3 presents templates from five datasets, where blue represents the fixed section, and black varies depending on the example content.
7

heading: Future Work
text: In the future, we will validate the effectiveness of our model on more large-scale language models, baselines, and datasets