Unified Demonstration Retriever for In-Context Learning
Year: 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 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 ∼ 175B), unseen datasets, varying demonstration quantities, etc.

heading: 1 Introduction
text: Large language models have shown an impressive 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 inputoutput pairs (demonstrations) and directly output
∗Equal Contribution †Corresponding Authors
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 (Devlin et al., 2019).
However, it has been shown that ICL’s performance highly depends on the provided demonstrations (Liu et al., 2022; Zhang et al., 2022; Li and Qiu, 2023a). This promotes the research of demonstration retrieval for in-context learning (Liu et al., 2022; Rubin et al., 2022; Shi et al., 2022): 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., BM25 (Robertson and Zaragoza, 2009) 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. Das et al. (2021) train the retriever for knowledge-based question answering, based on the logic form’s surface similarity. Hu et al. (2022) explore ICL on dialogue state tracking
4644
and design the similarity between dialogue’s states as the retriever’s training signal. Rubin et al. (2022) and Shi et al. (2022) 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(Rubin et al., 2022), 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.

heading: 2 Unified Demonstration Retriever
text: Provided a language model G, a training set Dtrain and a test case xtest, demonstration retrieval aims to retrieve xtest’s relevant demonstrations from Dtrain to help LM G decode the target output. Previous works (Das et al., 2021; Rubin et al., 2022; Shi et al., 2022) 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.

heading: 2.1 Bi-encoder with Task Instruction
text: UDR is based on the prevailing bi-encoder architecture, dense passage retriever (DPR) (Karpukhin et al., 2020), which encodes the query example and candidate examples separately and then calculates their similarity. To distinguish examples from different tasks, UDR encodes the example together with its task instruction, which is a short piece of text related to the task objective. Taking CNN/DailyMail (Hermann et al., 2015) as an example, its task instruction can be “Summarize the text”. Given an example query x and a candidate demonstration z = {x′, y′} from task Ti, UDR uses the query encoder Eq and demonstration encoder Ed to encode them respectively and calculates their similarity as:
sim(x, z) = Eq(Ii ⊕ x)⊤Ed(Ii ⊕ z), (1)
where Ii is Ti’s task instruction and ⊕ is the concatenation operator. Eq and Ed are two multilayer Transformer (Vaswani et al., 2017) encoders with “CLS” pooling and can be initialized with pre-trained models (Devlin et al., 2019).
Thus, we can not only get task-specific features by specifying the task instruction, but also retain the uniformity and parameter efficiency of ICL.

heading: 2.2 Learning from LM Feedback
text: To train the demonstration retriever, previous works (Das et al., 2021; Rubin et al., 2022; Hu et al., 2022) design task-specific training signals for several tasks separately, which makes their methods hard to transfer and scale on various tasks, and hinders systematic and compatible research on demonstration retrieval. For UDR’s training, we propose to cast various tasks’ training signals into
a unified list-wise ranking formulation. Then we introduce a multi-task list-wise ranking training framework, where we iteratively let the retriever itself to mine high-quality candidates and learn to rank them in turn, across various tasks, shown in Algorithm 1. We introduce the list-wise ranking training and iterative mining strategy as follows.

heading: 2.2.1 Ranking Candidates by LM
text: Given a training example (x, y) and its candidates Z = {zi}li=1, we first rank these candidates as:
r(zj) = rank(s(zj)|{s(zi)}li=1) (2) sgen (zj) = pG (y | zj , x) , (3)
scls (zj) = pG (y | zj , x)∑
y′∈Y pG (y ′ | zj , x)
, (4)
where s(zj) = sgen(zj) for generation tasks and s(zj) = scls(zj) for classification and multi-choice tasks. pG(·|·) is the LM G’s conditional likelihood. Y is the label space or choices of the classification or multi-choice task, respectively. For simplicity, we omit special tokens and classification tasks’ verbalizers in the equations above.
First we use G to score each candidate (Rubin et al., 2022) and calculate s(zj) as the ground truth y’s likelihood conditioned on the candidates zj and the query input x. s(zj) indicates the importance of zj for G to encode x and generate the ground truth y. Then we rank Z according to {s(zi)}li=1. The more important zj is for x, the higher zj’s rank will be. Thus we unify various tasks’ training signals into the same list-wise ranking formulation using LM’s feedback, instead of designing task-specific objectives (Das et al., 2021; Hu et al., 2022).

