MDR : Model-Specific Demonstration Retrieval at Inference Time for In-Context Learning (2024)

Abstract: Recently, retrieval-based in-context learning (ICL) methods for selecting demonstrations have been widely investigated. Existing methods train a dense retriever to retrieve the most appropriate demonstrations for a given test query, which improves ICL performance. However, we find that distinct LLMs exhibit different biases for “what is a good demonstration” since they possess differences in training data, model architectures and training methods. As a result, a demonstration suitable for one LLM may not be appropriate for others.Previous approaches ignore the model bias and fail to retrieve the most appropriate demonstrations for different inference LLMs, resulting in a degradation of ICL performance.To address this problem, we propose a simple yet effective metric to evaluate the appropriateness of demonstrations for a specific inference LLM. Furthermore, we introduce a Model-specific Demonstration Retrieval (MDR) method for ICL at inference time, which considers the biases of different LLMs. We test MDR on seen and unseen tasks with multi-scale inference LLMs, such as GPT-Neo-2.7B, LLaMA-7B and Vicuna-13B. Experiments on 23 datasets across 11 data domains highlight the remarkable effectiveness of MDR, showcasing improvements of up to 41.2% in comparison to methods that neglect model biases.

## introduction large language models (llms) such as @xcite , have emerged impressive abilities of handling a wide range of tasks. in-context learning (icl) @xcite , a new learning paradigm, allows llms to perform multi-tasks by observing a few demonstrations, without requiring any updates to the model parameters @xcite @xcite . in practice, icl typically uses a concatenation of a short sequence of annotated in-context examples known as demonstrations and a test query with its task-specific instruction. it offers a promising alternative to supervised fine-tuning since it greatly reduces the amount of required labeled data. the quality of the selected demonstrations plays a crucial role in icl @xcite . to improve icl's performance, given a test query, previous works mainly construct a demonstration pool and retrieve the most textual or semantically similar demonstrations by using bm25 @xcite , sentence-bert @xcite or other offthe-shelf sentence embeddings, which outperforms random selection. but these approaches are heuristic and sub-optimal as they lack guidance from task supervision. alternatively, another approaches utilize the feedback signals of a language model (lm), referred to as the scoring lm, to identify positive and negative demonstrations. subsequently, they employ a two-encoder retriever, where one encoder encodes the test query and the other encodes the demonstrations. the retriever is trained using contrastive learning to optimize the similarity between the test query and positive demonstrations, simultaneously minimizing it for negative demonstrations. during inference, the demonstrations exhibiting the highest similarity are retrieved and concatenated with the test query to form the model input. this approach has shown better performance in specific task domains, such as question answering @xcite and semantic parsing @xcite . recent efforts have shifted towards cross-task prompt retrieval @xcite and unified multi-task retrieval @xcite by incorporating signals from various tasks across multiple data domains to train one retriever, which has demonstrated impressive improvements. however, distinct llms exhibit different biases @xcite for "what is a good demonstration" since they possess differences in training data, model architectures and training methods. as a result, a demonstration suitable for one llm may not be appropriate for others. we conclude that existing methods fail to retrieve the most appropriate demonstrations for different inference llms. there are two key factors: (i) bias introduced in training phase: previous methods solely rely on one scoring lm to identify positive and negative demonstrations. the inherent bias of the scoring lm can impact this identification process, potentially introducing noise to the retriever training. consequently, this may lead to the retrieval of bad demonstrations when inference on different llms. (ii) bias neglected in inference phase: a high-quality demonstration should not only exhibit similarity to the test query, but also align with the bias of the inference model. previous methods employ the similarity-based extraction of demonstrations, which greatly neglect the inference model's bias. to test the bias issue, we train 3 tailored retrievers on rte dataset using three different scoring lms, including gpt-neo-2.7b, llama-7b and vicuna-13b. then we test the performance of each retriever on sst2 dataset under different inference llms, respectively. as shown in fig 1 , the one-to-one corresponding scoring lm and inference llm perform the best while performances on different inference llms can be inferior, which illustrates that using a single scoring lm for various inference llms limits the performance of icl. since training multiple retrievers for multiple inference llms is costly, retrieving demonstrations that meet inference model's bias at inference stage is crucial to improve the icl performance. however, directly estimating and evaluating the model bias is impractical, the evaluation of a demonstration's ap-propriateness for a specific inference llm remains a challenging and under-explored aspect in icl. to address this limitation, we propose a simple yet effective eigenvalue-based evaluation metric to assess the appropriateness of demonstrations for a specific inference llm. based on this metric, we further propose mdr, a model-specific demonstration retrieval (mdr) method for icl at inference time. given a pre-trained demonstration retriever, mdr retrieves the most appropriate demonstrations from a pre-constructed demonstration pool for a given test query and a specific inference llm. results in fig 1 verify that mdr is effective of mitigating the bias issue by reducing the performance gap between the corresponding and the non-corresponding scoring-inference llms. to evaluate mdr's generalization across diverse models, we test mdr on a variety of seen and unseen tasks with different llms at inference time. the outcomes highlight the effectiveness of mdr as a simple and straightforward method to select the most appropriate demonstrations for a specific llm without any additional training or human annotation. our contributions are summarized as follows: ## task definition we aim to improve the performance of icl by retrieving the most appropriate demonstrations for any given task input. specifically, given a pretrained demonstration retriever, a test query x test and an inference model lm, we retrieve the top-k most appropriate demonstrations {p j } k @xmath0. these demonstrations are concatenated with the test query in icl fashion. our objective is to optimize the model prediction y * test to align with the ground truth y test , where y * test can be denoted by: ## methodology provided by a frozen pre-trained demonstration retriever r trained on a demonstration pool p, and various inference models with their parameters, we aim to retrieve the most appropriate demonstrations from p for a test query to improve the icl performance without any retriever parameters updating. to achieve this, in this section, we first introduce a novel evaluation metric to evaluate the appropriateness of a demonstration to a specific inference llm. subsequently, we propose an inference framework named mdr, to re-rank demonstrations with our proposed evaluation metric. ## related work in this section, we introduce previous works on exploring different strategies for selecting in-context demonstrations for llms. two primary types of retrievers for icl are commonly used. one is off-the-shelf retrievers, such as finetuned bert @xcite , bm25 @xcite or sbert @xcite . the other approaches train a taskspecific retriever using designed signals, such as question answering @xcite , code generation @xcite and dialogue state tracking @xcite . in particular, @xcite introduce efficient prompt retriever (epr), which employs a lm to score demonstrations, and trains a dense retriever using contrastive learning. additionally, @xcite propose ceil to select diverse and helpful icl demonstrations by using determinantal point processes. recently, @xcite propose udr, a unified retriever designed for a wide range of tasks. unlike udr, which focuses on testing on seen tasks, @xcite utilize a single retriever for cross-task and cross-model scenario on unseen tasks. however, previous studies neglect the biases of different inference llms. in this paper, we fully consider llms biases and strive to retrieve the most appropriate demonstrations for a specific inference model during inference stage. ## conclusion in this paper, we introduce mdr, a simple yet effective method to evaluate and retrieve the modelspecific appropriate demonstrations for icl without re-training the retriever. mdr offers the feasibility of evaluating the appropriateness of a demonstration for llms without supervision. moreover, given a trained retriever, mdr makes it more applicable to use it on various larger inference models. in summary, mdr offers a promising insight for enhancing llms performance by leveraging their biases, presenting a notable perspective over existing retrieval-based methods and even prompt engineering.