Learning To Retrieve Prompts for In-Context Learning

In plain words

Large language models can learn a task from a few worked examples placed in the prompt, but which examples you pick changes the answer a lot. This paper trains a small search model that chooses those examples for each question.

The labels come from the language model itself. For each training example, the method tries several other training examples as the prompt and keeps the ones that make the correct answer most likely. A small retriever is then trained to find helpful examples from the question alone. At test time, one fast vector search returns the prompt. The method is tested on three datasets that map utterances to structured meaning representations, where it beats simple keyword search and earlier learned selectors.

TL;DR

• Learns a dense retriever (two BERT-base encoders, one for the question and one for example pairs) that picks in-context examples §3.2.
• Labels come from a scoring LM: for each training example, candidates are ranked by the probability the LM gives the gold output with that candidate as the prompt; top-k are positives, bottom-k hard negatives §3.1.
• Author-reported gains over the best baseline with GPT-Neo as both scorer and inference LM: BREAK 31.9 vs 26.0 (LF-EM), MTOP 64.2 vs 57.0 (EM), SMCalFlow 54.3 vs 51.4 (EM) tab. 2.
• Relevant when you pick few-shot examples for a fixed task and can afford a one-off labelling pass with an LM. Not established here: transfer across tasks beyond the three datasets, and multi-example dependence (left to future work §3.2).

Problem

Performance of in-context learning depends strongly on which training examples are in the prompt §1. Prior work either uses off-the-shelf similarity metrics or trains a retriever on surface similarity §1; the paper compares against BM25 and SBERT §4.2. Scoring every pair of training examples with the LM is quadratic in the training-set size and prohibitive §3.1. Choosing prompts at test time with the LM is slow, and the BruteForce variant scored below BM25 and SBERT on all three datasets §4.2 tab. 2. The method assumes labelled input-output pairs and access to a scoring LM §1.

Method

The method has two parts: labelling training data with a scoring LM §3.1, and training a dual-encoder retriever on those labels §3.2. Retrieval at test time uses maximum inner-product search over precomputed example vectors §3.2. Overview: fig. 1 (retrieval at test time) and fig. 2 (training data and training) §1.

Step 1: candidate sets (§3.1)

For each training pair (x, y), retrieve L candidates with an unsupervised retriever. The authors use BM25 over the output sequences y; using y instead of (x, y) gave slightly higher performance §3.1. The paper does not restate how the example itself is excluded from its own candidate list. The code drops the top-ranked hit find_bm25.py (see Pitfalls).

Step 2: score each candidate independently (§3.1)

Each candidate e is scored by the probability the scoring LM gives to the gold output y, given the candidate and the input x. Candidates are scored one at a time, not jointly. The score is an unnumbered equation in §3.1. The code computes the summed negative log-likelihood of the y tokens instead, where a lower value means a better candidate scorer.py scorer.py.

# Label one training example's candidates  [§3.1]
# Inputs: x and gold y of the training example; BM25 candidates; scorer g
# Output: positives (top-k) and hard negatives (bottom-k)

def label_candidates(x, y, candidates, scorer, k):
    # Score each candidate alone: P_g(y | candidate, x)  [§3.1]
    scores = []
    for candidate in candidates:
        scores.append(scorer.prob_of(y, given=(candidate, x)))
    # Sort from most to least helpful  [§3.1]
    ranked = sort_descending(candidates, scores)
    # Top-k become positives, bottom-k become hard negatives  [§3.1] [§4.2]
    positives = ranked[:k]
    hard_negatives = ranked[-k:]
    return positives, hard_negatives

Step 3: contrastive training (§3.2, eq. 1)

Both encoders are BERT-base and read the [CLS] vector §3.2. The input encoder E_X reads the question x. The prompt encoder E_P reads a candidate example, which is the concatenated input and output. Similarity is the inner product. Each training instance has one positive, one hard negative, and in-batch negatives: B-1 positives and B-1 hard negatives from the other instances in the mini-batch §3.2. The loss is the negative log-likelihood of the positive, over those candidates eq. 1. The paper says training follows DPR's contrastive procedure §3.2.

# One training step, DPR-style with in-batch negatives  [§3.2] [eq. 1]
# Shapes: queries (B, d); each example has a positive and a hard negative