heading: 2.2.2 Loss Function
text: With these candidates’ ranks from G’s feedback, we propose to use the following loss function to in-
ject the ranking signal into the retriever E, inspired by LambdaRank (Burges, 2010): Lrank = ∑
zi,zj∈Z w ∗ log(1 + esim(x,zj)−sim(x,zi)) (5)
where w = max(0, 1r(zi) − 1 r(zj) ).
For those zi and zj where r(zi) < r(zj), Lrank will draw sim(x, zi) up and optimize the retriever towards sim(x, zi) > sim(x, zj). Additionally, w adjusts the weight for each pair of demonstrations and inject list-wise ranking information into Lrank. When zi has a much higher rank than zj , e.g, r(zi) = 1 and r(zj) = 10, w will be a high weight and strongly draw sim(x, zi) up from sim(x, zj). Since we optimize the retriever on demonstration pairs under different w, Lrank can help UDR fully incorporate candidates’ listwise ranking signals from G’s feedback for various tasks and learn to retrieve those helpful demonstrations.
To fully leverage the computation of the same batch, we also use the in-batch negative loss as:
Lib = − log esim(x,z ∗) ∑
z∈Z e sim(x,z)
, (6)
where z∗ is the rank-1 candidate of x and Z is all candidates (x’s or not x’s) in the batch. Each batch is sampled from the same task, and to alleviate the bias towards high-resource tasks, we sample each task according to the multinomial distribution with probabilities {p(Ti)}Ti=1 as:
p(Ti) = qαi∑T j=1 q α j with qi = |DTi | ∑T j=1 |DTj | , (7)
where DTi is the ith task’s dataset. α is a predefined hyper-parameter and we follow Conneau and Lample (2019) to set α as 0.5.
The overall loss function of UDR is the integration of these two losses as follows,
L = λ ∗ Lrank + (1− λ) ∗ Lib, (8)
where λ is a pre-defined hyper-parameter.

heading: 2.2.3 Iterative Candidate Mining
text: The selection of candidates can be a key factor for retriever’s training (Karpukhin et al., 2020; Xiong et al., 2021). It is desirable for UDR to take the entire training set as candidates to provide abundant ranking signals. However, it is infeasible since scoring all pairs of training examples is quadratic in |D| and costly. Previous work (Rubin et al., 2022) selects those examples which have textually similar targets with x’s as candidates. However, it may bias the retriever to learn among candidates with highly similar targets. Meanwhile, it can probably miss important demonstrations. For instance, if an example z contains relevant logic with the query x but has a dissimilar target with x’s, the valuable z will not be selected as candidate to provide signal for the retriever. So, we propose an iterative mining strategy to select candidates by the retriever itself. Specifically, we iteratively train the retriever and use it to select candidates in turn. At each iteration, we update each training example’s candidates as:
Z∗ = top-Kz∈D sim(x, z) (9)
where D is the task’s entire training set. Then we will use LM G to score and rank Z∗. The new candidates in Z∗ can be divided into two categories. If a new candidate z has a low score, it means that we find a hard-negative candidate that can provide crucial negative signal for the retriever. If the score of z is high and even higher than all old candidates, it means that we find a valuable positive candidate that can help the retriever learn to find informative demonstrations. Thus, with iterative mining, we can explore the entire dataset, find high-quality candidates and improve training progressively. Before the first iteration, the retriever is untrained, so we initialize candidates based on surface similarity, inspired by Rubin et al. (2022).
For computational efficiency, we first update candidates and score Z∗ at each iteration, and then randomly sample l of Z∗ and rank them at each training step. In summary, Algorithm 1 shows the UDR’s overall training procedure.
Algorithm 1 Multitask List-wise Ranking Training Require: Bi-encoder Eq and Ed, language model
G, Training sets of T tasks {DTi}Ti=1 1: Initialize the bi-encoder. 2: Initialize candidates of each training example. 3: Score initialized candidates by G. 4: for Each iteration do 5: for Each training step, Ti ∼ p(T ) do 6: Sample a batch of examples. 7: For each example, sample l examples
z1∼l from its candidates and rank z1∼l by G’s score.
8: Update the bi-encoder’s parameters by L. 9: end for
10: Update candidates by new Eq and Ed. 11: Score new candidates by G. 12: end for

heading: 2.3 Inference
text: After training, we encode each task Ti’s training set using Ed(pi ⊕ ·). At the test stage, given a task Ti’s input, xtest, we use Eq(pi⊕·) to compute its encoding and then use FAISS (Johnson et al., 2021) to search over Ti’s training set to find the most relevant demonstrations, ascendingly sorted by sim(xtest, ·), D = (z1, z2, · · · , zL). For generation tasks, the number of final demonstrations, L, is determined by the LM G’s maximal input length C. Specifically, ∑L i=1 |zi|+ |xtest|+ |y| ≤ C, where |y| is the pre-defined maximal length of the generated target. For classification and multi-choice tasks, we observe that increasing L brings negligible performance improvement and thus we set L to a small value, 8. We conduct further analysis of the number of demonstrations in section 3.3.5. Finally, we use greedy decoding to get the result of G([z1; z2; · · · ; zL;xtest]). Notice that here D is ascendingly sorted by sim(xtest, ·) unless otherwise specified. Our analysis in section 3.3.4 shows that different orderings lead to similar performance. Thus we use the same ordering strategy with EPR (Rubin et al., 2022) for fair comparison.

heading: 3 Experiment

