PaperlayerView exact revision on arXiv

Source-linked AI summary

Unified Demonstration Retriever for In-Context Learning

Xiaonan Li, Kai Lv, Hang Yan, Tianyang Lin, Wei Zhu, Yuan Ni, Guotong Xie, Xiaoling Wang, Xipeng Qiu

arXiv:2305.04320v2cs.CL

TL;DR

Demonstration retrieval matters because in-context learning performance depends heavily on the examples provided, while task-specific retrievers are difficult to transfer and scale. UDR addresses this with one multi-task retriever trained through unified LM-feedback ranking and iterative candidate mining, and it significantly outperforms baselines across broad task and scenario evaluations.

  • Problem

    In-context learning depends heavily on its demonstrations, while separately trained task-specific retrievers are difficult to transfer and scale across diverse tasks.

  • Method

    UDR is a single multi-task retriever trained with unified list-wise ranking from language-model feedback and iterative mining of positive candidates and hard negatives.

  • Results

    UDR significantly outperforms baselines across 30+ tasks, improving over Instructor by 10 points on classification tasks and 6.7 points on generation tasks.

  • Takeaways & Limitations

    UDR retrieves informative demonstrations across a wide range of tasks and remains effective with different language models, unseen datasets, and varying demonstration quantities.

  • Takeaways & Limitations

    During training, UDR scores candidates separately even though inference supplies a sequence of demonstrations, leaving dependence between demonstrations unmodeled.

Abstract

