In-Context Retrieval-Augmented Language Models

In plain words

A language model writes the next words by looking at the text before them. This paper gives it a cheap way to look things up. Before the model writes, a search engine finds a passage that matches the last few words, and that passage is pasted in front of the text. The model itself stays exactly as it was. Nothing is retrained and no new parts are added, so the same trick works on a model you can only reach through an API.

The search can be tuned in a few ways. Searching often, using the last few dozen words as the query, and picking the best of several candidate passages all help the model predict what comes next. The authors report that a small model helped this way predicts text about as well as a much bigger model without the help. The results are measured on text prediction and a few question-answering sets, so they say little about other tasks.

TL;DR

In-Context RALM leaves a pretrained LM unchanged and prepends one retrieved document to its input §3.1. With BM25, a 32-token query and retrieval every 4 tokens, the authors report perplexity gains that match a model 2-3 times larger on five corpora §5 tab. 1. Reranking the top 16 BM25 hits with an LM or a trained RoBERTa reranker adds further gains §6 tab. 1. All numbers are author-reported on their own setups.

Problem

Retrieval-augmented language models usually change the LM. RETRO, for example, adds cross-attention layers and needs further training §1 §2. Such changes complicate deployment, and they are impossible when the LM is only available through an API §1 §8. The paper asks whether an off-the-shelf LM can use retrieved documents if they are simply placed in its context, and which retrieval choices matter for language modelling §1.

Method

Overview. The LM defines the usual next-token distribution over tokens eq. 1. In-Context RALM replaces the conditioning prefix by the retrieved document followed by the prefix, and does not change the LM weights §3.1 eq. 2. The paper uses exactly one document per prediction §3.1. It states that when the document plus prefix is longer than the model's limit, tokens are removed from the beginning of the prefix until the input fits §3.1.

1. Retrieval stride s. Retrieval runs once every s generated tokens rather than at every step §3.2 eq. 3. Each retrieval costs one search plus re-encoding of the prefix, so stride is a runtime/quality trade-off §3.2.
2. Query length l. The query is the last l tokens of the prefix that ends at the current stride boundary, not the whole prefix §3.2 eq. 4.
3. Retriever. The paper tests BM25 (sparse) and three dense retrievers: a frozen BERT-base with mean pooling, Contriever and Spider §4.2. BM25 gives the best perplexity §5.1.
4. Document. Passages are non-overlapping 100-word chunks from the retrieval corpus §4.3. Retrieved text is truncated to 256 tokens before it enters the input §4.3.
5. Reranking, zero-shot. Among the top k = 16 BM25 documents, pick the one that maximises the LM's probability of the last s' prefix tokens, given that document and the prefix before them §6.1 eq. 5 eq. 6. The scoring LM does not have to be the generating LM §6.1.
6. Reranking, predictive. A RoBERTa-base classifier scores each of the top k documents against the prefix §6.2. The scores are softmax-normalised over k eq. 7 and the best document is chosen eq. 8. The reranker is trained with a loss over the frozen LM's likelihood of the next stride eq. 9. Only the reranker is trained; the LM weights stay frozen §6.2.
7. Training data for the predictive reranker. Sample a prefix from the training data, run BM25 on its query, and compute the LM likelihood of the next stride given each retrieved document §6.2. The paper uses 300,000 such examples from WikiText-103 train, with GPT-2 110M as the LM §6.2.

# In-Context RALM decoding loop  [§3.1] [§3.2] [eq. 2] [eq. 4]
# Input:  prefix tokens, LM, retriever, stride s, query length l, limit n
# Output: tokens after generation

def in_context_ralm(prefix, lm, retriever, s, l, n):
    # generation proceeds in blocks of s tokens  [§3.2]
    tokens = list(prefix)
    while not finished(tokens):
        # query: last l tokens up to the current stride boundary  [§3.2]
        query = tokens[-l:]
        # one document from the corpus, top-scored by BM25  [§4.3]
        document = retriever.top_document(query)
        # the document is placed before the prefix, not in place of it  [§3.1]
        room = n - len(document)
        context = document + tokens[-room:]
        # generate the next s tokens with unchanged LM weights  [§3.1]
        tokens += lm.generate(context, num_tokens=s)
    return tokens