heading: 3.1 Experimental Settings
text: Dataset We train UDR on a wide range of NLP tasks, consisting of about 40 tasks across 13 task families and multiple data domains, including: Sentiment Classification: SST-2, SST-5 (Socher et al., 2013), Amazon (McAuley and Leskovec, 2013), Yelp (Zhang et al., 2015), MR (Pang and Lee, 2005)
and CR (Amplayo et al., 2022); Topic Classification: AGNews, Yahoo (Zhang et al., 2015), TREC (Voorhees and Tice, 2000) and DBPeida (Lehmann et al., 2015); Multi Choice: COPA (Roemmele et al., 2011), Cosmos QA (Huang et al., 2019), Commonsense Validation and Explanation (ComE and ComV) (Wang et al., 2019b); NLI: MNLI (Williams et al., 2018), SNLI (Bowman et al., 2015) and RTE (Bar-Haim et al., 2014); Subjectivity Classifcation: Subj (Pang and Lee, 2004); Lingustic Acceptibility: COLA; Semantic Parsing: BREAK (Wolfson et al., 2020), MTOP (Li et al., 2021) and SMCalFlow (Andreas et al., 2020); Text Summarization: CNN/DailyMail (Hermann et al., 2015), PubMed (Cohan et al., 2018) and Reddit (Kim et al., 2019); Commonsense Generation: CommonGen (Lin et al., 2020); Story Generation: Roc Story and Ending Generation (Mostafazadeh et al., 2016); Code Summarizaton: Go, Python, Java and PHP (Lu et al., 2021); Text Simplifiction: WikiAuto + Turk/ASSET (Jiang et al., 2020); Data to Text: DART (Nan et al., 2021) and E2E (Dušek et al., 2019). These tasks’ input/output, statistics, split and evaluation metrics are in Appendix A.
Implementation Details We follow EPR (Rubin et al., 2022) to use GPT-Neo-2.7B (Black et al., 2021) as the scoring LM and the inference LM for most experiments in the paper unless otherwise specified. We also explore UDR’s transferability across different inference LMs in section 3.3.2. Following EPR(Rubin et al., 2022), we initialize Eq and Ed as two separate “BERT-base-uncased” encoders (Devlin et al., 2019). We list the overall hyper-parameters and implementation details in Appendix B. On each task, we use one specific template for scoring and inference (see Appendix A). We evaluate UDR’s performance when inference templates are different with the scoring template in Appendix C, and the results show that UDR has stable performance across varying inference templates, which reflects UDR’s generality.
Model Comparison With the same inference LM, GPT-Neo-2.7B, we compare UDR with previous methods for demonstration retrieval by the downstream ICL performance, including: 1. Random: We randomly sample demonstrations from the corresponding task’s training set. 2. BM25 (Robertson and Zaragoza, 2009): A prevailing sparse retriever. For each test input xtest, we use BM25 to retrieve examples with the most
similar input. 3. SBERT (Reimers and Gurevych, 2019b): We use the Sentence-BERT as the dense demonstration retriever. Specifically, we follow Rubin et al. (2022) to take “paraphrase-mpnet-basev2” to encode the test input xtest and training set’s inputs, and retrieve the examples with the most similar input as demonstrations. 4. Instructor (Su et al., 2022): Instructor is a recently proposed competitive text embedding model trained on 330 tasks with instructions. By providing the specialized instruction, it can serve for demonstration retrieval. For fair comparison, we conduct experiments on its released base-size model. 5. DR-Target: This baseline is inspired by previous works on generation tasks like dialogue state tracking, question answering and code generation (Hu et al., 2022; Das et al., 2021; Poesia et al., 2022), which design the task-specific target’s similarity and use examples with similar targets to train the retriever. Here we use BM25 as the similarity function for each task’s target output. Specifically, we use BM25 to find positive pairs with similar targets and use DPR (Karpukhin et al., 2020) for training. 6. EPR (Rubin et al., 2022): EPR is a recently proposed representative method for training demonstration retriever. It uses the language model to assign candidate examples with positive and negative labels and thus trains a task-specific demonstration retriever by DPR. For fair comparison, we train EPR on each task using the same hyper-parameters of UDR. Specially, we discuss EPR’s candidate quantity in Appendix B.
Except that the performance of Random, BM25, SBERT and EPR on semantic parsing is from the previous paper (Rubin et al., 2022), other results are from our implementation since they are not explored previously.

heading: 3.2 Main Results
text: We show the performance comparison of classification tasks and generation tasks in Table 1 and Table 2, respectively. We can see that UDR outperforms baselines significantly on most tasks, which shows UDR’s best overall demonstration retrieval ability on a wide range of NLP tasks. Specially, compared with DR-Target and EPR, UDR has better overall performance and this shows the effectiveness of our unification of various tasks’ training signals. Meanwhile, compared with Instructor (Su et al., 2022), the text embedding model trained on 330 tasks’ text pairs, UDR has an improvement
of 10 and 6.7 points for classification and generation tasks respectively with less training data. This straightly demonstrates that our proposed training framework can help UDR incorporate LM’s feedback through a unified ranking formulation and better retrieve informative demonstrations.
Additionally, we find the random baseline shows the worst performance on most tasks and this reflects the necessity to retrieve high-quality relevant demonstrations. Meanwhile, EPR and UDR have better performance than other methods, which reflects the importance of LM’s feedback. Among these datasets, we notice a different trend on text summarization datasets like CNN/DailyMail and Reddit, on which these methods have similar performance. We conjecture that the LM can already have the knowledge of summarization since there are a lot of “[Article, TL;DR, Abstract]” texts in its pre-
training corpus (Radford et al., 2018), thus random demonstrations can well activate LM’s summarization ability without example-specific information.