def training_loss(batch):
    # Encode the questions with the input encoder  [§3.2]
    queries = encode_input(batch.questions)           # (B, d)
    # Encode positives and hard negatives with the prompt encoder  [§3.2]
    positives = encode_prompt(batch.positives)        # (B, d)
    hard_negatives = encode_prompt(batch.hard_negatives)  # (B, d)
    # Every query is scored against all B positives and B hard negatives  [§3.2]
    candidates = concat(positives, hard_negatives)    # (2B, d)
    scores = queries @ candidates.T                   # (B, 2B) inner products
    # Negative log-likelihood of each query's own positive  [eq. 1]
    return cross_entropy(scores, targets=arange(B))

Step 4: index and inference (§3.2)

Encode all training examples with E_P and store the vectors in FAISS §3.2. At test time, encode the test input with E_X and take the L most similar examples, sorted from highest to lowest inner product §3.2. Keep the longest prefix that fits the context budget C, where the prefix length counts the example tokens, the test input and the desired output length §3.2. Place the examples before the test input, with the highest-ranked one closest to it §3.2. Decode greedily §3.2.

# Build the test prompt and decode  [§3.2]
# Inputs: test input x_test; vector index of E_P(e); budget C; output budget
# Output: generated text

def answer(x_test, index, context_budget, output_budget):
    # Encode the test input, then take the L most similar examples  [§3.2]
    query = encode_input(x_test)
    ranked = index.top_k(query, L)            # highest inner product first
    # Keep the longest prefix that fits the context  [§3.2]
    prompt_examples = []
    for example in ranked:
        used = token_count(prompt_examples + [example])
        if used + token_count(x_test) + output_budget > context_budget:
            break
        prompt_examples.append(example)
    # The most similar example is placed closest to the test input  [§3.2]
    prompt = reversed(prompt_examples) + [x_test]
    return greedy_decode(inference_lm, prompt)

Hyperparameters

Results

Author-reported. Setup: development sets, GPT-Neo as scoring and inference LM, full context C=2,048 §4.3 tab. 2.

Development results (GPT-Neo scores and infers) tab. 2

Test results (GPT-Neo scores and infers) tab. 3

Proxy setting: GPT-Neo scores, larger LM infers. Sample of 1,000 dev examples tab. 5

One-shot setting (only the top example in the prompt), BREAK dev tab. 4

Prompt copying, dev, GPT-Neo LM-as-a-service tab. 7

Ablations over k and L on the GPT-Neo LM-as-a-service setup tab. 8, app. A:

Other claims in §4.4 and the appendix: the LM-Oracle is substantially higher than EPR §4.4; EPR is comparable to BM25-Oracle on BREAK and higher on MTOP and SMCalFlow §4.4; copying is mostly from highly ranked prompts §4.4 fig. 4; the t-SNE clusters capture both lexical and structural similarity §4.4 fig. 3. The distribution of the number of examples that fit the budget is in fig. 5 app. A.

Paper vs code