# Zero-shot reranking of k retrieved documents  [§6.1] [eq. 6]
# Input:  prefix tokens, k documents, frozen LM, scoring length s'
# Output: index of the chosen document

def rerank_zero_shot(prefix, documents, lm, s_prime):
    # target: the last s' prefix tokens, which are known at test time  [§6.1]
    target = prefix[-s_prime:]
    # condition on each document plus the prefix before the target  [§6.1]
    context_prefix = prefix[:-s_prime]
    scores = []
    for document in documents:
        scores.append(log_prob(lm, target, document + context_prefix))
    return argmax(scores)

# Predictive reranker loss  [§6.2] [eq. 7] [eq. 9]
# Input:  prefix, k documents, next-stride tokens y, reranker f, frozen LM
# Output: scalar loss for one training example

def predictive_loss(prefix, documents, next_tokens, reranker, lm):
    # relevance scores f(prefix, d_i), normalised over the k documents  [eq. 7]
    scores = [reranker(prefix, d) for d in documents]
    weights = softmax(scores)
    # LM probability of the next stride given each document  [§6.2]
    log_likes = [
        log_prob(lm, next_tokens, d + prefix)
        for d in documents
    ]
    # the loss is -log of a sum of products; compute it in log space  [eq. 9]
    return -logsumexp([log(w) + ll for w, ll in zip(weights, log_likes)])

Hyperparameters

Results

Author-reported. Setup: WikiText-103 test (word-level perplexity), RealNews, ArXiv, Stack Exchange and FreeLaw from The Pile (token-level perplexity), with GPT-2 models tab. 1. Lower is better. "BM25, first" is the top BM25 document §5; "zero-shot" and "predictive" are the rerankers in §6.

Source: tab. 1. Dashes mean the paper reports no number for that cell tab. 1.

Other results tab. 2 tab. 3 tab. 4 tab. 5:

Headline claims, as the authors state them:

• On five corpora, In-Context RALM with BM25 matches the perplexity gain of a 2-3x larger GPT-2 model §1 §5.
• A 345M GPT-2 with BM25 In-Context RALM beats a 762M GPT-2 without retrieval; with the trained reranker it beats a 1.5B GPT-2 §1.
• A 6.7B OPT with In-Context RALM matches a 66B OPT on WikiText-103 and RealNews §1 §5. The OPT values are in a figure, not in the text, so they are not transcribed here.
• LLaMA-13B zero-shot on NQ rises from 12.0% to 31.0% exact match, and on TriviaQA from 54.8% to 60.1%, with two DPR documents §7 tab. 4.

Paper vs code

Code: https://github.com/AI21Labs/in-context-ralm at commit d39d329673c007fc9a5dc3b61a22f6bc853df84f, Apache-2.0. The repository is linked from the paper's abstract, and the owner is the authors' organisation.

The released code covers BM25 retrieval, perplexity evaluation with optional zero-shot and oracle reranking, and the QA evaluation. It does not contain:

• the dense retrievers (BERT, Contriever, Spider); selecting dense raises an error retrieval_factory.py;
• the predictive reranker, its training code or its training data pipeline; the README only lists logprob and oracle reranking README.md.

Discrepancies and mismatches found:

