Unified Demonstration Retriever for In-Context Learning

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 4644–4668 July 9-14, 2023 ©2023 Association for Computational Linguistics Unified Demonstration Retriever for In-Context Learning Xiaonan Li1∗, Kai Lv1∗, Hang Yan1, Tianyang Lin1, Zhu Wei2†, Yuan Ni3, Guotong Xie3, Xiaoling Wang2, Xipeng Qiu1† 1 Shanghai Key Laboratory of Intelligent Information Processing, Fudan University 1 School of Computer Science, Fudan University 2East China Normal University 3 Pingan Health Tech 1{lixn20, klv21, hyan19, tylin20, xpqiu}@fudan.edu.cn, 2wzhu@stu.ecnu.edu.cn, xlwang@cs.ecnu.edu.cn 3{niyuan442, xieguotong}@pingan.com.cn 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 highly dependent on the pro- vided 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 demon- strations for in-context learning. While pre- vious works focus on training task-specific retrievers for several tasks separately, these methods are often hard to transfer and scale on various tasks, and separately trained re- trievers incur a lot of parameter storage and deployment cost. In this paper, we propose Unified Demonstration Retriever (UDR), a sin- gle 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 feed- back. 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 base- lines. 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. 1 Introduction Large language models have shown an impres- sive in-context learning ability for various Natural Language Processing (NLP) tasks (Brown et al., 2020; Dong et al., 2022). In-context learning (ICL) is a recent learning paradigm where a language model (LM) learns a task by observing a few input- output pairs (demonstrations) and directly output ∗Equal Contribution †Corresponding Authors Retriever LM … Training set Demonstrations xtest X: Flights from Atlanta to Denver? Y: 1)flights 2)#1 from Atlanta 3)…… X: All flights on Friday? Y: 1)flights 2)#1 on Friday Prediction Xtest: Figure 1: Demonstration retrieval: Given a test input xtest, relevant demonstrations are retrieved from the training set. Then the inference LM takes demonstra- tions and xtest as input and generates the 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 up- dates, which makes it a promising alternative to supervised fine-tuning (Devlin et al., 2019). However, it has been shown that ICL’s perfor- mance highly depends on the provided demonstra- tions (Liu et al., 2022; Zhang et al., 2022; Li and Qiu, 2023a). This promotes the research …

Content truncated to 500 words