Code: official repository OhadRubin/EPR (first author's account), clone at commit 8dd8f1adb6a802e02c817d72add6469cb3374504, linked from the introduction §1. The repository root has no LICENSE file, so its license is unknown. The DPR/ subdirectory is CC-BY-NC-4.0 LICENSE. No code is quoted in this card; every code reference is a path and line citation.

Discrepancies found:

• Batch size. Paper: 120 app. A. The config file has 120 biencoder_local.yaml, but run.py defaults dpr_bs to 60 and passes it as train.batch_size, which overrides the config run.py run.py. Running run.py with defaults trains with batch 60.
• Learning rate. Paper: 1e-4 app. A. Code: 1.3416e-4 biencoder_local.yaml. The paper says it applied "minimal learning rate tuning" on the BREAK validation set app. A, so the code value may be the tuned one.
• Epochs. Paper: 30 app. A. The code default is 30 biencoder_local.yaml, but the README example uses dpr_epochs=120 README.md.
• Positive choice. Paper: positives are sampled from the top-k set §3.1. Code: the positive is always the first (highest-scoring) candidate, because shuffling of positives is off by default biencoder.py. The hard negative is drawn at random from the bottom-k each step biencoder.py biencoder_data.py.
• Random negatives. The EPR dataset class adds 50 random training examples as negatives biencoder_data.py, but other_negatives is 0 in the training config biencoder_local.yaml, so those negatives are not used in the default run.
• Context budget. Paper: C=2,048 tokens, and the chosen prefix must fit with the test input and output budget §3.2 §4.3. Code: the example budget defaults to 1,600 tokens few_shot_dsr.py few_shot_dsr.py, and run.py sets the generation length to 1,950 run.py. The inference config says 2,048 inference.yaml. run.py does not pass the example budget, so the 1,600 default applies.
• Score form. Paper: a probability §3.1. Code: the summed negative log-likelihood of the target tokens, sorted ascending so that the lowest loss ranks first scorer.py scorer.py. This ranks candidates the same way as the log-probability, but the code does not normalize by length.
• Scoring input. Code scores the sequence "candidate question, tab, candidate output, newline, test question, tab, test output" and counts loss only on the test output tokens scorer_dsr.py. The paper gives no delimiter format.
• Candidate retrieval query. Code queries BM25 with the output only, using whitespace tokens after splitting on semicolons break.py break.py. This matches §3.1.
• Inference retrieval depth. The paper says "the L most similar" without a number §3.2. The DPR retrieval default is 100 dense_retriever.yaml, and run.py does not override it.
• Prompt format. Code writes each example as question, tab, numbered steps ("1#) ..."), joined by newlines, then the test question, tab, and the postfix "1#)" break.py few_shot_dsr.py. The paper does not specify this format.
• Budget accounting. The code counts each example's own tokens and does not count the joining newlines or the postfix few_shot_dsr.py break.py. The real prompt can therefore exceed the budget by a few tokens.

No discrepancies found: encoders (BERT-base) §3.2, hard negatives per instance (1) §3.2, greedy decoding §3.2, and prompt order (highest-ranked example nearest the test input) §3.2.

Pitfalls

• Labelling cost. Each training example needs L scoring passes with a 2.7B LM. For BREAK's 44K training examples that is 44K × 50 = 2.2M forward passes at L=50 (arithmetic from §4.1 and §4.2, not a figure from the paper). Budget this before starting.
• Train and test use different signals. Candidates for training are retrieved with the gold output y, which is only available at training time §3.1. At test time, the query is the input alone §3.2. The retriever has to learn that mapping. The authors note that modelling the dependence between the chosen examples is left to future work §3.2.
• Self-exclusion is by position. Both the BM25 and the kNN finder drop the first hit when building training candidates find_bm25.py find_knn.py. If the example is not ranked first (ties, duplicate outputs), it stays in its own candidate list. The paper does not describe this step.
• Metric and split differences. BREAK dev uses LF-EM and the test set uses NEM §4.3. Dev and test numbers are not directly comparable.
• "8.5%" is absolute points against CBR. The one-shot text says EPR beats "the best baseline by 8.5%" §4.4. In tab. 4, 23.0 − 14.5 = 8.5 points against CBR. BM25 (15.2) is the higher baseline, so the gap to the best baseline is 7.8 points. The "5%" over BM25-Oracle is 5.0 points (23.0 − 18.0).
• The "2 to 8 points" summary is loose. The proxy text in §4.4 describes EPR's margin over the best baseline as roughly 2 to 8 points across datasets and models. tab. 5 gives gaps from 1.6 (SMCalFlow, Codex) to 8.2 (MTOP, GPT-Neo).
• Copy-analysis totals do not match the main table. tab. 7 lists total accuracy 32.3% for BREAK and 54.5% for SMCalFlow. tab. 2 gives EPR 31.9 and 54.3 for the same setting. The values 32.3 (BREAK) and 54.5 (SMCalFlow) are not the EPR results in tab. 2; the source of the difference is not stated.
• Single runs. No seeds, variances or confidence intervals are reported §4.3. The proxy setting uses 1,000 examples, evaluated once per model §4.3. Treat gaps of a few points as uncertain.
• Ablations are narrow. k and L were varied only on the GPT-Neo setting app. A. The paper calls performance "generally robust" app. A, but the table shows a 2.1-point drop on SMCalFlow at k=10 (52.2 vs 54.3) and a 3.3-point drop at L=100 (51.0 vs 54.3).
• Retrieval cost is one vector search per query. Inference retrieval uses a FAISS flat inner-product index over all training examples faiss_indexers.py. The paper does not discuss memory or index size.
• Bias. The appendix notes that EPR is trained on signals from large LMs and may inherit their biases app. A.
• Licensing for implementers. The code is unlicensed at the repository root, and the DPR subtree is non-commercial (CC-BY-NC-4.0). Reimplementing from the paper and this card avoids that, but copying code from the repository needs permission from the authors.

Lineage

• Builds on: DPR (Karpukhin et al., 2020) for dual-encoder contrastive training §3.2; BERT-base as the encoder §3.2; BM25 and SBERT for the candidate sets §3.1.
• Compared against: Random, SBERT, BM25, BruteForce (few-shot prompt selection by scoring 200 random examples with the inference LM) §4.2; DR-BM25 (the same retriever trained on BM25 labels) and CBR (Das et al., 2021, adapted by using F1 over output tokens) §4.2; BM25-Oracle and LM-Oracle §4.2.
• Superseded by: not stated in this card.