SlowGuess/NAACL_2022_OCR huggingface dataset:
Ohad Rubin Jonathan Herzig Jonathan Berant
The Blavatnik School of Computer Science, Tel Aviv University
{ohad.rubin, jonathan.herzig, joberant}@cs.tau.ac.il
In-context learning is a recent paradigm in natural language understanding, where a large pretrained 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.
The striking language skills and world knowledge embedded in large pre-trained language models (LMs) (Devlin et al., 2019; Petroni et al., 2019; Raffel et al., 2020; Brown et al., 2020) 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 GPT-3 (Brown et al., 2020), but has quickly spread to other LMs (Lieber et al., 2021; Du et al., 2021; Rae et al., 2021).
An attractive property of in-context learning is that it provides a single model for multiple language understanding tasks. However, Liu et al.
Figure 1: An overview of prompt retrieval: Given a question from BREAK, one retrieves similar training examples from an index of the training set. The question and training examples (the prompt) are passed to an inference LM that decodes the output.
(2021a) showed that downstream performance can vary widely depending on the choice of in-context examples. This has sparked interest in prompt retrieval (see Fig. 1), where given a test instance, training examples are chosen for the prompt based on some similarity metric. Recent work has either used off-the-shelf unsupervised similarity metrics, or trained a prompt retriever to select examples based on surface similarity (Das et al., 2021).
In this work, we suggest to use language models themselves to label examples that can serve as good prompts, and train a prompt retriever from this signal. To train the retriever (see Fig. 2), we assume access to a training set of input-output pairs and to a scoring LM, i.e., a language model that will be used to score prompts. For each training example  $(x,y)$ , we go over other candidate training examples, and estimate the probability, according to the scoring LM, of  $y$  conditioned on  $x$  and the candidate prompt. We label training examples that lead to high probability as positive and low probability as negative and train a prompt retriever
Human Language Technologies, pages 2655 - 2671
Figure 2: An overview of our approach for training EPR. Given a training example, we use an unsupervised retriever  $R_{u}$  to obtain a set of candidates. We then pass the candidates to a scoring LM and label the top-  $k$  and the bottom-  $k$  as positive and negative examples, respectively. Last, we use this training data to train a dense retriever.
from this data using contrastive learning. We argue that using an LM for labeling examples is a better proxy for training a retriever compared to previously-proposed surface similarity heuristics. Importantly, when creating the training data, we have access to the gold label  $y$ , which can be used to obtain a high-quality set of candidate prompts. This leads to good positive examples and hard negative examples, which are beneficial for training with a contrastive objective.
Using a scoring LM to train an efficient retriever for a potentially different test time inference LM is beneficial in two scenarios. First, when the scoring LM is smaller than the inference LM and serves as a proxy for it. This results in cheap and efficient data generation for the retriever, accessible to a wide range of researchers. Second, our approach can be used even when the scoring and inference LMs are identical (e.g., both are GPT-3). This is beneficial when we do not have access to model parameters and can only use it as a service, an increasingly popular paradigm. In this case, we use the LM to train a light-weight retriever that is only tasked with learning a similarity function. More generally, given that the scale of LMs is likely to keep increasing in the foreseeable future, one can view our approach for Efficient Prompt Retrieval, or EPR, as a method for interfacing and learning to interact with large LMs.
We empirically test EPR on three structured sequence-to-sequence tasks, where input natural language utterances are mapped to a meaning representation: MTOP (Li et al., 2021) and SM-CALFLOW(Andreas et al., 2020), which focus on task-oriented dialogue, and BREAK (Wolfson et al.,
2020), a benchmark for mapping questions to a language-based meaning representation. We observe that EPR substantially improves performance compared to prior work on prompt retrieval. When the scoring LM and inference LM are identical (using GPT-NEO (Black et al., 2021)), performance compared to the best baseline improves from  $26\%$  to  $31.9\%$  on BREAK, from  $57\%$  to  $64.2\%$  on MTOP, and from  $51.4\%$  to  $54.3\%$  on SMCALFLOW. When using GPT-NEO as a proxy for larger LMs (GPT-J, GPT-3, and CODEX), we observe similar gains, where performance improves substantially in all cases.
To conclude, we propose an approach for retrieving training examples for in-context learning in large language models, and show it substantially outperforms prior methods. Given recent developments in scaling LMs, designing efficient methods for interacting with LMs is an important direction for future research.
2 Background: Prompt Retrieval
Problem setup Given a training set  $\mathcal{D} = \{(x_i, y_i)\}_{i=1}^n$  of input-output sequences, and a test example  $x_{\mathrm{test}}$ , our goal is to train a retriever,  $R(x_{\mathrm{test}}, \mathcal{D})$ , that will retrieve a subset of training examples  $\mathcal{P} = \{(x_j, y_j)\}_{j=1}^m \subset \mathcal{D}$ , where  $m \ll n$ . We succinctly refer to  $\mathcal{P}$  as the prompt.
Given an inference LM,  $g$ , a good prompt should lead to the target output sequence when the test example  $x_{\mathrm{test}}$  is concatenated to the prompt  $\mathcal{P}$  and passed as a prefix to  $g$ . Specifically, decoding from
1 Prompt often refers to a natural language template filled by an input example (Liu et al., 2021b), but here it denotes the sequence of training examples provided as input to the LM.
the LM  $g([\mathcal{P}; x_{\mathrm{test}}])$  should yield  $y_{\mathrm{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 Liu et al. (2021a) investigated the effect of different prompts on the performance of GPT-3 and demonstrated that the choice of in-context examples strongly affects downstream performance. They used an unsupervised sentence encoder to encode training examples, and retrieved for every test instance its nearest neighbors.
Das et al. (2021) 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.
Shin et al. (2021) 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.
3 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.
3.1 Generating the Training Data
Our approach relies on finding which training examples can serve as good prompts for other training examples. Scoring all pairs of training examples is quadratic in  $|\mathcal{D}|$ , and thus prohibitive. Hence, we present a method for choosing a set of candidate examples  $\bar{\mathcal{E}} \subset D$ , from which we will choose positive and negative examples for training. Importantly, since we are not at test time and are only generating data for training, we can use the target sequence  $y$  to retrieve a good set of candidates. This can be approached using a simple retrieval method, given that our goal is to retrieve examples that are similar to the input in terms of their output sequence,  $y$ .
To obtain a high-quality candidate set of training examples, we take advantage of an unsupervised retriever,  $\bar{\mathcal{E}} = R_u((x,y),\mathcal{D})$ . For the choice of the unsupervised retriever, we experiment with BM25 (Robertson and Zaragoza, 2009), a sparse retriever that relies on surface text similarity, and SBERT (Reimers and Gurevych, 2019), which is based on dense sentence encoding. For both, we experimented with passing the retriever the training pair  $(x,y)$  or the target sequence  $y$  only, and found that using  $y$  leads to slightly higher performance.
Scoring the candidate set Once we retrieve the set of candidates  $\bar{\mathcal{E}} = \{\bar{e}_1,\dots ,\bar{e}_L\}$  for a training example  $(x,y),^2$  we score each candidate  $\bar{e}_l\in \bar{\mathcal{E}}$  independently with a scoring LM,  $\hat{g}$  , which serves as a proxy for the inference LM,  $g$  . Specifically, the score for a candidate prompt is
$$
s (\bar {e} _ {l}) = \mathrm {P r o b} _ {\hat {g}} (y \mid \bar {e} _ {l}, x),
$$
which is the probability under the LM,  $\hat{g}$ , of the output sequence conditioned on the candidate prompt and input sequence. This indicates how helpful this candidate is for decoding the target (independent of all other candidates). We argue this score is a better proxy for the utility of a training example at inference time compared to prior approaches.
We apply this scoring function to all training examples, and define for each training example a set of positive examples  $\mathcal{E}_{\mathrm{pos}}$ , which includes the top- $k$  candidates in  $\bar{\mathcal{E}}$  according to  $s(\bar{e}_l)$ , and a set of negative examples  $\mathcal{E}_{\mathrm{neg}}$ , which includes the bottom- $k$  candidates in  $\bar{\mathcal{E}}$  according to  $s(\bar{e}_l)$ . This should lead to relevant positive examples, assuming that the set of candidates,  $\bar{\mathcal{E}}$  includes good prompt candidates and hard negatives, since all candidates have high similarity with  $(x,y)$  according to  $R_u(y,\mathcal{D})$ . With positive and negative examples at our disposal, we can now apply contrastive learning, which we describe next.
Training Our training procedure proceeds exactly like the contrastive learning procedure from DPR (Karpukhin et al., 2020). This procedure results in an input encoder  $E_{X}(\cdot)$ , which receives the sequence of input tokens,  $x$ , and a prompt encoder  $E_{P}(\cdot)$ , which receives a candidate prompt, namely, a concatenation of the tokens in an input-output pair. Both encoders are initialized with BERT-base
2We omit the dependence of  $\bar{\mathcal{E}}$  on  $(x,y)$  for simplicity.
(Devlin et al., 2019), and the output vector representation is given by the CLS token, as usual. The goal of training is to learn a similarity metric such that given a test example  $x_{\mathrm{test}}$ , it will be similar to training examples that lead to decoding of  $y_{\mathrm{test}}$ .
Our training instances are of the form  $\langle x_i, e_i^+, e_{i,1}^-, \ldots, e_{i,2B-1}^- \rangle$ . Where the positive example  $e_i^+$  is sampled from the set  $\mathcal{E}_{\mathrm{pos}}^{(i)}$ , and our negative examples consist of one hard negative example sampled from  $\mathcal{E}_{\mathrm{neg}}^{(i)}$ ,  $B - 1$  positive examples from the other instances in the same mini-batch, and the  $B - 1$  hard negatives from those instances. We define the similarity score between an input and an input-output pair to be the inner product  $\operatorname{sim}(x, e) = E_X(x)^\top E_P(e)$ . We can now define the typical contrastive learning objective and minimize for each example the negative log likelihood of the positive example:
$$
\begin{array}{l} L \left(x _ {i}, e _ {i} ^ {+}, e _ {i, 1} ^ {-}, \dots e _ {i, 2 B - 1} ^ {-}\right) \tag {1} \\ { = } { - \log \frac { e ^ { \mathrm { s i m } ( x _ { i } , e _ { i } ^ { + } ) } } { e ^ { \mathrm { s i m } ( x _ { i } , e _ { i } ^ { + } ) } } + \sum _ { j = 1 } ^ { 2 B - 1 } e ^ { \mathrm { s i m } ( x _ { i } , e _ { i , j } ^ { - } ) } . } \\ \end{array}
$$
An advantage of this approach is that for batch size  $B$  the effective batch size is of order  $B^2$ , with the in-batch negatives trick (Henderson et al., 2017).
Inference After training the input encoder and prompt encoder, we encode the entire set of training examples with  $E_P(\cdot)$  in a pre-processing step using FAISS (Johnson et al., 2017). At test time, given an input sequence,  $x_{\mathrm{test}}$ , we compute its encoding  $E_{X}(x_{\mathrm{test}})$ , and then use maximum innerproduct search over the training data to find the  $L$  most similar training examples, sorted by their inner product (from high to low):  $\mathcal{P} = (e_1,\dots ,e_L)$ . The final prompt  $\mathcal{P}'$  is determined by  $C$ , the maximal context size supported by the inference LM,  $g$ . Specifically,  $L^{\prime}\leq L$  is the largest  $L^{\prime}$  such  $\sum_{i = 1}^{L'}|e_i| + |x_{\mathrm{test}}| + |y'|\leq C$ , where  $|y^{\prime}|$  is the desired maximal length of the generated output. Finally, we return the output of greedy decoding on  $g([e_{L'};e_{L' - 1};\ldots ;e_1;x_{\mathrm{test}}])$ .
We note that while at training time we score each training example independently, at test time the language model observes a prompt, i.e., a sequence of examples. We leave modeling the dependence between different training examples to future work.
We now compare EPR to a wide range of unsupervised and supervised baselines, both when the
scoring LM,  $\hat{g}$ , is smaller than the inference LM,  $g$ , and when they are identical.
We focus on tasks that map utterances to meaning representations, where in-context examples can be used to learn the mapping from inputs to outputs. Examples from each dataset and the number of examples are in Table 1.
- BREAK (Wolfson et al., 2020): A dataset mapping complex natural language questions into a language-based meaning representation, where a question is decomposed into an ordered list of atomic steps. We use the low-level BREAK subset, containing  $44\mathrm{K} / 7\mathrm{K} / 8\mathrm{K}$  examples in its training/development/test sets.
- MTOP (Li et al., 2021): A semantic parsing dataset, focused on task-oriented dialogue, where commands are mapped to complex nested queries across 11 domains. Similar to past work (Pasupat et al., 2021), we use the English subset of MTOP, containing  $16\mathrm{K} / 2\mathrm{K} / 4\mathrm{K}$  examples in its training/development/test sets.
- SMCALFLOW (Andreas et al., 2020): A large English-language task-oriented dataset that covers tasks such as calendar, weather, places, and people. The meaning representation is a dataflow program, which includes API calls, function composition and complex constraints. SMCALFLOW includes 15K development set examples and 134K training examples, from which we sample a random set of 44K examples for training.
4.2 Baselines and Oracles
We consider the following unsupervised baselines, which are applied at test time only.
- RANDOM: we randomly sample examples from the training set  $\mathcal{D}$ .
- SBERT: We use Sentence Transformers, a library providing BERT-based sentence embeddings. Specifically, we use paraphrase/mpnet-base-v2, a 110M parameter model to encode the test utterance  $x_{\text {test }}$  and retrieve the examples with the most similar utterances as in-context examples.
- BM25: We use the classical sparse retrieval method BM25 (Robertson and Zaragoza, 2009), which is an extension of TF-IDF, to retrieve for
3https://www.sbert.net/index.html.
Dataset	Size	Utterance	Meaning Representation
BREAK	60K	There are more birds in the image on the right than in the image on the left.	1) return right image;2) return birds in #1;3) return number of #2;4) return left image;5) return birds in #46) return number of #5;7) return if #3 is higher than #6;
MTOP	22K	call Zoey's wife.	[IN:CREATE_CALL = [SL:CONTACT = [IN:GET_contact = [SL:CONTACT_RELATED = Zoey] [SL:TYPE_RELATION = wife]])]
SMCALFLOW	148K	Can you create me a new meeting on thursday morning?	(Yield (CreateCommitEventWrapper (CreatePreflightEventWrapper (Event.start_?(DateTimeConstraint(Morning)(NextDOW(Thursday))))))
Table 1: Examples from each of the datasets we evaluate on.
each test utterance  $x_{\mathrm{test}}$  the training examples with the most similar utterance.
- BRUTEFORCE: We apply the prompt selection method for few-shot semantic parsing from Shin et al. (2021). Given a test example  $x_{\mathrm{test}}$ , we sample 200 training examples. For each training example  $(x_i, y_i)$ , compute  $\operatorname{Prob}_g(x_{\mathrm{test}} \mid x_i)$ , and use the highest scoring examples for the prompt. Similar to us, this approach uses the inference LM to choose prompts. However, it does so at test time, which results in slow inference.
Next, we describe baselines that use the training set,  $\mathcal{D}$ , to train a prompt retriever. All supervised methods share the following template. First, a candidate set  $\bar{\mathcal{E}}$  of  $L = 50$  examples is retrieved with the unsupervised retriever  $R_{u}(y,\mathcal{D})$ . We use BM25 as an unsupervised retriever, since it outperformed SBERT (see §4.4). We then score each candidate prompt  $\bar{e}_l\in \bar{\mathcal{E}}$  with some scoring function, and label the top- $k$  prompts as positive examples and the bottom- $k$  as negative examples ( $k = 5$ ). Different supervised methods only differ in the scoring function itself.
- DR-BM25: Here, we use the original BM25 scores for labeling positive and negative examples and training a dense retriever.
- CASE-BASED REASONING (CBR): We adapt the scoring function from Das et al. (2021), which focused on knowledge-base question answering. They define the weight for a pair of logical forms to be the  $\mathrm{F}_1$  score between the two sets of relations appearing in those logical forms, and use this weight to softly label their data. Since in our setting we do not assume logical forms,
we define the score between two output sequence  $y_{i}$  and  $y_{j}$  to be the  $\mathrm{F}_1$  between the two sets of tokens in  $y_{i}$  and  $y_{j}$ , omitting stop words.
- EFFICIENT PROMPT RETRIEVAL (EPR): Our full approach from §3.
Last, we also consider two oracle models.
- BM25-ORACLE: We score test examples with BM25 using the gold output sequence  $R_{\mathrm{BM25}}(y_{\mathrm{test}}, \mathcal{D})$ . This provides an upper-bound on what can be learned by DR-BM25. EPR can potentially outperform this oracle, since its training signal goes beyond surface text similarity.
- LM-ORACLE: We use the procedure for labeling training data at test time. Given a test example  $(x_{\mathrm{test}}, y_{\mathrm{test}})$ , we first retrieve  $L$  candidate training examples with  $R_{\mathrm{BM25}}(y_{\mathrm{test}}, \mathcal{D})$ , we then sort the candidates with the scoring LM  $\hat{g}$ , estimating the probability of  $y_{\mathrm{test}}$  given  $x_{\mathrm{test}}$  and the candidate prompt. This provides an upper bound for EPR, since EPR is trained to emulate this behaviour.
Language models In this work, we only train a dense retriever, but use scoring and inference LMs. For our scoring LM,  $\hat{g}$ , we use GPT-NEO (Black et al., 2021), a 2.7B-parameter LM trained on The Pile (Gao et al., 2021), an 825 GB English text corpus, constructed from a wide range of high-quality resources.
In addition, we consider the following inference LMs: (a) GPT-J (Wang and Komatsuzaki, 2021): a 6B-parameter LM, also trained on The Pile. The advantage in this setup, is that GPT-J was trained on the same corpus as GPT-NEO. (b) GPT-3 (Brown et al., 2020): A 175B-parameter
4Results for  $k\in \{1,5,10\}$  and  $L\in \{50,100\}$  are in Appendix A.
Model	BREAK	MTOP	SMCALFLOW
Unsuper.	RANDOM	1.7	7.3	8.9
SBERT	21.6	48.7	43.6
BM25	26.0	52.9	46.1
BRUTEFORCE	7.7	18.1	11.1
Super.	DR-BM25	23.6	50.2	43.1
CBR	25.7	57.0	51.4
EPR (ours)	31.9	64.2	54.3
Oracle	BM25-ORACLE	32.3	58.9	47.3
LM-ORACLE	43.1	71.6	73.7
Table 2: Development results when GPT-NEO is the scoring and inference LM. Numbers for BREAK are LF-EM, and for MTOP and SMCALFLOW are EM.
Model	BREAK	MTOP
Unsuper.	BM25	17.6	49.0
Super.	CBR	18.4	57.5
EPR (ours)	23.9	64.4
model, trained mostly on a filtered subset of common crawl. (c) CODEX (Chen et al., 2021): A GPT-3 175B-parameter model finetuned on code from GitHub. Since our tasks involve mapping from utterances to programs or meaning representations, CODEX might potentially perform well at in-context learning.
For all LMs, we use a maximum context size of  $C = 2,048$  tokens.
Evaluation On BREAK, we evaluate performance on the development set with LF-EM (Hasson and Berant, 2021), which is a better metric compared to Normalized Exact Match (NEM), the official metric, as it measures whether two meaning representations are semantically equivalent. On the test set, we use NEM. On MTOP and SM-CALFLOW, we evaluate with Exact Match (EM), i.e., whether the string output by the inference LM is identical to the reference string.
We evaluate EPR in two settings: (a) LM-as-a-service, and (b) LM-as-a-proxy. In the first setting, we use GPT-NEO as both the scoring LM and inference LM. In this setting, we evaluate on the full development sets of BREAK, MTOP, and SMCALFLOW. In the latter setting, as we access GPT-3 and CODEX through a paid API, we sample a random subset of 1,000 development examples from each dataset and evaluate each model once on this subset.
Table 3: Test results where GPT-NEO is the scoring and inference LM. Numbers for BREAK are NEM, the official metric, and for MTOP are EM.
Model	One-shot	Full-context
Unsuper.	RANDOM	1.1	1.7
BM25	15.2	26.0
Super.	DR-BM25	14.1	23.6
CBR	14.5	25.7
EPR	23.0	31.9
Oracle	BM25-ORACLE	18.0	32.3
LM-ORACLE	33.3	43.1
ANYCORRECT-ORACLE	53.6	-
Table 4: Development results on BREAK with GPT-NEO in the one-shot setting. Numbers are LF-EM. Full-context is the corresponding numbers from Table 2.
LM-as-a-service Table 2 reports results where the scoring and inference LMs are identical. EPR substantially outperforms all other baselines. Specifically, when comparing to the best baseline, it improves performance from 26.0 to 31.9 on BREAK, from 57.0 to 64.2 on MTOP, and from 51.4 to 54.3 on SMCALFLOW. This shows that using the LM itself to label examples is an effective approach for obtaining a strong prompt retriever. Table 3 shows test results on BREAK and MTOP corroborating that EPR substantially improves performance compared to BM25 and CBR.