REPLUG : Retrieval-Augmented Black-Box Language Models (2024)

Abstract: 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 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%. Code is publicly released at github.com/swj0419/REPLUG.

## introduction large language models (lms) such as @xcite and codex @xcite , have demonstrated impressive performance on a wide range of language tasks. these models are typically trained on very large datasets and store a substantial amount of world or domain knowledge implicitly in their parameters. however, they are also prone to hallucination and cannot represent the full long tail of knowledge from the training corpus. retrieval-augmented language models @xcite @xcite , in contrast, can retrieve knowledge from an external datastore when needed, potentially reducing hallucination and increasing coverage. previous approaches of retrieval-augmented language models require access to the internal lm representations (e.g., to train the model @xcite or to index the datastore @xcite figure 1 : different from previous retrieval-augmented approaches @xcite that enhance a language model with retrieval by updating the lm's parameters, replug treats the lm as a black box and augments it with a frozen or tunable retriever. this black-box assumption makes replug applicable to large lms, which are often served via apis. et al., 2020)), and are thus difficult to be applied to very large lms. in addition, many best-in-class llms can only be accessed through apis. internal representations of such models are not exposed and fine-tuning is not supported. in this work, we introduce replug (retrieve and plug), a new retrieval-augmented lm framework where the language model is viewed as a black box and the retrieval component is added as a tuneable plug-and-play module. given an input context, replug first retrieves relevant documents from an external corpus using an off-theshelf retrieval model. the retrieved documents are prepended to the input context and fed into the black-box lm to make the final prediction. because the lm context length limits the number of documents that can be prepended, we also adopt an ensemble scheme that encodes the retrieved documents in parallel with the same black-box lm, allowing us to easily trade compute for accuracy. as shown in figure 1 , replug is extremely flexible and can be used with any existing black-box lm and retrieval model. we also introduce replug lsr (replug with lm-supervised retrieval), a training scheme that can further improve the initial retrieval model in replug with supervision signals from a blackbox language model. the key idea is to adapt the retriever to the lm, which is in contrast to prior work @xcite that adapts language models to the retriever. we use a training objective which prefers retrieving documents that improve language model perplexity, while treating the lm as a frozen, black-box scoring function. our experiments show that replug can improve the performance of diverse black-box lms on both language modeling and downstream tasks, including mmlu @xcite and open-domain qa @xcite . for instance, replug can improve codex (175b) performance on mmlu by 4.5%, achieving comparable results to the 540b, instruction-finetuned flan-palm. furthermore, tuning the retriever with our training scheme (i.e., replug lsr) outperforms various off-theshelf retrievers and leads to additional improvements, including up to 6.3% increase in gpt-3 175b language modeling. to the best of our knowledge, our work is the first to show the benefits of retrieval to large lms (>100b model parameters), for both reducing lm perplexity and and improving in-context learning performance. we summarize our contributions as follows: black-box language models large language models, such as gpt-3 (brown et al., 2020), codex @xcite , are not open-sourced due to commercial considerations and are only available as black-box apis, through which users can send queries and receive responses. on the other hand, even open sourced language models such as bloom-176b (scao et al., 2022) require significant computational resources to run and finetune locally. for example, finetuning bloom-176b requires 72 a100 gpus (younes belkda, 2022), making them inaccessible to researchers and developers with limited resources. traditionally, retrieval-augmented model frameworks @xcite @xcite @xcite have focused on the white-box setting, where language models are fine-tuned to incorporate retrieved documents. however, the increasing scale and blackbox nature of llms makes this approach infeasible. to address these challenges, we investigate retrieval-augmentation in the black-box setting, where users only have access to the model predictions and cannot access or modify its parameters. retrieval-augmented models augmenting language models with relevant information retrieved from knowledge stores has shown to be effective in improving performance on various nlp tasks, including language modeling @xcite @xcite and open-domain question answering @xcite @xcite . specifically, using the input as query, (1) a retriever first retrieves a set of documents from a corpus and then (2) a language model incorporates the retrieved documents as additional information to make a final prediction. previous retrieval-augmented lms require updating the model parameters , which cannot be applied to black-box lms, which cannot be applied to black-box lms. for example, atlas @xcite finetunes an encoder-decoder model jointly with the retriever by modeling documents as latent variables, while retro @xcite changes the decoder-only architecture to incorporate retrieved texts and pretrains the language model from scratch. another line of then it prepends each document separately to the input context and ensembles output probabilities from different passes ( §3.2 input reformulation). retrieval-augmented lms such as knn-lm (khandelwal et al., 2020; zhong et al., 2022) retrieves a set of tokens and interpolates between the lm's next token distribution and knn distributions computed from the retrieved tokens at inference. knn-lm requires access to internal lm representations to compute the knn distribution, which are not available for black-box lms such as gpt-3. in this work, we investigate ways to improve large blackbox language models with retrieval. while concurrent work @xcite has demonstrated that using a frozen retriever can improve gpt-3 performance on open-domain question answering, we approach the problem in a more general setting, including language modeling and understanding tasks. we additionally adopt an ensemble method to incorporate more documents and a training scheme to further adapt the retriever to large lms. ## replug we introduce replug (retrieve and plug), a new retrieval-augmented lm paradigm where the lm is treated as black box and the retrieval component is added as a potentially tuneable module. as shown in figure 2 , given an input context, replug first retrieves a small set of relevant documents from an external corpus using a retriever ( §3.1). then we pass the concatenation of each retrieved document with the input context through the lm in parallel, and ensemble the predicted probabilities ( §3.2). ## replug lsr: training the dense retriever instead of relying only on existing neural dense retrieval models @xcite @xcite , we further propose replug lsr (replug with lm-supervised retrieval), which adapts the retriever in replug by using the lm itself to provide supervision about which documents should be retrieved. inspired by @xcite , our approach can be seen as adjusting the probabilities of the retrieved documents to match the probabilities of the output sequence perplexities of the language model. in other words, we would like the retriever to find documents that result in lower perplexity scores. as shown in figure 3 , our training algorithm consists of the four steps: (1) retrieving documents and computing the retrieval likelihood ( §4.1), (2) scoring the retrieved documents by the language model ( §4.2), (3) updating the retrieval model parameters by minimizing the kl divergence between the retrieval likelihood and the lm's score distribution ( §4.3), and (4) asynchronous update of the datastore index ( §4.4). ## experiments we perform evaluations on both language modeling ( §6.1) and downstream tasks such as mmlu ( §6.2) and open-domain qa ( §6.3). in all settings, re-plug ĩmprove the performance of various blackbox language models, showing the effectiveness and generality of our approach. ## conclusion we introduce replug, a retrieval-augmented lm paradigm that augments black-box lms with a tuneable retriever. this work opens up new possibilities for integrating retrieval into large black-box lms and is the first to demonstrate even the stateof-the-art llms could benefit from retrieval.