Source-linked AI summary
In-Context Retrieval-Augmented Language Models
Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton-Brown, Yoav Shoham
TL;DR
Existing RALM methods improve language modeling and grounding but commonly require changes to the LM architecture and dedicated retraining, complicating deployment. This paper introduces In-Context RALM, which prepends retrieved documents to the input of an unchanged LM and uses retrieval and ranking tailored to language modeling. The approach achieves substantial gains across model sizes and diverse corpora, while specialized document selection provides additional improvements and supports use with frozen or API-accessed LMs.
Problem
Existing RALM approaches commonly modify the LM architecture and require dedicated retraining, complicating deployment.
Method
In-Context RALM prepends retrieved grounding documents to the input of an unchanged, off-the-shelf LM and adapts retrieval and ranking to language modeling.
Results
In-Context RALM provides substantial language-modeling gains across model sizes and diverse corpora, with specialized document selection producing additional gains.
Takeaways & Limitations
The framework can make LM grounding more deployable with frozen pretrained models, including models accessed through APIs.
Takeaways & Limitations
The study prepends only one external document and retrieves at fixed token intervals, leaving multi-document retrieval and sparser retrieval for future work.
Abstract
from arXiv · showhide
Retrieval-Augmented Language Modeling (RALM) methods, which condition a language model (LM) on relevant documents from a grounding corpus during generation, were shown to significantly improve language modeling performance. In addition, they can mitigate the problem of factually inaccurate text generation and provide natural source attribution mechanism. Existing RALM approaches focus on modifying the LM architecture in order to facilitate the incorporation of external information, significantly complicating deployment. This paper considers a simple alternative, which we dub In-Context RALM: leaving the LM architecture unchanged and prepending grounding documents to the input, without any further training of the LM. We show that In-Context RALM that builds on off-the-shelf general purpose retrievers provides surprisingly large LM gains across model sizes and diverse corpora. We also demonstrate that the document retrieval and ranking mechanism can be specialized to the RALM setting to further boost performance. We conclude that In-Context RALM has considerable potential to increase the prevalence of LM grounding, particularly in settings where a pretrained LM must be used without modification or even via API access.
1 Introduction
In-Context RALM addresses the difficulty of grounding language models without modifying or retraining them by prepending retrieved documents to the input. Across diverse corpora and model sizes, this simple approach yields substantial gains, with task-specialized ranking providing further improvements.
- RALM grounds language models with external documents to address factual inaccuracies, missing provenance, and limited access to external knowledge.
- Existing RALM systems often modify the LM architecture and require dedicated retraining, hindering wider adoption.
- In-Context RALM prepends selected documents to the LM input while leaving the LM architecture and weights unchanged.
- Experiments cover five diverse corpora and open-source LMs ranging from 110M to 66B parameters.
- In-Context RALM with off-the-shelf retrievers produced gains equivalent to increasing LM size by 2–3× across examined corpora.
- Task-specialized rerankers added gains equivalent to a further 2× LM size increase, including improvements from 345M GPT-2 beyond larger GPT-2 baselines.
2 Related Work
In-Context RALM belongs to retrieve-and-read RALM but differs from prior systems by using frozen off-the-shelf LMs and emphasizing document selection for language modeling. It therefore avoids further LM training while investigating retrieval choices that improve LM performance.
- RALM comprises nearest-neighbor language models and retrieve-and-read models; In-Context RALM belongs to the latter family.
- Prior retrieve-and-read approaches train the LM, whereas In-Context RALM reads documents with an off-the-shelf LM without further training.
- RETRO modifies an autoregressive LM with chunked cross-attention and new parameters to attend to retrieved documents.
- In-Context RALM focuses on choosing documents that improve language-modeling performance.
3 Our Framework
In-Context RALM retrieves documents from an external corpus and prepends them to the language model’s input without changing the LM weights. Its design exposes retrieval stride and query length as tunable choices balancing relevance, runtime, and performance.
- RALM conditions next-token prediction on documents retrieved from an external corpus using the current prefix.
- In-Context RALM concatenates retrieved documents before the prefix in the Transformer input without altering LM weights.
- Retrieval Stride: Retrieval stride s controls how many tokens separate retrieval operations, with larger strides reducing retrieval frequency and runtime costs.
- Retrieval Stride: Each retrieval operation incurs retriever and prefix-embedding costs, making stride selection a runtime–performance tradeoff.
- Retrieval Query Length: The retrieval query uses the last ℓ prefix tokens, because overly long queries can dilute information near the generation point.
4 Experimental Details
The experiments evaluate In-Context RALM across diverse language-modeling and question-answering datasets, open-source models from 110M to 66B parameters, and sparse and dense retrievers. Retrieval corpora and sequence-length choices are standardized to support controlled comparisons.
- Datasets: The evaluation covers five language-modeling datasets and two open-domain question-answering datasets.
- Datasets: The language-modeling corpora span WikiText-103, RealNews, ArXiv, Stack Exchange, and FreeLaw.
- Language Models: Experiments use open-source GPT-2, GPT-Neo, GPT-J, OPT, and LLaMA models ranging from 110M to 66B parameters.
- Implementation: All models use a maximum sequence length of 1,024 tokens to enable direct comparisons.
- Retrievers: The study compares BM25 with mean-pooled BERT, Contriever, and Spider dense retrievers.
- Retrieval Corpora: Retrieval uses Wikipedia for WikiText-103 and question answering, while the remaining datasets use their training data as retrieval corpora.
5 The Effectiveness of In-Context RALM with Off-the-Shelf Retrievers
In-Context RALM produces substantial language-modeling gains with off-the-shelf retrieval across diverse corpora and model sizes. BM25 is the strongest tested general-purpose retriever, while frequent retrieval and a query length near 32 tokens provide favorable configurations.
- In-Context RALM yields substantial language-modeling gains across the diverse evaluation suite.
- Off-the-shelf retrieval improves perplexity enough to match models 2–3× larger across the examined corpora.
- BM25 Outperforms Off-the-Shelf Neural Retrievers: BM25 outperforms Contriever, Spider, and mean-pooled BERT retrievers in the tested In-Context RALM setting.
- BM25 Outperforms Off-the-Shelf Neural Retrievers: A 6.7B-parameter OPT model with off-the-shelf In-Context RALM matches the performance of a 66B-parameter OPT model.
- Retrieval Frequency: More frequent retrieval improves language-modeling performance, but the authors use s = 4 to balance performance against runtime cost.
- Retrieval Query Length: BM25 has a query-length sweet spot around 32 tokens, whereas excessively long queries can dilute relevance to the language-modeling task.
6 Improving In-Context RALM with LM-Oriented Reranking
This section improves In-Context RALM by specializing document selection: rerankers choose among BM25’s top documents using signals aligned with upcoming language-model predictions. Zero-shot and trained reranking both improve performance, while smaller reranking models can support API-accessed LMs.
- Motivation: Reranking the top-k BM25 documents targets a major limitation of selecting only the first lexical match.BM25 offers limited semantic understanding and cannot weight query tokens by their relevance to generated text.
- 6.1 LMs as Zero-Shot Rerankers: Zero-shot reranking uses an LM to select the document that best supports the upcoming text, approximated at test time with preceding prefix tokens.The method reranks the top 16 BM25 documents; the development-set optimum for the prefix stride is s′ = 16.
- 6.1 LMs as Zero-Shot Rerankers: A reranking LM need not match the generation LM, enabling smaller and faster rerankers and use when the generation model’s log probabilities are unavailable through an API.A GPT-2 110M reranker achieved roughly the same performance as larger generation models reranking themselves, with only minor degradation.
- 6.2 Training LM-dedicated Rerankers: Predictive Reranking trains a classifier on target-corpus data to score each BM25 candidate by its relevance to the continuation.The classifier receives the prefix and candidate document, producing a scalar relevance score while the LM remains frozen.
- 6.2 Training LM-dedicated Rerankers: Predictive Reranking improved GPT-2 110M perplexity from 29.6 to 26.8 and GPT-2 1.5B perplexity from 16.6 to 15.4.The authors report that domain-specific training was more effective than zero-shot reranking, while top-16 BM25 oracle results still leave room for improvement.
7 In-Context RALM for Open-Domain Question Answering
The paper evaluates In-Context RALM on open-domain question answering using LLaMA models with and without retrieved documents, showing substantial gains from in-context retrieval.
- Evaluation setup: The evaluation tests whether LMs can use retrieved documents for open-domain question answering without further training or training examples.Experiments use the LLaMA family on Natural Questions and TriviaQA.
- Document count: Most of the improvement can be obtained with only two documents, or even one document in some cases.The analysis examines the effect of document count on the Natural Questions and TriviaQA development sets.
- Results: Showing relevant documents significantly improves question-answering performance on Natural Questions and TriviaQA.The test-set evaluation uses two retrieved documents, motivated by the document-count analysis.
- Results: 18+ points on Natural Questions: LLaMA-13B zero-shot performance rises from 12.0% to 31.0% with retrieved documents.This is the paper’s concrete Natural Questions example.
- Results: 5+ points on TriviaQA: LLaMA-13B zero-shot performance rises from 54.8% to 60.1% with retrieved documents.This is the corresponding concrete TriviaQA example.
8 Discussion
The discussion presents In-Context RALM as a way for frozen, off-the-shelf LMs to benefit from retrieval while avoiding architectural modification, and identifies directions for further improvement.
- Motivation: In-Context RALM addresses factual inaccuracies and missing provenance by grounding language-model generation in retrieved external documents.These issues remain common in machine-generated text and make language modeling harder to trust.
- Discussion: The framework enables frozen, off-the-shelf LMs to benefit from retrieval using general-purpose retrievers, with further gains from tailoring document selection to the LM setting.The discussion also cites evidence that In-Context RALM can improve the factuality of large LMs.
- Future work: Future work could add more external documents and retrieve more sparsely to reduce latency and cost.The paper studied prepending a single document and retrieving at fixed token intervals.
- Open research: The authors released the code to support community use and further improvement of RALM.They hope this will encourage research and wider adoption.
A Query Length Ablations
The query-length ablations compare optimal query lengths for off-the-shelf dense retrievers and BM25, revealing different preferred lengths.
- Dense retrievers: The ablations examine optimal query length ℓ for the off-the-shelf dense retrievers BERT and Contriever.Spider results are omitted because they are almost identical to Contriever’s.
- Dense retrievers: 64 tokens is consistently optimal for the off-the-shelf dense retrievers BERT and Contriever.The finding is reported across the dense-retriever ablations.
- BM25 comparison: 32 tokens is optimal for BM25, contrasting with the 64-token optimum for the dense retrievers.The BM25 comparison is based on the corresponding Figure 6 experiments.
B GPT-Neo Results
This appendix reports In-Context RALM results for the GPT-Neo model family on WikiText-103 and RealNews.
- Models: Table 5 reports results for applying In-Context RALM to models in the GPT-Neo family.The passage identifies the model family evaluated in this appendix.
- Datasets: The GPT-Neo evaluation uses WikiText-103 as one of its datasets.The dataset is listed alongside RealNews in the table description.
- Datasets: The GPT-Neo evaluation also uses RealNews as a corpus.The passage describes Table 5 as covering both named datasets.
C Open-Domain Question Answering Experiments: Further Details
The section details closed-book and open-book question-answering prompts and identifies analyses of perplexity across query lengths and model families.
- Closed-Book Setting: The closed-book setting uses a question-answering prompt ending with “A:” after the question.The example asks who received the first Nobel Prize in Physics.
- Open-Book Setting: The open-book setting extends the closed-book prompt by adding retrieved texts before the question-answering instruction.The prompt includes passages about Nobel-related topics and asks the same physics question.
- Model Evaluation: Table 5 reports GPT-Neo-family performance using word-level perplexity on WikiText-103 test data and token-level perplexity on RealNews development data.The table compares model performance across two datasets and perplexity granularities.
- Retriever Analysis: Figure 9 analyzes perplexity as a function of query-token count for an off-the-shelf BERT retriever on WikiText-103 development data.The figure caption specifies the retriever, dataset split, and independent variable.
- Retriever Analysis: Figure 10 analyzes perplexity as a function of query-token count for Contriever on WikiText-103 development data.This provides a parallel query-length analysis using Contriever.