Shall We Pretrain Autoregressive Language Models with Retrieval? A Comprehensive Study (2023)

Abstract: Large decoder-only language models (LMs) can be largely improved in terms of perplexity by retrieval (e.g., RETRO), but its impact on text generation quality and downstream task accuracy is unclear. Thus, it is still an open question: shall we pretrain large autoregressive LMs with retrieval? To answer it, we perform a comprehensive study on a scalable pre-trained retrieval-augmented LM (i.e., RETRO) compared with standard GPT and retrieval-augmented GPT incorporated at fine-tuning or inference stages. We first provide the recipe to reproduce RETRO up to 9.5B parameters while retrieving a text corpus with 330B tokens. Based on that, we have the following novel findings: i) RETRO outperforms GPT on text generation with much less degeneration (i.e., repetition), moderately higher factual accuracy, and slightly lower toxicity with a nontoxic retrieval database. ii) On the LM Evaluation Harness benchmark, RETRO largely outperforms GPT on knowledge-intensive tasks, but is on par with GPT on other tasks. Furthermore, we introduce a simple variant of the model, RETRO++, which largely improves open-domain QA results of original RETRO (e.g., EM score +8.6 on Natural Question) and significantly outperforms retrieval-augmented GPT across different model sizes. Our findings highlight the promising direction of pretraining autoregressive LMs with retrieval as future foundation models. We release our implementation at: https://github.com/NVIDIA/Megatron-LM/tree/main/tools/retro.

