Source-linked AI summary
REPLUG: Retrieval-Augmented Black-Box Language Models
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettlemoyer, Wen-tau Yih
TL;DR
REPLUG addresses the challenge of augmenting large language models that cannot be modified or accessed internally. It treats the LM as a black box, prepends retrieved documents, and tunes the retriever using LM supervision; this improves performance across tasks, including GPT-3 language modeling and Codex MMLU. The approach remains limited by unclear interpretability and, in open-domain QA, performance gaps associated with few-shot training data.
Problem
Large or API-accessed language models cannot readily use conventional retrieval-augmented methods that require modifying LM parameters or accessing internal representations.
Method
REPLUG prepends documents retrieved from an external corpus to a frozen black-box LM, while REPLUG LSR tunes the retriever using LM-based supervision.
Results
REPLUG improves diverse black-box LMs across language modeling and downstream tasks, with REPLUG LSR reporting a 6.3% increase in GPT-3 175B language modeling.
Takeaways & Limitations
Retrieval can improve performance even for large-scale black-box language models, including models accessed through APIs.
Takeaways & Limitations
REPLUG lacks interpretability because it is unclear when the model relies on retrieved knowledge versus parametric knowledge.
Abstract
from arXiv · showhide
We introduce REPLUG, a retrieval-augmented language modeling framework that treats the language model (LM) as a black box and augments it with a tuneable retrieval model. Unlike prior retrieval-augmented LMs that train language models with special cross attention mechanisms to encode the retrieved text, REPLUG simply prepends retrieved documents to the input for the frozen black-box LM. This simple design can be easily applied to any existing retrieval and language models. Furthermore, we show that the LM can be used to supervise the retrieval model, which can then find documents that help the LM make better predictions. Our experiments demonstrate that REPLUG with the tuned retriever significantly improves the performance of GPT-3 (175B) on language modeling by 6.3%, as well as the performance of Codex on five-shot MMLU by 5.1%.
1. Introduction
REPLUG addresses the difficulty of retrieval augmentation for large or API-accessed language models by treating the LM as a black box and adding retrieval externally. It prepends retrieved documents, can tune the retriever using LM supervision, and improves performance across language modeling and downstream tasks.
- Motivation: Large language models can hallucinate and miss long-tail knowledge, while prior retrieval-augmented methods often require access to or modification of LM parameters.These constraints make retrieval augmentation difficult for very large or API-only models.
- REPLUG: REPLUG treats the LM as a black box and prepends retrieved documents to the input, making the framework applicable to existing black-box LMs and retrieval models.The approach is designed for large models whose internal representations or fine-tuning interfaces are unavailable.
- REPLUG LSR: REPLUG LSR tunes the retriever with supervision from a frozen black-box LM, preferring documents that improve LM perplexity.This adapts the retriever to the LM rather than adapting the LM to the retriever.
- Results: 4.5% improvement in Codex (175B) performance on MMLU is reported for REPLUG, reaching results comparable to 540B instruction-finetuned Flan-PaLM.The reported result concerns MMLU performance.
- Results: 6.3% increase in GPT-3 175B language modeling is reported when the retriever is tuned with REPLUG LSR.The paper also reports improvements across diverse black-box LMs on language modeling, MMLU, and open-domain QA.
2. Background and Related Work
The background motivates retrieval augmentation as a way to supplement language models with external information, but emphasizes that conventional white-box approaches are difficult to use with large black-box models. REPLUG investigates retrieval augmentation under this black-box constraint.
- Black-box Language Models: Large models may be available only through black-box APIs, while locally running or fine-tuning open models can require substantial computational resources.The paper gives BLOOM-176B fine-tuning as an example requiring 72 A100 GPUs.
- Black-box Language Models: Traditional retrieval-augmented frameworks commonly fine-tune language models to incorporate retrieved documents, an approach the paper describes as infeasible for increasingly large black-box models.The black-box setting provides access only to model predictions, not parameters or internal representations.
- Retrieval-augmented Models: Retrieval-augmented models retrieve documents from a corpus using the input as a query and then incorporate those documents into the language model.This general pattern has been applied to language modeling and open-domain question answering.
- REPLUG: REPLUG performs retrieval at inference by obtaining relevant documents from an external corpus, prepending each separately to the input, and ensembling output probabilities.This inference design is presented as part of the paper’s black-box retrieval-augmentation approach.
3. REPLUG
REPLUG retrieves documents for an input context and feeds each retrieved document with that context through a black-box LM, combining the resulting predictions. Its ensemble design supports a compute–accuracy trade-off while avoiding additional computational overhead relative to prepending all documents together.
- Framework: REPLUG treats the language model as a black box and adds retrieval as a potentially tuneable module.The framework is intended to integrate with existing language models without modifying their parameters.
- Inference: Given an input context, REPLUG retrieves a small set of relevant documents from an external corpus and processes each document with the context through the LM in parallel.The predicted probabilities from these passes are ensembled.
- Document Retrieval: The dense retriever encodes documents and the input context, then retrieves the top-k documents with the highest cosine-similarity scores.Document embeddings are computed using mean pooling of the encoder’s final-token representations, and a FAISS index supports efficient retrieval.
- Input Reformulation: The ensemble strategy processes retrieved documents separately because the LM context window limits how many documents can be prepended together.This design allows compute to be traded for accuracy by changing the number of retrieved-document passes.
- Input Reformulation: The ensemble method requires k LM runs but does not incur additional computational-cost overhead compared with prepending all retrieved documents.Each pass performs cross attention between one retrieved document and the input context.
4. REPLUG LSR: Training the Dense Retriever
REPLUG LSR trains the retriever to align its document probabilities with scores from a frozen language model, favoring documents that improve prediction likelihood.
- REPLUG LSR adapts the retriever by using the black-box language model to supervise which documents should be retrieved.
- The procedure retrieves top-scoring documents, computes retrieval likelihoods, scores them with the language model, and updates the retriever by minimizing distributional divergence.
- The retrieval likelihood marginalizes over the retrieved subset rather than the full corpus because full-corpus computation is intractable.
- Higher language-model probability for the ground-truth continuation identifies documents that better improve language-model perplexity.
- Only retrieval-model parameters are updated during training; the language-model parameters remain fixed under the black-box assumption.
- Document embeddings and the efficient search index are recomputed every T training steps after retriever updates make earlier embeddings stale.
5. Training Setup
The training setup uses Contriever-based retrieval, GPT-3 Curie supervision, sampled Pile data, and periodic retrieval-model training procedures.
- The training section separates the REPLUG model setting from the REPLUG LSR procedure for training the retriever.
- REPLUG can use dense or sparse retrievers, while the experiments use Contriever because of its demonstrated strong performance.
- REPLUG LSR initializes the retriever with Contriever and uses GPT-3 Curie as the supervision language model for computing likelihoods.
- Training queries comprise 800K sampled Pile sequences of 256 tokens, split into 128-token input contexts and 128-token continuations.
- The external corpus contains sampled 128-token Pile documents that do not overlap with documents used for training queries.
- Figure 3 depicts retriever training supervised by outputs from a frozen language model.
- For efficiency, the setup precomputes corpus embeddings, builds a FAISS index, retrieves the top 20 documents, and uses temperature 0.1 for likelihood computation.
6. Experiments
Across language modeling, MMLU, and open-domain QA, REPLUG and REPLUG LSR improve black-box language-model performance, with gains reported across model sizes and evaluation settings.
- Experiments: REPLUG and REPLUG LSR consistently improve language-model performance across language modeling, MMLU, and open-domain QA evaluations.The experiments cover GPT, OPT, and BLOOM models, including models with up to 175B parameters.
- 6.1. Language Modeling: 7.7% improvement over baselines is achieved by REPLUG LSR averaged over 8 models, compared with 4.7% for REPLUG in language modeling.The reported gain is relative to the original language-model baselines.
- 6.2. MMLU: 5.1% improvement is obtained by REPLUG LSR over original Codex on MMLU, compared with 4.5% for REPLUG.MMLU results are evaluated with 5-shot in-context learning and are averaged across four categories.
- 6.2. MMLU: 1.9% improvement over the original model is reported for REPLUG LSR in MMLU’s STEM category.The paper presents this result as evidence that retrieval may improve problem-solving abilities.
- 6.3. Open Domain QA: 12.0% improvement on NQ and 5.0% on TriviaQA is reported for REPLUG LSR over original Codex in the few-shot setting.The model uses 16-shot in-context learning with top-10 retrieved documents incorporated through the ensemble method.
- 6.3. Open Domain QA: REPLUG LSR achieves a new few-shot state of the art on NQ and TriviaQA, while still lagging retrieval-augmented models fine-tuned on full training data.The reported comparison includes Atlas, which is fine-tuned with 64 training examples.
7. Analysis
REPLUG’s gains persist across retrieval settings and language-model families, with relevant-document ensembling outperforming random ensembling and helping especially with rare entities. The method improves perplexity across model sizes and families on Wikitext-103.
- 7.1. REPLUG performance gain does not simply come from the ensembling effect: Relevant-document ensembling improves performance, whereas random-document ensembling performs worse; increasing the number of relevant documents improves results monotonically.About 10 documents suffice for large performance gains.
- 7.2. REPLUG is applicable to diverse language models: GPT-2, BLOOM, and OPT models of varying sizes consistently benefit from REPLUG on Wikitext-103 perplexity.Figure 5 compares model size on the x-axis with perplexity on the y-axis.
- 7.2. REPLUG is applicable to diverse language models: 6.9% perplexity improvement is achieved by OPT with 125M parameters, compared with 5.6% for OPT with 66B parameters.REPLUG improves perplexity across GPT-2, OPT, and BLOOM model families.
8. Conclusion
REPLUG integrates retrieval with black-box language models to improve language modeling and downstream-task performance. The authors also identify limited interpretability as an open issue for future work.
- 8. Conclusion: REPLUG treats the language model as a black box and augments it with a tuneable retrieval model.The framework is presented as a retrieval-augmented language modeling paradigm.
- 8. Conclusion: REPLUG can be integrated with existing language models to improve performance on language modeling or downstream tasks.The conclusion describes this result as opening possibilities for integrating retrieval into large-scale black-box language models.
- 8. Conclusion: REPLUG lacks interpretability because it is unclear when the model relies on retrieved knowledge or parametric knowledge.The authors identify developing more interpretable retrieval-augmented language models as future work.