Unified Demonstration Retriever for In-Context Learning (2023)

Abstract: In-context learning is a new learning paradigm where a language model conditions on a few input-output pairs (demonstrations) and a test input, and directly outputs the prediction. It has been shown sensitive to the provided demonstrations and thus promotes the research of demonstration retrieval: given a test input, relevant examples are retrieved from the training set to serve as informative demonstrations for in-context learning. While previous works train task-specific retrievers for several tasks separately, these methods are hard to transfer and scale on various tasks, and separately trained retrievers will cause a lot of parameter storage and deployment cost. In this paper, we propose Unified Demonstration Retriever (UDR), a single model to retrieve demonstrations for a wide range of tasks. To train UDR, we cast various tasks’ training signals into a unified list-wise ranking formulation by language model’s feedback. Then we propose a multi-task list-wise ranking training framework with an iterative mining strategy to find high-quality candidates, which can help UDR fully incorporate various tasks’ signals. Experiments on 30+ tasks across 13 task families and multiple data domains show that UDR significantly outperforms baselines. Further analyses show the effectiveness of each proposed component and UDR’s strong ability in various scenarios including different LMs (1.3B 175B), unseen datasets, varying demonstration quantities, etc. We will release the code and model checkpoint after review.

## introduction large language models have shown an impressive in-context learning ability for various natural language processing (nlp) tasks @xcite . in-context learning (icl) is a recent learning paradigm where a language model (lm) learns a task by observing a few inputoutput pairs (demonstrations) and directly output the prediction of the given test input. thus icl can unify a wide range of nlp tasks through one language model's inference without parameter updates, which makes it a promising alternative to supervised fine-tuning @xcite . however, it has been shown that icl's performance highly depends on the provided demonstrations @xcite @xcite . this promotes the research of demonstration retrieval for in-context learning @xcite @xcite : as shown in figure 1 , given a test input, relevant examples are retrieved from an annotated training set, to serve as informative demonstrations for icl. there are about two lines of methods to retrieve demonstrations. one is to leverage off-the-shelf retrievers, e.g., @xcite or sentence-bert (reimers and gurevych, 2019a). they can retrieve demonstrations that are textually or semantically similar to the test input and achieve empirical improvements. thanks to their versatility, they can serve for extensive nlp tasks, but they are heuristic and sub-optimal since they are not guided by task supervision. another line is to train a task-specific retriever by a specially designed task signal. @xcite train the retriever for knowledge-based question answering, based on the logic form's surface similarity. @xcite explore icl on dialogue state tracking and design the similarity between dialogue's states as the retriever's training signal. @xcite and @xcite leverage the lm's feedback to train demonstration retrievers for semantic parsing in english and cross-lingual scenarios, respectively. these task-specialized retrievers show better performance than the former, but they still face two challenges: 1. these explorations are limited to a small range of tasks and demonstrated separately on each task, e.g., semantic parsing or dialogue state tracking, which restricts systematic and compatible research on demonstration retrieval for icl while icl is a unified framework for extensive tasks. 2. it is costly for these methods to transfer and scale on various tasks and the reason is twofold: (i) they need to design a specialized training signal for each task. (ii) the number of retrievers will scale up with increasing tasks, which results in massive parameter storage and deployment costs. to address these limitations, we explore learning various tasks' demonstration retrieval in a unified formulation and propose unified demonstration retriever (udr), a single multi-task model for demonstration retrieval of a wide range of tasks. to train udr, we cast various tasks' training signals into a unified list-wise ranking formulation. for a training example from task t , we select a list of candidate examples from t 's training set and rank them by lm's feedback. then we propose a multitask list-wise ranking training framework, with an iterative mining strategy to find high-quality candidates. specifically, we iteratively train the retriever to rank candidates and use itself to find high-quality positive candidates and hard negatives. compared with the representative method for demonstration retrieval, epr @xcite , which trains the retriever by the binary label from lm's feedback and selects candidates in a manually limited range, our training framework can explore the entire dataset to get high-quality candidates and help udr fully incorporate the lm's feedback through list-wise ranking training. experiments on 30+ tasks across 13 task families and multiple data domains show that udr significantly outperforms baselines and further analyses show the effectiveness of each proposed component and udr's strong ability under various scenarios including different lms (1.3b ∼ 175b), unseen datasets, varying demonstrations quantities, etc. we release the code and model checkpoint at https://github.com/kailv69/udr . ## unified demonstration retriever provided a language model g, a training set d train and a test case x test , demonstration retrieval aims to retrieve x test 's relevant demonstrations from d train to help lm g decode the target output. previous works @xcite @xcite propose task-specialized methods for several tasks separately, but they are hard to transfer and scale on various tasks. in this work, we focus on learning various tasks' demonstration retrieval in a unified formulation and propose udr , a single model for demonstration retrieval of a wide range of tasks, as shown in figure 2 . we introduce its architecture, training, and inference as follows. ## related work in this section, we introduce previous demonstration retrievers for in-context learning, and explain the difference between udr and them. in general, there are two kinds of demonstration retrievers for icl. one is to leverage off-the-shelf retrievers. for example, @xcite propose to use a fine-tuned bert to encode examples and use a 2023 ) propose claif to enhance the sentence embedder by the gigantic language model's feedback. specifically, they use @xcite to generate the data of sentence pairs and then score them by the output of gpt-3, which depends on the strong natural language understanding ability of gpt-3. different from them, we leverage the conditional probability to measure the helpfulness of an example, which only needs a small language model, and is more efficient and environmental-friendly. recently, @xcite propose mot (memoryof-thought) to let the llm self-improve in two stages: 1. before test stage, the llm generate reasoning paths and answers on an unlabeled dataset for itself, 2. at test stage, the llm retrieves relevant reasoning paths (memory) to help itself answer the given test question. while mot focuses on the scenario with unlabeled dataset and uses the llm for retrieval, we train a small retriever by a lm's feedback from tasks' supervision and thus the proposed method is more lightweight. we leave demonstration retrieval with reasoning paths or unlabeled datasets as future work. ## conclusion in this paper, we propose udr, a single multi-task model for a wide range of tasks' demonstration retrieval. to train udr, we cast various tasks' training into a unified list-wise ranking formulation by language model's feedback, and propose a multitask list-wise ranking training framework, with an iterative mining strategy to find high-quality candidates. experiments on 30+ tasks show that udr significantly outperforms baselines. further analyses show the effectiveness of each proposed component and udr's strong ability in various scenarios including different lms (1.3b ∼ 175b), unseen datasets, varying demonstration quantities, etc.