heading: 3.3 Analysis

heading: 3.3.1 Ablation Study
text: To evaluate the effect of UDR’s each component, we conduct ablation study on SMCalFlow, SST-2 and Java code summarization, shown in Table 3. When removing list-wise ranking training, we use EPR’s training strategy (Rubin et al., 2022). We can see that removing task instructions cause slight performance degradation, which indicates that they can help UDR distinguish examples from various tasks and thus get better task-specific features. Meanwhile, we can see that UDR has a slightly better performance than the single-task counterpart on SST-2 and Java. We suppose that is because
there are several relevant tasks in UDR’s training tasks and our multi-task ranking unification can help UDR fully share these tasks’ knowledge. The performance of single-task UDR still outperforms EPR significantly and this straightly reflects that our training components, i.e., list-wise ranking formulation and iterative candidate mining strategy, can 1. help UDR better incorporate LM’s feedback than EPR 2. serve as a competitive universal training method for a task-specific retriever. Removing list-wise ranking training and iterative candidate mining both cause performance degradation, which straightly indicates their effectiveness.

heading: 3.3.2 Transferability across Different LMs
text: In this section, we evaluate UDR’s transferability across different inference LMs on SMCalFlow and E2E. Specifically, we compare BM25, EPR and UDR on inference LMs with different sizes, including: GPT-Neo-1.3B (Black et al., 2021), GPT-J (6B) (Wang and Komatsuzaki, 2021), CodeDavinci-002 (175B) (Chen et al., 2021) and TextDavinci-003 (175B) (Brown et al., 2020; Ouyang et al., 2022) and we show the result in Table 4. When comparing UDR with baselines, the trends are similar with using GPT-Neo-2.7B (the scoring LM) as inference LM. UDR outperforms BM25 and EPR significantly and it shows UDR’s strong transferability across different inference LMs. Meanwhile, we find that UDR with larger inference LM can improve performance such as Text-Davinci-003 on SMCalFlow and GPT-J on E2E, which shows UDR’s potential utility in the
future where more competitive large-scale LM is built. When we demonstrate the example-specific demonstration transferability across different inference LMs in this paper, Li and Qiu (2023a) show that task-level demonstrations also exhibit such transferability. We leave the analysis of the transferablity of ICL’s demonstrations across different LMs as future work.

heading: 3.3.3 Performance on Unseen Datasets
text: In this section we explore UDR’s zero-shot transferability and evaluate it on unseen datasets including: 1. Twitter sentiment classification (Naji, 2012) 2. question-answering NLI (QNLI) (Wang et al., 2019a) 3. Ruby and JavaScript code summarization (Lu et al., 2021). These domains or programming languages (Twitter, NLI on QA, Ruby and Javascript) are never seen during UDR’s training and thus can straightly reflect UDR’s zero-shot transferability. We compare UDR with two powerful universal retrievers, BM25 and SBERT, and show the result in Table 5. We can see UDR significantly outperforms BM25 and SBERT on these unseen datasets by about 10 points on average, which shows that the learned ranking knowledge inside UDR can be well transferred and generalized to unseen datasets.

heading: 3.3.4 The Order of Demonstrations
text: Previous work (Lu et al., 2022) has revealed that ICL is sensitive to demonstrations’ order when using random examples. Specifically, the same randomly sampled demonstrations with different orders can lead to the performance between random guess and near state-of-the-art. Here we explore the effect of ordering on example-specific demonstrations retrieved by UDR. We compare 3 demonstrations’ orders: 1. random, for this setting, we run experiments with 10 different random seeds and report the best and worst performance. 2. descending sorted by UDR’s score, i.e, the demonstration which has the highest similarity with xtest is put at the beginning of LM’s input. 3. ascending sorted by UDR’s score, opposite to “2”. The result is shown in Table 6. We
observe a different phenomenon from that in previous work (Lu et al., 2022). In general, The performance of UDR’s demonstrations with different orders is more stable than previously investigated random examples. Across these tasks, different orders’ performance gap is within 1 point, and it is far less than the performance fluctuation of up to tens points when using random examples (Lu et al., 2022). This indicates that high-quality demonstrations are less sensitive to the ordering and stabilize in-context learning, which is consistent with the analysis in previous work (Chen et al., 2022; Li and Qiu, 2023a).

