Learning to Retrieve In-Context Examples for Large Language Models (2024)

Abstract: Large language models (LLMs) have demonstrated their ability to learn in-context, allowing them to perform various tasks based on a few input-output examples. However, the effectiveness of in-context learning is heavily reliant on the quality of the selected examples. In this paper, we propose a novel framework to iteratively train dense retrievers that can identify high-quality in-context examples for LLMs. Our framework initially trains a reward model based on LLM feedback to evaluate the quality of candidate examples, followed by knowledge distillation to train a bi-encoder based dense retriever. Our experiments on a suite of 30 tasks demonstrate that our framework significantly enhances in-context learning performance. Furthermore, we show the generalization ability of our framework to unseen tasks during training. An in-depth analysis reveals that our model improves performance by retrieving examples with similar patterns, and the gains are consistent across LLMs of varying sizes.

## introduction in-context learning (icl) @xcite is an emerging learning paradigm that allows llms to perform tasks with few-shot examples, without requiring any updates to the model parameters. this approach stands in stark contrast to traditional machine learning, where models are typically trained on large datasets of labeled examples @xcite . in-context learning offers a significant advantage in domains where labeled data is scarce or expensive to obtain, as it greatly reduces the amount of required labeled data. there are several challenges associated with understanding and enhancing the effectiveness of incontext learning. one such challenge is that llms can be highly sensitive to the quality of the incontext examples provided @xcite . if the examples are not representative of the target task, then the model may not be able to learn effectively. empirical studies @xcite have demonstrated that using bm25 algorithm or off-the-shelf sentence embeddings @xcite to retrieve examples from the training set can substantially enhance the performance of in-context learning over random selection. another approach involves training dense retrievers based on the feedback signals from llms, which has shown promising results in semantic parsing @xcite , cross-task prompt retrieval @xcite , and unified multi-task retrieval @xcite . however, existing methods either focus on a relatively small language model @xcite , or fail to exploit the fine-grained feedback information from llms in a principled manner @xcite . in this paper, we propose a novel framework, llm-r (llm retriever), which aims to retrieve high-quality in-context examples for large language models. given an initial set of retrieved candidates, our framework ranks them based on the conditional llm log probabilities of the groundtruth outputs. subsequently, a cross-encoder based reward model is trained to capture the fine-grained ranking signals from llms. finally, a bi-encoder based dense retriever is trained using knowledge distillation. the reward model plays a crucial role in providing more informative soft-labels that are suitable for distillation, instead of using heuristically constructed one-hot labels. this pipeline can be iterated multiple times by retrieving a new set of candidates based on the latest dense retriever. for evaluation purposes, we assemble a diverse set of 30 nlp tasks, which span 9 categories, including question answering, natural language inference, commonsense reasoning, and summarization, among others. experimental results obtained using llama-7b @xcite demonstrate that our model improves the in-context learning performance by an average of 7.8% compared to random selection. similar improvements are also observed on held-out tasks and llms of varying sizes. further analysis reveals that the top-retrieved examples share similar input patterns or the same labels as the testing example. our model is particularly effective for classification tasks with ample training examples. in contrast, tasks such as closed-book question answering and commonsense reasoning rely more on the inherent capabilities of llms and are less sensitive to the quality of in-context examples. in-context learning is an emergent property of large language models (llms) that enables them to perform various tasks conditioned on a few inputoutput examples, without any parameter updates or fine-tuning. this property has been demonstrated in llms such as @xcite , gpt-neo @xcite , and llama @xcite another area of research is to explore different strategies for selecting and designing in-context examples for llms. recent studies @xcite @xcite have shown that using bm25 algorithm or fine-tuning dense retrievers based on llm feedback to retrieve from the training set can improve the performance of in-context learning. our work also falls into this area by proposing a novel training method. to model the interaction between in-context examples, determinantal point process @xcite and sequential decision-making @xcite are introduced as preliminary explorations. in contrast, structured prompting @xcite breaks the limitation of input context length and scales the number of in-context examples to thousands. dense retrieval is a widely used information retrieval approach that utilizes dense vectors to perform semantic matching between queries and documents in the latent space @xcite . compared to sparse retrieval methods such as bm25, dense retrieval exploits the powerful modeling capacity of pre-trained language models (plms) @xcite to learn relevance functions and has the potential to overcome the vocabulary mismatch problem. various techniques such as hard negative mining @xcite , knowledge distillation @xcite , and continual pre-training @xcite have been proposed to enhance the performance of dense retrieval. retrieval augmented llms combine the generative power of llms with the ability to retrieve relevant information from external sources @xcite @xcite . this paradigm has the potential to enhance the factual consistency of generated texts, make llms aware of the up-to-date knowledge, as well as provide a natural way for source attribution @xcite . the retrieved information can be incorporated into llms through various mechanisms, such as input concatenation @xcite , intermediate attention fusion @xcite , and output interpolation @xcite . for in-context learning, the goal of retrieval augmentation is to improve the performance of llms on downstream tasks by retrieving informative examples @xcite . ## preliminaries in this section, we provide a brief introduction to the problem setting of in-context example retrieval. given a test example x test from a target task and k in-context examples {(x i , y i )} k @xmath0, a frozen language model m is employed to predict an output y test through autoregressive decoding. the primary objective of in-context example retrieval is to retrieve k examples from p such that the predicted output y test is as close as possible to the ground-truth output y test based on some task-specific metrics. in this paper, the example pool p is the union of the training set for all the tasks in our evaluation. straightforward solutions include utilizing the bm25 algorithm or readily available text embedding models @xcite to retrieve examples from p by treating x test as a query. despite their simplicity, these methods ## methodology our proposed framework is depicted in figure 1 . it includes four main components: training data generation, reward modeling, dense retriever training, and inference, which are described in detail in the following subsections. ## analysis in this section, we examine the performance of llm-r across various tasks, llms, and model variants. unless explicitly specified, "llm-r" refers to the model with 2 training iterations. ## conclusion in this paper, we introduce an iterative training framework named llm-r to retrieve high-quality in-context examples for large language models. this framework generates training data by utilizing a frozen llm to rank the top retrieved candidates, and then learns a cross-encoder based reward model to capture the ranking preference. biencoder based dense retrievers are trained to distill the knowledge from the reward model. we conduct a comprehensive evaluation of llm-r on a diverse set of tasks and demonstrate that it consistently outperforms various strong baselines. our model also generalizes well to held-out tasks and llms of varying sizes.