• Stride default. The paper uses s = 4 §5.2. The retrieval script defaults to 32 prepare_retrieval_data.py, while the evaluation script defaults to 4 eval_lm.py. The README passes 4 to both README.md README.md. The stride must match in both scripts, because the evaluation asserts the stored spans eval_lm.py.
• Documents per query. The paper uses one document for LM experiments §3.1. The retrieval script defaults to one prepare_retrieval_data.py, but the README builds 16 for reranking README.md.
• How the document enters the window. The paper says tokens are removed from the start of the prefix until the input fits §3.1. The code uses a fixed window of max_length tokens, writes the document's tokens over the first positions of the window, and so drops the same number of prefix tokens eval_lm.py. The effect matches the paper's description; an implementation that concatenates first and then trims from the left would drop document tokens instead.
• Passage truncation. Passages are truncated to 256 tokens §4.3. The code uses the same default eval_lm.py eval_lm.py.
• Query length. The paper's l = 32 for BM25 §5.3. The code default is 32 retrieval_factory.py. The query is decoded text, so the search sees re-decoded tokens, not the original IDs sparse_retrieval.py.
• Reranking window. The zero-shot reranker scores the last 16 prefix tokens, matching s' = 16 §6.1 eval_lm.py. The code always uses the evaluated LM for reranking eval_lm.py; the paper's cheaper GPT-2 110M reranker tab. 3 cannot be reproduced with the released CLI.
• Word perplexity. The code's word count is the number of spaces in the raw string eval_lm.py. The paper says "word perplexity" tab. 1 without giving a counting rule, so word-level numbers depend on this choice.
• Question normalisation (QA). The code lowercases the first letter of the question and adds a question mark eval_qa.py. The paper's example prompt capitalises "Who" app. C. Expect small differences in QA results from this.
• Corpus date. The paper uses the 20 Dec 2018 Wikipedia dump §4.3. The code downloads a prebuilt index named wikipedia-dpr retrieval_factory.py; it does not check the dump date. The paper's rules for removing WikiText articles are implemented by a title list sparse_retrieval.py wikitext103_forbidden_titles.txt.
• QA prompt. The paper's open-book prompt gives documents as "title, blank line, text" with a sentence before the question app. C. The code matches that for two documents eval_qa.py.

No discrepancies found in: the LM-side reranking formula §6.1 (the code computes the chosen document as the argmin of summed negative log-likelihood over the window eval_lm.py); the top-k of 16 in the reranking experiments §6 README.md; the query formula of l tokens up to the stride boundary §3.2 sparse_retrieval.py.

Pitfalls

• The predictive reranker is not reproducible from this repository. Its data generation, training loop and optimizer are missing, and the paper leaves the optimizer, warmup, weight decay and dropout unstated §6.2. Any implementation has to fill these in, and the results can change with them.
• The loss needs log-space arithmetic. The loss is the negative log of a sum over k documents of a softmax weight times an LM probability over a whole stride §6.2 eq. 9. Multiplying the raw probabilities underflows. Use log-sum-exp.
• Retrieval every few tokens is expensive. Each retrieval runs a search and re-encodes the prefix §3.2. The paper states this runtime cost but reports no timings in the text. Reranking with k = 16 needs k forward passes §6.1.
• Stride and query length are separate knobs. The paper says that tying them (s = l) degrades LM performance §3.2. Prior work such as RETRO couples them §3.2.
• Query length is retriever-specific. BM25 wants 32 tokens, dense retrievers 64 §5.3 app. A. The dense numbers cannot be reproduced from the released code.
• Results depend on the retrieval corpus and its preprocessing. The WikiText-103 numbers use a 2018 Wikipedia corpus with 120 articles removed §4.3. The other corpora use each dataset's training split §4.3. Results on another corpus are not covered.
• Title parsing fails on extra newlines. The sparse retriever expects each document to contain exactly one newline between title and text sparse_retrieval.py. Any other passage raises an error. The paper does not discuss this.
• The window trims the prefix, not the document. Documents take up to 256 tokens of a 1,024-token window §4.2 §4.3, so the earliest prefix tokens are dropped first §3.1.
• Perplexity is not generation quality. All LM results are perplexity tab. 1 tab. 2 tab. 5. The QA results use exact match on DPR-provided passages §7 tab. 4. The paper does not measure generated text quality.
• Sequence length. The paper caps input at 1,024 tokens for comparability, although some models allow 2,048 §4.2. Using 2,048 may change the numbers §4.2 fn. 5.
• Version pins. The code pins torch==1.12.1 and transformers==4.28.1 requirements.txt, and BM25 needs Pyserini with Java 11 README.md.

Lineage

• Builds on: the retrieve-and-read family and kNN-LM §2. Closest is RETRO, which modifies the LM with chunked cross-attention and uses a frozen BERT retriever §2 §4.2.
• Compared against: RETRO's BERT-based retriever and the dense retrievers Contriever and Spider §4.2 §5.1.
• Concurrent work: REPLUG also prepends retrieved text to an LM, but trains a dedicated retriever §1.
• Future directions named by the paper: several documents per prompt, and sparser retrieval triggered by a model §8. The paper also cites a recent study that reports factuality gains from In-Context RALM §8.