heading: 3.3.5 The Impact of Demonstration Quantity
text: We compare UDR with BM25 and EPR under different amounts of demonstrations on two classification tasks: Yelp and RTE, and two generation tasks: WikiAuto and Java code summarization. We show results in Figure 3. We can see that UDR outperforms baselines consistently across varying amounts of demonstrations. Meanwhile, we can draw two conclusions from the results: 1. The number of demonstrations has a greater impact on generation tasks than classification tasks. Specifically, as the number of demonstrations increases, generation tasks’ performance gets significant improvements while classification tasks’ has slight or no improvements. 2. The quality of demonstrations can be more important than their quantity. In detail, UDR with the quota of 2 demonstrations still outperforms BM25 and EPR with 8 demonstrations. This also reflects the strong demonstration retrieval ability of UDR. Li and Qiu (2023b) observe the similar trends in the CoT-retrieval scenario, indicating that the relevance of the used reasoning paths is more important than their quantity.

heading: 4 Related Work
text: 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, Liu et al. (2022) propose to use a fine-tuned BERT to encode examples and use a
KNN-based method to retrieve semantically similar demonstrations to improve ICL. Agrawal et al. (2022) use BM25 to retrieve demonstrations for machine translation. Compared with them, UDR incorporates various tasks’ supervision by unified LM’s feedback and thus can better retrieve informative demonstrations. Another approach is to train a task-specific retriever by a designed task-specific signal. Das et al. (2021) explore demonstration retrieval for knowledge-based question answering and define the F1 score of logic forms as soft-label to train the retriever. Poesia et al. (2022) train a demonstration retriever for code generation, based on the edit distance of abstract syntax trees. Hu et al. (2022) define the similarity between dialogue states, and use it to train a demonstration retriever for dialogue state tracking. Rubin et al. (2022) propose Efficient Prompt Retriever (EPR) for semantic parsing, which is to use the language model to score examples, assign positive and negative labels for them and use DPR (Karpukhin et al., 2020) to train a demonstration retriever. Shi et al. (2022) explore demonstration retrieval for cross-lingual semantic parsing using a similar example scoring method with EPR. These task-specific methods serve for each task separately and are hard to transfer and scale on various tasks. For other tasks, it requires to redesign the similarity function or training signal. Compared with them, we introduce a unified training framework based on list-wise ranking and propose a single multi-task retriever UDR to serve for a wide range of tasks. Compared with EPR, besides UDR’s versatility on various tasks, UDR can incorporate LM’s feedback by ranking-based training in a more fine-grained way and receive more
crucial candidates’ signals by the iterative mining strategy. Cheng et al. (2023) propose CLAIF to enhance the sentence embedder by the gigantic language model’s feedback. Specifically, they use GPT-3 (Brown et al., 2020) 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, Li and Qiu (2023b) 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.

heading: 5 Conclusion
text: 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.
Limitations
We illustrate this paper’s limitations from the following three aspects:
1) Limited by the computational resources, we only train UDR from the initialization of “BERT base uncased” following EPR (Rubin et al., 2022). We regard explorations based on other competitive pre-trained models like RoBERTa (Liu et al., 2019) and DeBERTa (He et al., 2021) as future work.
2) Most of current dense demonstration retriev-
ers, including UDR, are black-box models. Although they lead to significantly better performance than BM25, how they find informative demonstrations is still unknown. Therefore, a better understanding of the principle of informative demonstration’s retrieval or an interpretable and transparent demonstration retriever may be the next stage of improving demonstration retrieval. Xu et al. (2023) propose a more explainable method, beyond-context learning, which first uses the language model to get training data’s next word probability distribution, then assigns test instances with labels of their nearest neighbors with similar next word’s probability distribution. We leave demonstration retrieval with better explainability as future work.
3) In the training stage we use LM to score candidates separately but in the inference stage LM is provided with a sequence of demonstrations. Although experimental results demonstrate UDR’s effectiveness, we think it is a promising direction to model the dependence between different demonstrations and leave it to future work.

heading: 6 Acknowledgements
text: This work was supported by the National Natural Science Foundation of China (No. 62236004 and No. 62022027) and Shenzhen City’s Science and Technology Plan Project (No. JSGG20210802153806021).