## introduction large language models (lms), including masked lms (e.g., bert @xcite ), autoregressive lms (e.g., gpt @xcite ), and encoder-decoder lms (e.g., t5 @xcite , bart @xcite ), have obtained state-of-the-art results for various nlp tasks. among them, the autoregressive lms like gpt-3 @xcite and @xcite demonstrate noticeable in-context learning ability and excellent long-form text generation results. due to its importance, the community has spent considerable efforts to scale up such autoregressive generative lms with more data and parameters and observed significant breakthroughs in a variety of real-world applications (e.g., @xcite ), including open-ended text generation and various downstream tasks (e.g., question answering). the successful public examples include gpt-3 (w/ 170b parameters) @xcite , gopher (280b) @xcite , megatron-turing (530b) @xcite , and palm (540b) @xcite . although large-scale autoregressive lms have achieved huge successes, they also suffer from several weaknesses. first, it requires a huge number of model parameters to memorize the world knowledge, which makes it costly for deployment. second, it is difficult to safeguard factual accuracy, which may provide users with incorrect information @xcite . third, it is expensive to update the model knowledge learned during pretraining with up-to-date facts @xcite , yielding outdated answers @xcite . to mitigate the problems above, one line of research proposes to improve language models with retrieval. the retrieval process can be integrated into lms at: i) fine-tuning stage @xcite @xcite , or ii) pretraining stage @xcite . most previous work augments bert or encoder-decoder lms with retrieval at fine-tuning stage, demonstrating successes for knowledge-intensive nlp tasks @xcite @xcite . however, it re-mains relatively underexplored to pretrain autoregressive (decoder-only) lms with retrieval, especially considering the noticeable success of chat-gpt (openai, 2022) that underscores the extreme importance of the autoregressive lms. most recently, retro @xcite proposes to pretrain autoregressive lms with a retrieval module, which is practically scalable to large-scale pretraining from scratch by retrieving billions of token and largely reduces model parameters while achieving lower perplexity than standard gpt. it also provides the flexibility to update the knowledge stored in lms @xcite by updating the retrieval database without training lms again. the success of pretraining lms with retrieval raises an important question for the community if we want to pretrain autoregressive lms in the future: shall we pretrain autoregressive (decode-only) lms with retrieval by default or not? however, previous work @xcite misses the important evaluation on whether the model like retro could obtain comparable or even better results in terms of open-ended text generation and various nlp downstream tasks, apart from lower perplexity on the held-out dataset compared to standard gpt. to answer the above question and bridge the missing gap, we perform an extensive study on retro, as to the best of our knowledge, retro is the only retrieval-augmented autoregressive lm that supports large-scale pretraining with retrieval on the massive pretraining corpus with hundreds of billion or trillion tokens. our comprehensive study sheds light on the promising direction of pertaining autoregressive lms with retrieval to serve as future foundation models, as they overall outperform standard gpt models in terms of perplexity, text generation quality, and downstream task performances, especially for knowledge-intensive tasks, including open-domain qa. ## key findings we successfully reproduce and pretrain retro @xcite from scratch foot_0 , with parameter sizes ranging from 148m up to 9.5b by retrieving from a text corpus with over 330b tokens. in addition, we discuss the inference strategy of retro for text generation that is not covered in @xcite , and perform a large-scale evaluation in different scenarios. to minimize the discrepancy variables between retro and gpt, we use the same decoder architecture, same hyper-parameters, and same pre-training corpus to pre-train retro and gpt given the same number of pre-training steps. we highlight our novel findings for retro and gpt as follows: ## related work retrieval has been applied in various nlp tasks for years, including question answering (qa) (e.g., @xcite , machine translation (e.g., @xcite , and conversation @xcite @xcite . in particular, language models have been augmented with retrieval at different stages, including inference time @xcite , fine-tuning stage @xcite @xcite , and pretraining stage @xcite . lms have been augmented with retrieval at the fine-tuning stage for downstream tasks, primarily for open-domain qa. dpr @xcite finetunes one bert to encode questions and the other bert to encode answers within a dual encoder framework, using a contrastive loss to align the hidden representations of question and corresponding answer. rag @xcite studies the fine-tuning recipe for retrieval-augmented generation models, especially on open-domain qa tasks. fid @xcite improves rag with a better lm backbone t5, and fuses multiple retrieved passages to the decoder during fine-tuning to further improve qa accuracy. we-bgpt @xcite leverages web search engine and fine-tunes gpt using reinforcement learning with human feedback (rlhf) for reference generation and factuality improvement, which is orthogonal to our work that focuses on pretraining with retrieval. the proposed rlhf can be applied to retro as well. realm @xcite performs both unsupervised pretraining and supervised fine-tuning strategies for retrieval-augmented bert model in open-domain qa. their pretraining involves asynchronous re-embedding and re-indexing all documents every several hundred training steps, which quickly becomes impractical for training corpus with trillion tokens. atlas @xcite uses a similar approach but augments the t5 architecture @xcite with retrieval at both pre-training and fine-tuning. before pretraining, it first initializes the encoder-decoder lm backbone with pretrained t5, and the dense retriever with pretrained contriever @xcite . during pretraining, it also applies asynchronous index refresh every 1000 steps. in contrast, retro @xcite embeds and indexes the whole training corpus at chunk-level (e.g., chuck @xmath0. during pretraining, the model relies on a trainable bidirectional encoder to embed the retrieved chunks of raw text. the gpt decoder further "select" the relevant piece of evidence from the encoder side by a chunk-wise cross-attention. this architecture design enables lm pretraining on hundreds of billion tokens by retrieving from trillion tokens. see ## model and implementation in this section, we first introduce preliminaries of retro, then provide detailed recipe of our implementation, including retrieval database, pretraining, and retrieval-augmented finetuning and generation. ## open-ended text generation in this section, we delve into the problem of openended text generation, which refers to tasks of generating coherent continuation given the preceding prompt. given that this problem for retro has never been studied before, we manage to bridge the gap and evaluate the open-ended text generation of retro compared to gpt from three aspects: a) text quality, b) factuality, and c) toxicity. ## lm evaluation harness benchmark besides the open-ended text generation, it is also important to examine the generalization of retro on various downstream tasks, which is also missing from the literature. therefore, we use lm evaluation harness benchmark @xcite and consider the following nine representative nlp downstream tasks. see more details in §g. zero-shot evaluation. we present the zero-shot evaluation results in ## open-domain question answering in this section, we study two widely used opendomain qa datasets, natural question (nq) and triviaqa. ## conclusion in this work, we perform a comprehensive study of pretrained retrieval-augmented llm to answer the question: shall we pretrain decoder-only lms with retrieval? we observe consistent improvements in text generation quality, factual accuracy, lower toxicity, and downstream task accuracy, especially for knowledge-intensive tasks, including open-domain qa. given the ∼ 25% percentage of additional gpu hours for pretraining (see