Learning To Retrieve Prompts for In-Context Learning (2022)

Abstract: 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 prompts). In this work, we propose an efficient method for retrieving prompts for in-context learning using annotated data and an 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.

## introduction the striking language skills and world knowledge embedded in large pre-trained language models (lms) @xcite @xcite have recently led to in-context learning, a new paradigm in natural language understanding. under this paradigm, a language model is given a prompt, which typically contains a few training examples, as well as a test instance as input, and generates the output for the test instance directly, without any update to its parameters. this approach was first introduced in @xcite , but has quickly spread to other lms @xcite @xcite . an attractive property of in-context learning is that it provides a single model for multiple language understanding tasks. however, liu et al. ## background: prompt retrieval problem setup given a training set @xmath0, y i )} n @xmath1, and a test example x test , our goal is to train a retriever, r(x test , d), that will retrieve a subset of training examples @xmath2, y j )} m @xmath3, where m ≪ n. we succinctly refer to p as the prompt. 1 given an inference lm, g, a good prompt should lead to the target output sequence when the test example x test is concatenated to the prompt p and passed as a prefix to g. specifically, decoding from the lm g([p; x test ]) should yield y test . in this work, we focus on structured tasks, such as semantic parsing, where x is a natural language utterance and y is a meaning representation for that utterance. prior work @xcite investigated the effect of different prompts on the performance of gpt-3 and demonstrated that the choice of incontext examples strongly affects downstream performance. they used an unsupervised sentence encoder to encode training examples, and retrieved for every test instance its nearest neighbors. @xcite trained a supervised prompt retriever for knowledge-base question answering. the retriever was trained with supervision that is tailored for knowledge-base queries, and relies on surface similarity between formal queries. conversely, our approach takes advantage of the generative lm itself and is thus more general. @xcite used gpt-3 to select examples for the prompt for few-shot semantic parsing. however, rather than training a retriever, they randomly sample a large set of utterance-program pairs from the training set, and choose those that are similar to the target instance question according to gpt-3. this results in an expensive inference procedure, where gpt-3 is run hundreds of times for each test instance, unlike our approach, which is based on a light-weight sub-linear retriever. ## efficient prompt retriever we now describe our method for training epr, an efficient prompt retriever for in-context learning. we first describe how to generate labeled data ( §3.1), and then how to use the training data for training and inference ( §3.2). fig. 2 provides an overview of the training procedure. ## experimental results we now compare epr to a wide range of unsupervised and supervised baselines, both when the scoring lm, ĝ, is smaller than the inference lm, g, and when they are identical. ## related work in-context learning our understanding of incontext learning has grown substantially recently. @xcite suggests that by conditioning on a prompt, the task of predicting the next word approaches linear separability. @xcite suggests that in-context learning occurs when the model infers a shared latent concept between examples in a prompt. @xcite present a pre-training scheme theoretically motivated by the bias of in-context learning, that gives significant improvements. recently, @xcite showed that the model does not rely on the ground truth input-label mapping provided in the demonstrations as much as previously thought. retrieval research on training dense retrievers has skyrocketed recently, propelled by interest in open-domain question answering @xcite @xcite @xcite . work on retrieval-based methods has also spread more widely to other knowledge-intensive tasks @xcite , e.g., fact verification @xcite . similar to us, @xcite proposed to use retrieval in semantic parsing. however, they focus on controlling the output generated by a model. retrieval methods have also been successfully used in language modeling @xcite @xcite and machine translation @xcite . prompts developing methods for interacting with lms and extracting desired behaviours has attracted considerable attention, under the umbrella term prompting. in this work, prompts are a set of in-context training examples, but substantial effort has also been devoted to casting natural language tasks as language modeling by phrasing the target task in natural language (see survey in @xcite ). such approaches include prompt engineering through manual patterns @xcite , decoding methods @xcite @xcite , and methods for extracting either hard @xcite or soft @xcite @xcite prompts automatically. prompt retrieval for supervised models in parallel to this work, adding training examples as additional input has been shown to be useful for supervised models as well. @xcite and @xcite used bm25 to retrieve and augment the input with similar examples from the training set. fine-tuning the model with the additional inputs improved performance on tasks such as summarization and question answering. such methods can also potentially benefit from a stronger retriever. ## conclusions large pre-trained lms are becoming an inseparable part of the natural language understanding ecosystem. however, accessing their weights or updating them can be prohibitive for many researchers. in this work, we propose epr, a method for learning to retrieve good prompts for in-context learning, by using language models themselves as the scoring function. this allows us to train a light-weight retriever and substantially improve performance on three challenging tasks. more broadly, given that large lms models are likely to play a prominent role in developing language understanding models, it is important to develop approaches for interacting with such models effectively. epr can be viewed as a step in this direction. council (erc) under the european union horizons 2020 research and innovation programme (grant erc delphi 802800). this work was completed in partial fulfillment for the ph.d degree of ohad rubin.