heading: A Task Overview
text: We show each task’s 1. input/output domain 2. statistics and evaluation metric 3. instruction, inference template and example cases in Table 9, 10 and 11, respectively. For the dataset which has publicly available test data, we use the test data for evaluation, like SST-2, SST-5, MTOP, etc. For the others like BREAK and SMCalFlow, we follow previous work (Rubin et al., 2022) and use the dev data for evaluation. For training efficiency, we manually limit the training examples of UDR. Specifically, for the classification task whose training set size is > 30000, we randomly sample a 30000 subset for UDR’s training. For the generation task whose training set size is > 100000, we randomly sample a 100000 subset for UDR’s training. In the pilot experiment, we find such a strategy will not cause significant performance degradation. At the inference stage, we use the full training set as demonstrations’ pool. Restricted by computational resources, we randomly sample a test set of 3000 samples for evaluation on these tasks: Amazon, Yelp, AGNews DBPedia and Yahoo.
B Implementation Details and Hyper-Parameters
We follow Rubin et al. (2022) to use GPT-Neo2.7B (Black et al., 2021) as the scoring LM and the inference LM for most experiments in the paper unless otherwise specified. Following EPR (Rubin et al., 2022) and DPR (Karpukhin et al., 2020), we
initialize Eq and Ed as two separate “BERT base uncased” encoders (Devlin et al., 2019). Thus the total number of parameters of UDR is about 220M. We use 8 NVIDIA A100s-80GB to train UDR for up to 30 epochs before iteratively mining candidates. And then we train UDR for 10 epochs at each iteration. The whole training pipeline including scoring candidates takes about 8 days. In the pilot experiment, we select the number of training epochs through the average performance on validation set on single-task SST-2, TREC, MTOP, Java code summarization, WikiAuto and DART. We set the number of iterations as 3. We follow EPR (Rubin et al., 2022) to set learning rate and batch size as 1e-4 and 128 and we use AdamW (Loshchilov and Hutter, 2019) as the optimizer. We list the overall hyper-parameters in Table 8. On each task, we use one specific template for scoring and inference (see Table 11). For fair comparison, we train DR-Target, EPR and UDR under the same hyperparameter and report their average performance under three random seeds.
The initialization of UDR’s candidates For classification and multi-choice tasks, we initialize candidates as those examples that have similar input with x by BM25. For generation tasks, similarly, we initialize candidates as those of similar targets with x’s, inspired by previous work (Rubin et al., 2022).
The Quantity of EPR’s Candidates Since UDR’s training needs to score iteratively mined candidates and thus has to score more candidates than EPR, we also run experiments on EPR with the same candidate quantities of UDR. But we find increasing the candidates of EPR instead slightly hurts its overall performance, which is consistent with its original paper (Rubin et al., 2022). Thus for EPR, we use the same number of candidates as its original paper.

heading: C Performance across varying inference templates
text: For UDR, we use one specific template when scoring candidates and here we evaluate UDR’s transferability across different inference templates on MR, Yahoo and Subj. The results are shown in Table 7. We can see that the performance gap across various inference templates is smaller than 1 point and this reflects UDR’s stability and transferability across different inference templates.