from arXiv · show

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 (\textbf{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.

1 Introduction

In-context learning uses a few demonstrations to predict a test input, but its performance depends strongly on which demonstrations are provided. UDR addresses the transfer and scaling limits of task-specific retrievers with one multi-task retriever trained through unified ranking signals.

  • Motivation: In-context learning lets a language model learn tasks from a few input-output demonstrations without parameter updates.
  • Motivation: Demonstration retrieval selects relevant training examples for a test input because demonstration choice strongly affects in-context learning performance.
  • Prior approaches: Off-the-shelf retrievers generalize across tasks but rely on textual or semantic similarity rather than task supervision.
  • Prior approaches: Task-specific retrievers can perform better but are difficult to transfer and scale because each task needs specialized signals and additional retriever parameters.
  • UDR: UDR is a single multi-task retriever that unifies task training signals as list-wise rankings from language-model feedback and mines high-quality candidates iteratively.
  • Results: Experiments cover 30+ tasks across 13 task families and multiple domains, with UDR significantly outperforming baselines and supporting different language-model sizes and unseen datasets.

2 Unified Demonstration Retriever

UDR is a single task-instruction-aware bi-encoder that retrieves demonstrations across tasks by unifying language-model feedback into multi-task list-wise ranking. Iterative candidate mining supplies hard negatives and valuable positives, while FAISS performs inference-time retrieval.

  • Unified Demonstration Retriever: UDR is a single multi-task model that retrieves demonstrations for a wide range of tasks in a unified formulation.It addresses the limited transferability and scalability of separately trained task-specific retrievers.
  • Bi-encoder with Task Instruction: The bi-encoder separately encodes query and candidate demonstrations after concatenating each example with its task instruction.Similarity between the two encodings determines retrieval relevance.
  • Learning from LM Feedback: UDR converts task-specific supervision into list-wise rankings by ordering candidate demonstrations according to the language model’s conditional likelihood of the ground-truth output.Generation tasks use sequence likelihood, while classification and multi-choice tasks use their corresponding scoring functions.
  • Loss Function: The ranking loss weights demonstration pairs by their rank differences, encouraging higher-ranked candidates to receive greater retriever similarity.A larger rank gap produces a larger pairwise weight and stronger optimization pressure.
  • Loss Function: Task-balanced batches combine list-wise ranking and in-batch negative losses, with task sampling adjusted by dataset size using α = 0.5.The combined objective integrates both losses through a predefined λ.
  • Iterative Candidate Mining: Iterative mining repeatedly updates candidates with the retriever, rescoring them with the language model to discover hard negatives and valuable positives.The procedure progressively explores the task’s training set while avoiding quadratic scoring of all example pairs.
  • Inference: At inference, FAISS searches task-specific training embeddings, using an input-length budget for generation and eight demonstrations for classification and multi-choice tasks.Retrieved demonstrations are sorted by similarity before greedy decoding.

3 Experiment

UDR is evaluated across a broad collection of NLP tasks and implementation settings. The experiments compare it with standard and task-specific retrieval baselines and examine template stability.

  • Experimental Settings: The training and evaluation suite covers about 40 tasks across 13 task families and multiple data domains.Families include sentiment, topic, multi-choice, natural-language inference, subjectivity, linguistic acceptability, semantic parsing, and text-to-text generation.
  • Experimental Settings: The primary experiments use GPT-Neo-2.7B as both the scoring and inference language model, unless otherwise specified.UDR’s transferability across different inference language models is also examined.
  • Experimental Analysis: UDR maintains stable performance when inference templates differ from the scoring template, indicating generality across templates.The alternative templates are reported in Appendix C.
  • Model Comparison: UDR is compared with Random, BM25, SBERT, Instructor, and prior task-specific retrievers using downstream in-context-learning performance.Most comparison results are implemented by the authors, while semantic-parsing results for several baselines come from prior work.

3.2 Main Results

UDR outperforms baselines across classification, multi-choice, and generation tasks, supporting unified demonstration retrieval across a broad task range. Its advantage is especially clear against Instructor, while summarization datasets show a different pattern.

  • Overall performance: UDR significantly outperforms baselines on most classification and generation tasks.The results are reported separately for classification and multi-choice tasks in Table 1 and generation tasks in Table 2.
  • Overall performance: UDR achieves better overall performance than DR-Target and EPR, indicating the value of unifying different tasks’ training signals.The comparison is presented as evidence for the unified ranking formulation.
  • Analysis: Random retrieval performs worst on most tasks, reflecting the necessity of retrieving high-quality relevant demonstrations.This pattern contrasts with the stronger performance of EPR and UDR, which use language-model feedback.
  • Analysis: CNN/DailyMail and Reddit summarization show similar performance across retrieval methods.The authors conjecture that language-model pretraining already provides summarization knowledge, allowing random demonstrations to activate that ability without example-specific information.

3.3 Analysis

UDR’s analyses show that its components support transfer across inference LMs and unseen datasets, while high-quality demonstrations reduce sensitivity to ordering and quantity.

  • Ablation Study: Removing list-wise ranking training or iterative candidate mining causes performance degradation, supporting the effectiveness of both components.The ablation study evaluates SMCalFlow, SST-2, and Java code summarization.
  • Transferability across Different LMs: UDR outperforms BM25 and EPR significantly across inference LMs ranging from GPT-Neo-1.3B to Text-Davinci-003 at 175B.The performance trends remain similar when GPT-Neo-2.7B is replaced as the inference LM.
  • Performance on Unseen Datasets: UDR significantly outperforms BM25 and SBERT by about 10 points on average on datasets unseen during training.The evaluation covers Twitter sentiment classification, QNLI, and Ruby and Java code summarization.
  • The Order of Demonstrations: Different orders of UDR-retrieved demonstrations have performance gaps within 1 point, versus fluctuations of up to tens of points for random examples.This indicates that high-quality demonstrations are less sensitive to ordering.
  • The Impact of Demonstration Quantity: UDR consistently outperforms baselines across demonstration quantities, with two demonstrations outperforming BM25 and EPR using eight.Generation-task performance benefits more from additional demonstrations than classification-task performance.

4 Related Work

Prior demonstration retrievers use generic similarity or task-specific supervision, whereas UDR unifies task feedback and applies finer-grained ranking with iterative candidate mining.

  • Off-the-shelf Retrievers: Off-the-shelf retrievers such as BM25 and KNN-based BERT methods retrieve textually or semantically similar demonstrations across tasks.These methods provide versatility but are not guided by task supervision.
  • UDR: UDR incorporates supervision from various tasks through unified LM feedback and list-wise ranking training.Its iterative mining strategy captures signals from more crucial candidates than EPR.

5 Conclusion

The paper presents UDR as a single multi-task retriever trained with unified LM-feedback ranking and iterative candidate mining, and reports broad gains across tasks and scenarios.

  • Conclusion: UDR is a single multi-task model for demonstration retrieval across a wide range of tasks.It replaces separately trained task-specific retrievers with one unified model.
  • Conclusion: UDR unifies task training signals through LM-feedback list-wise ranking and iteratively mines high-quality candidates.The framework is designed to incorporate signals from multiple tasks.
  • Conclusion: UDR significantly outperforms baselines on 30+ tasks and shows strong ability across different LMs, unseen datasets, and demonstration quantities.The reported LM range is 1.3B–175B.

Limitations

The paper identifies constraints in UDR’s training resources, interpretability, and modeling of demonstration dependence. It leaves stronger encoder exploration, explainable retrieval, and cross-demonstration dependence for future work.

  • Model capacity: UDR is trained only from the “BERT base uncased” initialization because of limited computational resources.The paper names RoBERTa and DeBERTa as alternative pretrained models for future exploration.
  • Interpretability: UDR and other dense demonstration retrievers are black-box models whose retrieval of informative demonstrations remains unknown.The paper leaves better understanding and interpretable, transparent retrievers as future work.
  • Demonstration dependence: Training scores candidates separately, whereas inference provides the language model with a sequence of demonstrations.The paper proposes modeling dependence between demonstrations as a future direction.

A Task Overview

The task overview organizes each task by its input/output domain, statistics, evaluation metric, instructions, inference template, and example cases. Evaluation uses public test data when available and development data otherwise, while training examples are manually capped for efficiency.

  • Evaluation data: Evaluation uses public test data for datasets such as SST-2, SST-5, and MTOP, while BREAK and SMCalFlow use development data following prior work.The evaluation split depends on whether publicly available test data exists.
  • Training data: For training efficiency, UDR manually limits the number of training examples, including a 30,000-example cap for large classification datasets.The cited passage continues with separate handling for generation tasks.

B Implementation Details and Hyper-Parameters

The implementation uses GPT-Neo-2.7B for candidate scoring and inference in most experiments, separate BERT encoders for UDR, and iterative candidate mining. Initial candidates come from BM25-style similarity, while EPR uses its original candidate count after larger sets slightly hurt performance.

  • Language models: GPT-Neo-2.7B serves as the scoring and inference language model for most experiments unless otherwise specified.The implementation follows Rubin et al. (2022) for this choice.
  • Hyper-parameters: Table 8 lists the hyper-parameters used in the experiments.The supplied passage identifies the table but does not enumerate its entries.
  • Model configuration: UDR initializes separate Eq and Ed BERT base uncased encoders, yielding about 220M total parameters.The encoders are initialized separately following EPR and DPR.
  • Candidate initialization: Classification and multi-choice candidates are initialized with BM25 input similarity, while generation candidates use similar targets.The generation-task initialization is inspired by prior work.
  • Baseline settings: Increasing EPR’s candidate quantity slightly hurts overall performance, so EPR retains the candidate count from its original paper.The comparison accounts for UDR’s need to score iteratively mined candidates.

C Performance across varying inference templates

UDR remains stable and transferable across the evaluated inference templates. On MR, Yahoo, and Subj, performance differences between templates are smaller than one point.

  • Template transferability: Performance gaps across inference templates are smaller than 1 point on MR, Yahoo, and Subj.The paper interprets this as evidence of UDR’s stability and transferability across different inference templates.

D Potential Risk

Because UDR is trained from large language model feedback, it may also contain the varied biases present in those models.

  • D Potential Risk: UDR may inherit biases from the large language models whose feedback provides its training signal.The passage identifies this as a potential risk of the training approach.
Loading 2305.04320v2…