Source-linked AI summary
Learning To Retrieve Prompts for In-Context Learning
Ohad Rubin, Jonathan Herzig, Jonathan Berant
TL;DR
Prompt selection substantially affects in-context learning, motivating efficient retrieval methods that avoid relying only on surface similarity. EPR uses a language model to label useful prompts and trains a lightweight dense retriever, substantially improving results across three structured sequence-to-sequence tasks while leaving interactions among multiple prompt examples unmodeled.
Problem
Prompt choice can cause wide variation in downstream performance, while existing retrieval methods rely on surface similarity or expensive repeated language-model inference.
Method
EPR scores candidate prompts with a language model, labels high- and low-scoring examples, and trains a lightweight dense retriever from these labels.
Results
EPR substantially outperforms prior prompt-retrieval methods and baselines across BREAK, MTOP, and SMCALFLOW, including improvements from 26% to 31.9% on BREAK, 57% to 64.2% on MTOP, and 51.4% to 54.3% on SMCALFLOW.
Takeaways & Limitations
EPR provides an efficient way to retrieve prompts when large language-model weights are inaccessible or when a smaller scoring model serves as a proxy.
Takeaways & Limitations
During training, examples are scored independently, while test-time inference uses a sequence of examples whose dependence is not modeled.
Abstract
from arXiv · showhide
In-context learning is a recent paradigm in natural language understanding, where a large pre-trained language model (LM) observes a test instance and a few training examples as its input, and directly decodes the output without any update to its parameters. However, performance has been shown to strongly depend on the selected training examples (termed prompt). In this work, we propose an efficient method for retrieving prompts for in-context learning using annotated data and a LM. Given an input-output pair, we estimate the probability of the output given the input and a candidate training example as the prompt, and label training examples as positive or negative based on this probability. We then train an efficient dense retriever from this data, which is used to retrieve training examples as prompts at test time. We evaluate our approach on three sequence-to-sequence tasks where language utterances are mapped to meaning representations, and find that it substantially outperforms prior work and multiple baselines across the board.
1 Introduction
In-context learning depends strongly on which training examples are provided as prompts. EPR uses language-model scoring to train an efficient prompt retriever and substantially improves performance across three structured sequence-to-sequence tasks.
- In-context learning gives a language model training examples and a test instance, then decodes the output without updating model parameters.
- Prompt retrieval addresses widely varying downstream performance by selecting training examples for each test instance using similarity or learned retrieval.
- EPR labels candidate examples using the probability assigned by a scoring LM to the correct output, then trains a dense retriever with contrastive learning.
- EPR uses gold outputs during training to obtain strong positive and hard negative prompt examples, improving the retriever’s contrastive supervision.
- EPR supports lightweight retrieval when scoring and inference LMs are identical or when a smaller scoring LM proxies for a larger inference LM.
- 26% to 31.9% on BREAK, 57% to 64.2% on MTOP, and 51.4% to 54.3% on SMCALFLOW versus the best baseline when scoring and inference use GPT-NEO.
2 Background: Prompt Retrieval
Prompt retrieval selects a small subset of training examples to serve as the prompt for an inference language model. Prior approaches use nearest-neighbor or surface-similarity retrieval, while other methods rely on expensive repeated language-model inference.
- A retriever maps a test example and training set to a prompt containing m training examples, where m is much smaller than the training-set size n.
- The prompt is concatenated with the test input, and the inference LM should decode the corresponding target output.
- Here, prompt denotes the sequence of training examples supplied to the LM, rather than a natural-language template filled by an input example.
- Earlier work retrieved nearest neighbors with an unsupervised sentence encoder after showing that prompt choice strongly affects GPT-3 performance.
- Other approaches train task-specific retrievers from surface similarity or repeatedly run GPT-3 over many sampled examples, making inference expensive.
3 Efficient Prompt Retriever
EPR generates prompt-training labels with language-model scores, trains a dense retriever using contrastive learning, and retrieves prompts efficiently at inference time. Its procedure uses target outputs during training-data construction, approximate proxy scoring, and maximum-inner-product search at test time.
- 3.1 Generating the Training Data: EPR first retrieves candidate training examples whose output sequences resemble the target output, avoiding quadratic scoring over all training-example pairs.The candidate set can be produced with BM25 or SBERT, using either the full training pair or the target sequence alone.
- 3.1 Generating the Training Data: Training-time candidate construction may use the target sequence y, because target outputs are available when generating retriever supervision but not at test time.This distinction enables high-quality candidate prompts during data generation while the trained retriever operates from test inputs.
- 3.1 Generating the Training Data: A scoring language model independently assigns each candidate a probability for the target output conditioned on the candidate prompt and input.This score estimates how helpful each candidate is for decoding the target and serves as a proxy for the inference language model.
- 3.1 Generating the Training Data: The top-k scored candidates become positive examples and the bottom-k become negative examples, yielding relevant positives and hard negatives for retriever training.The candidate set is designed to contain both good prompts and difficult alternatives with high similarity to the training instance.
- 3.2 Training and Inference: The dense retriever encodes inputs and concatenated input-output prompts with separate BERT-initialized encoders, learning similarity through an inner-product contrastive objective.Training samples positives and hard negatives, while in-batch negatives increase the effective batch size to order B^2.
- 3.2 Training and Inference: At inference, FAISS indexes prompt encodings; the input encoder retrieves the L highest-inner-product training examples for the test input.The retrieved sequence is constrained by a maximum total token budget and is passed to the inference language model for greedy decoding.
- 3.2 Training and Inference: The method scores training examples independently during training, whereas test-time prompts contain sequences of examples whose dependence is not modeled.The authors leave modeling interactions between different prompt examples to future work.
4 Experimental Results
EPR is evaluated on three utterance-to-meaning-representation datasets against unsupervised and supervised baselines in both shared-model and proxy-LM settings. It substantially improves retrieval and decoding performance, while analyses indicate that its gains reflect structural as well as lexical matching and support both copying and generalization.
- Experimental setup: EPR is evaluated on BREAK, MTOP, and SMCALFLOW against random, SBERT, BM25, and supervised retrieval baselines.The datasets map natural-language utterances to meaning representations, including ordered steps, nested queries, and dataflow programs.
- LM-as-a-service: 31.9 versus 26.0 on BREAK, 64.2 versus 57.0 on MTOP, and 54.3 versus 51.4 on SMCALFLOW when GPT-NEO scores and performs inference.These comparisons are against the best baseline in the LM-as-a-service setting.
- Oracle comparisons: EPR outperforms BM25-ORACLE on MTOP and SMCALFLOW and is comparable on BREAK, despite BM25-ORACLE accessing the test output sequence.LM-ORACLE performs substantially higher, indicating that the scoring LM supplies strong supervision for retriever training.
- Retrieval analysis: EPR captures lexical and structural similarity, and its retrieved examples can support correct outputs when competing retrieval examples lack matching operations or concepts.A BREAK example links EPR’s correct decoding to retrieved examples involving argmax or argmin and “a code,” unlike the compared CBR example.
- Prompt copying: 84.5% of MTOP examples were abstractly copied, with 71.6% EM on that subset versus 64.2% overall, while non-copying cases still show some generalization.Copying is substantially more frequent in MTOP and SMCALFLOW than in BREAK, and accuracy is higher when copying occurs.
5 Related Work
Related work situates EPR at the intersection of in-context learning, retrieval, prompting, and supervised-model augmentation. Prior approaches retrieve examples using similarity or use retrieval to influence model inputs or outputs.
- In-context learning: In-context learning conditions next-word prediction on prompts, with prior accounts emphasizing linear separability or shared latent concepts.
- Retrieval: Dense retrieval research has expanded from open-domain question answering to knowledge-intensive tasks such as fact verification.
- Retrieval: Retrieval has been applied to semantic parsing, language modeling, and machine translation, including methods that control generated outputs.
- Prompts: Prompting includes in-context training examples as well as manually designed natural-language patterns for casting tasks as language modeling.
- Prompt retrieval for supervised models: Supervised models can use BM25-retrieved training examples as additional input, and fine-tuning with them improves summarization and question answering.
6 Conclusions
The paper presents EPR as a lightweight way to learn prompt retrieval using language models as scoring functions. It reports substantial improvements on three challenging tasks and frames efficient interaction with large models as an important direction.
- EPR learns to retrieve good in-context prompts by using language models as scoring functions, enabling a lightweight retriever.
- EPR substantially improves performance on three challenging tasks.
- Developing effective approaches for interacting with large language models is important as these models become prominent in language understanding.
A Appendix
The appendix documents dynamic prompt counts, hyperparameter robustness, training details, examples, visualizations, and a bias-related risk assessment. It also reports that EPR may inherit biases from the large language models providing its training signal.
- Distribution of the number of in-context examples: The dynamic selection procedure produces different numbers of in-context examples across test instances under a 2,048-token context limit.
- Effect of hyperparameters: Performance is generally robust to k, the number of positive or negative prompt labels, and L, the number of prompts retrieved by the unsupervised retriever.
- Training details: EPR training uses Adam with batch size 120, learning rate 1e-4, eight RTX 3090 GPUs, and 30 epochs.
- Risk assessment: Because EPR trains on signals from large language models, it might also exhibit their various biases.
- Additional examples: The appendix includes examples where EPR is correct while CBR is incorrect, alongside the top-three prompts retrieved by each method.
- Visualization: Tables 12–16 provide examples of clusters from EPR’s t-SNE projection on BREAK.