heading: D Potential Risk
text: Previous works have shown Large language models can have various kinds of bias (Bender et al., 2021). Since UDR is trained from the feedback of large language models, it can also contain such bias.
Task Family: Sentiment Classification
Task: SST-2 Task Instruction: Sentiment of the sentence: Inference Verbalizer: {great, terrible} Inference Template: Input: A three-hour cinema master class. It was terrible. A pretensions – and disposable story — sink the movie. It was great. · · · The movie ’s blatant derivativeness is one reason it ’s so lackluster. It was Output: terrible.
Task: SST-5 Task Instruction: Sentiment of the sentence: Inference Verbalizer: {great, good, okay, bad, terrible} Inference Template: Same as SST-2
Task: Amazon Task Instruction: Sentiment of the sentence: Inference Verbalizer: {great, good, okay, bad, terrible} Inference Template: Same as SST-2
Task: Yelp Task Instruction: Sentiment of the sentence: Inference Verbalizer: {great, good, okay, bad, terrible} Inference Template: Same as SST-2
Task: MR Task Instruction: Sentiment of the sentence: Inference Verbalizer: {great, terrible} Inference Template: Same as SST-2
Task: CR Task Instruction: Sentiment of the sentence: Inference Verbalizer: {great, terrible} Inference Template: Same as SST-2
Task Family: Topic Classification
Task: AGNews Task Instruction: Topic of the text: Inference Verbalizer: {World, Sports, Business, Technology} Inference Template: Input: ’LONDON, Oct 26 (AFP) - World oil prices will be driven down over the next two years due to there being enough crude to meet soaring demand, Claude Mandil, executive director of the International Energy Agency (IEA), said here Tuesday. Topic: Business. WASHINGTON - This year’s surge in energy prices is likely to have far less of an impact on the economy than the oil shocks of the 1970s, Federal Reserve Chairman Alan Greenspan said Friday. Greenspan predicted that the global economy will adjust to the recent surge in prices, which has seen oil topping \\$50 per barrel, by boosting energy exploration and production and by increasing fuel efficiency... Topic: World. · · · Oil demand is rising faster than predicted this year as OPEC pumps more low-quality oil in a failed bid to reduce record prices, according to International Energy Agency, an adviser to 26 industrialized nations. Topic: Output: Business.
Task: TREC Task Instruction: Topic of the question: Inference Verbalizer: {Description, Entity, Expression, Human, Location, Number} Inference Template: Same as AGNews
Task: DBPedia Task Instruction: Topic of the text: Inference Verbalizer: {Company, Educational Institution, Artist, Athlete, Office Holder, Mean of Transportation, Building, Natural Place, Village, Animal, Plant, Album, Film, Written Work} Inference Template: Same as AGNews
Task: Yahoo Task Instruction: Topic of the text: Inference Verbalizer: {Society & Culture, Science & Mathematics, Health, Education & Reference, Computers & Internet, Sports, Business & Finance, Entertainment & Music, Family & Relationships, Politics & Government} Inference Template: Same as AGNews
Task Family: Multi-Choice
Task: COPA Task Instruction: Answer the question based on the text. Inference Template: Input: I scratched my skin. What happened as a result? My itch went away. misplaced my wallet. What happened as a result? I retraced my steps. · · · I emptied my pockets. What happened as a result? Output: I retrieved a ticket stub.
Task: Cosmos QA Task Instruction: Answer the question based on the text. Inference Template: Same as COPA
Task: ComV Task Instruction: Which statement of the two is against common sense? Inference Template: Same as COPA
Task: ComE Task Instruction: Select the most corresponding reason why this statement is against common sense. Inference Template: Same as COPA
Task Family: NLI
Task: MNLI Task Instruction: Recognizing textual entailment between these 2 texts. Inference Verbalizer: {Entailment, Inconclusive, Contradiction} Inference Template: Input: uh-huh exactly not what color you are how old you are what if your male or female that would be wonderful i guess it’s kind of an ideal world though huh Based on that information, is the claim The world would be better if race and gender did not matter. People would get along much better "Entailment", "Contradiction", or "Inconclusive"? Answer: Inconclusive. uh-huh exactly not what color you are how old you are what if your male or female that would be wonderful i guess it’s kind of an ideal world though huh Based on that information, is the claim The world would be better if race and gender did not matter. "Entailment", "Contradiction", or "Inconclusive"? Answer: Entailment. · · · It’s that kind of world. Based on that information, is the claim The world is getting better. "Entailment", "Contradiction", or "Inconclusive"? Answer: Output: Inconclusive
Task: SNLI
Task Instruction: Recognizing textual entailment between these 2 texts. Inference Verbalizer: {Entailment, Inconclusive, Contradiction} Inference Template: Same as MNLI
Task: RTE Task Instruction: Recognizing textual entailment between these 2 texts. Inference Verbalizer: {True, False} Inference Template: Same as MNLI
Task Family: Subjective Classification
Task: Subj Task Instruction: Subjectivity of the sentence: Inference Verbalizer: {subjective, objective} Inference Template: Input: thirteen conversations about one thing lays out a narrative puzzle that interweaves individual stories , and , like a mobius strip , elliptically loops back to where it began . It’s subjective. a small gem of a movie that defies classification and is as thought-provoking as it is funny , scary and sad . It’s subjective. · · · smart and alert , thirteen conversations about one thing is a small gem . It’s Output: subjective
Task Family: Lingustic Acceptibility
Task: COLA Task Instruction: The grammaticality of this sentence: Inference Verbalizer: {not grammatical, grammatical} Inference Template: Input: The sea monster drowned the sailors. It is grammatical. He rode out the storm. It is grammatical. · · · The sailors rode the breeze clear of the rocks. It is Output: grammatical
Task Family: Semantic Parsing
Task: BREAK Task Instruction: Parse the sentence into logical form: Inference Template: Input: Parse the sentence into logical form: what flights are available from pittsburgh to boston on saturday 1#) return flights 2#) return #1 from pittsburgh 3#) return #2 to boston 4#) return #3 on saturday 5#) return #4 that are available Parse the sentence into logical form: what flights are available wednesday afternoon from denver to san francisco 1#) return flights 2#) return #1 from denver 3#) return #2 to san francisco 4#) return #3 on wednesday afternoon 5#) return #4 that are available · · · Parse the sentence into logical form: what flights are available tomorrow from denver to philadelphia 1#) Output: return flights ;return #1 from denver ;return #2 to philadelphia ;return #3 if available
Task: MTOP Task Instruction: Parse the sentence into logical form: Inference Template: Same as BREAK
Task: SMCalFlow Task Instruction: Parse the sentence into logical form: Inference Template: Same as BREAK
Task Family: Text Summarization
Task: CNN/DailyMail Task Instruction: Summarize the text: Inference Template: Input: Summarize the text: JERUSALEM (CNN) – Israel moved to defend itself in the face of international criticism Monday over its eviction of dozens of Palestinian families from a neighborhood of Jerusalem they have lived in for generations.· · · TL;DR: Israel incurs international criticism over eviction of Palestinian families . Two Jewish families moved in after evictions in East Jerusalem . Israeli spokesman says dispute is a legal one between private parties . Summarize the text: (CNN)The International Criminal Court opened an inquiry into attacks in Palestinian territories, paving the way for possible war crimes investigation against Israelis.· · · TL;DR: An inquiry allows the court to review evidence and determine whether to file charges . The U.S. calls for negotiations between Palestinian, Israeli officials . · · · Summarize the text: (CNN)The Palestinian Authority officially became the 123rd member of the International Criminal Court on Wednesday, a step that gives the court jurisdiction over alleged crimes in Palestinian territories.· · · TL;DR: Output: Membership gives the ICC jurisdiction over alleged crimes committed in Palestinian territories since last June . Israel and the United States opposed the move, which could open the door to war crimes investigations against Israelis .
Task: PubMed Task Instruction: Summarize the text: Inference Template: Same as CNN/DailyMail
Task: Reddit Task Instruction: Summarize the text: Inference Template: Same as CNN/DailyMail
Task Family: Commensense Generation
Task: Commen Gen Task Instruction: Generate a sentence using these concepts: Inference Template: Input: Generate a sentence using these concepts: counter, pizza, restaurant Generated sentence: Two men standing at counters assembling pizzas in a restaurant. Generate a sentence using these concepts: counter, restaurant, stand Generated sentence: A man stands behind the counter of a restaurant. · · · Generate a sentence using these concepts: field, look, stand Generated sentence: Output: The player stood in the field looking at the batter.
Task Family: Story Generation
Task: Roc Story Task Instruction: Beginning of the story: Inference Template: Input: Beginning of the story: Taylor had been up all night memorizing lines for the play. Rest of the story: She knew that most of the girls in her class would be auditioning too. She watched the other girls stumble over their lines. She took a deep breath before going up on stage with a smile. All her lines were delivered perfectly and she got the part. Beginning of the story: Gabby was proud to be given the lead role in the school play. Rest of the story: She had worked hard on her audition piece. She worked hard to memorize her lines for the play. When the show opened, she stood on the stage and took it all in. She loved the feeling of performing. · · · Beginning of the story: Natalie had auditioned for the lead in the school play.
Rest of the story: Output: She won the part and was super excited. She rehearsed for weeks and weeks. On opening night, she acted her little heart out. The play was a huge success!
Task: Roc Story Ending Task Instruction: An unfinished story: Inference Template: Same as ROC Story
Task Family: Code Summarization
Task: Go Task Instruction: Comment on the code. Inference Template: Input: Comment on the code. Code: func NewSTM ( c * v3 . Client , apply func ( STM ) error , so ... stmOption ) ( * v3 . TxnResponse , error ) {
opts := & stmOptions { ctx : c . Ctx ( ) } for _ , f := range so {
f ( opts ) } if len ( opts . prefetch ) != 0 {
f := apply apply = func ( s STM ) error { s . Get ( opts . prefetch ... ) return f ( s )
} } return runSTM ( mkSTM ( c , opts ) , apply )
}
Comment: RunContainer runs a fake Docker container Comment on the code. Code: func ( s Subnet ) EnsureDead ( ) ( err error ) {
defer errors . DeferredAnnotatef ( & err , " " , s ) if s . doc . Life == Dead {
return nil } ops := [ ] txn . Op { {
C : subnetsC , Id : s . doc . DocID , Update : bson . D { { " " , bson . D { { " " , Dead } } } } , Assert : isAliveDoc ,
} } txnErr := s . st . db ( ) . RunTransaction ( ops ) if txnErr == nil {
s . doc . Life = Dead return nil } return onAbort ( txnErr , subnetNotAliveErr ) }
Comment: EnsureDead sets the Life of the subnet to Dead if it s Alive . If the subnet is already Dead no error is returned . When the subnet is no longer Alive or already removed errNotAlive is returned . · · · Comment on the code. Code: func NewSTM ( c * v3 . Client , apply func ( STM ) error , so ... stmOption ) ( * v3 . TxnResponse , error ) {
opts := & stmOptions { ctx : c . Ctx ( ) } for _ , f := range so { f ( opts ) } if len ( opts . prefetch ) != 0 {
f := apply apply = func ( s STM ) error { s . Get ( opts . prefetch ... ) return f ( s ) } } return runSTM ( mkSTM ( c , opts ) , apply )
}
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: 3 A1. Did you describe the limitations of your work?
text: the "limitations" section after "conclusion" section

