Unified Demonstration Retriever for In-Context Learning
Xiaonan Li,Kai Lv,Hang Yan,Tianyang Lin,Zhu Wei,Yuan Ni,Guotong Xie,Xiaoling Wang,Xipeng Qiu
DOI: https://doi.org/10.18653/v1/2023.acl-long.256
2023-01-01
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 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 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 retrievers incur 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 similar to 175B), unseen datasets, varying demonstration quantities, etc.