heading: 3 A2. Did you discuss any potential risks of your work?
text: Appendix D

heading: 3 A3. Do the abstract and introduction summarize the paper’s main claims?
text: in the end of abstract and introduction
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 7 Did you use or create scientific artifacts? Left blank.
B1. Did you cite the creators of artifacts you used? No response.
B2. Did you discuss the license or terms for use and / or distribution of any artifacts? No response.
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? No response.
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? No response.
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? No response.
B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. No response.
C 3 Did you run computational experiments? section Experiments

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? Appendix B
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? Appendix B

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Appendix B

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Appendix B
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: , 2021) as future work. 2) Most of current dense demonstration retriev-
ers, including UDR, are black-box models. Although they lead to significantly better performance than BM25, how they find informative demonstrations is still unknown. Therefore, a better understanding of the principle of informative demonstration’s retrieval or an interpretable and transparent demonstration retriever may be the next stage of improving demonstration retrieval. Xu et al. (2023) propose a more explainable method, beyond-context learning, which first uses the language model to get training data’s next word probability distribution, then assigns test instances with labels of their nearest neighbors with similar next word’s probability distribution. We leave demonstration retrieval with better explainability as future work. 3) In the training stage we use LM to score candidates separately but in the inference stage LM is provided with a sequence of demonstrations. Although experimental results demonstrate UDR’s effectiveness, we think it is a promising direction to model the dependence between different demonstrations and leave it to future work