Retrieval-Pretrained Transformer: Long-range Language Modeling with Self-retrieval (2024)

Abstract: Retrieval-augmented language models (LMs) have received much attention recently. However, typically the retriever is not trained jointly as a native component of the LM, but added post-hoc to an already-pretrained LM, which limits the ability of the LM and the retriever to adapt to one another. In this work, we propose the Retrieval-Pretrained Transformer (RPT), an architecture and training procedure for jointly training a retrieval-augmented LM from scratch and applying it to the task of modeling long texts. Given a recently generated text chunk in a long document, the LM computes query representations, which are then used to retrieve earlier chunks in the document, located potentially tens of thousands of tokens before. Information from retrieved chunks is fused into the LM representations to predict the next target chunk. We train the retriever component with a semantic objective, where the goal is to retrieve chunks that increase the probability of the next chunk, according to a reference LM. We evaluate RPT on four long-range language modeling tasks, spanning books, code, and mathematical writing, and demonstrate that RPT improves retrieval quality and subsequently perplexity across the board compared to strong baselines.

## introduction large language models (lms) have had immense success recently @xcite @xcite , becoming a useful tool across disciplines. however, their success comes at a computational cost, due to increasing parameter counts for storing world knowledge @xcite and growing context lengths that enable access to distant information, but incur a quadratic complexity penalty. retrieval-augmented language modeling (ralm) alleviates this cost @xcite @xcite , as precise retrieval of relevant information can reduce memory and computation requirements. moreover, ralm is beneficial for factuality, freshness, and generalization without necessitating retraining, simply by swapping the retrieval index @xcite @xcite . however, past work on ralm has by and large not trained the retriever as a first-class component of the lm. in some cases @xcite @xcite , the retriever was used only at test time, or remained fixed throughout training, preventing it from adapting to the lm generator. in other cases, the retriever component was jointly trained but only after a separate pretraining phase for both the retriever and lm @xcite @xcite . thus, the retriever was not pre-trained from scratch with the lm, and only a fraction of the training budget was allocated for joint training. recently, @xcite presented a retrieval-augmented lm that trains a retriever from scratch jointly with the lm, but (a) the retriever was trained to exploit lexical information only, and (b) the retrieved information was not fused at the representation level back into the lm. in this work, we present the retrieval-pretrained transformer (rpt), a retrievalaugmented lm, where the retriever is a first-class component, trained jointly from scratch with the lm. rpt relies on two technical contributions. first, on the architecture side (see figure 1 ), input representations for the retriever are computed from the lm representations themselves (a concept we dub self-retrieval), and retrieved representations are fused back into the lm decoder for making next word predictions. second, we train the retriever with an auxiliary loss function that encourages retrieving text fragments that increase figure 1 : retrieval-pretrained transformer (rpt) is a language trained from scratch with a native retrieval ability that can be applied to long texts (e.g., books). rpt takes a chunk of text as input, retrieves semantically relevant chunks from the past to better predict the next chunk, and fuses these retrieved chunks into its representations. on top of a standard lm loss, the retriever is trained to retrieve chunks that increase the probability of the next chunk according to a reference lm. the probability of generating the subsequent text. specifically, given a recently generated chunk c t , the retriever is trained to retrieve chunks c i that increase the probability of p scoring (c t+1 | c i , c t ) according to a reference scoring lm. figure 1 provides an illustrative example for a case where a crime scene is described, and a scoring lm shows the benefit of retrieving a chunk thousands of tokens away (chunk 13) compared to lexical retrieval, which leads to a chunk that is only superficially related (chunk 100). unlike existing retrieval-augmented models that use an auxiliary encoder for retrieval @xcite @xcite , rpt is able to leverage its internal hidden states for retrieval after a single pre-training stage, greatly simplifying joint training. we apply rpt to the problem of modeling long documents, such as books, articles, and code, as those are naturally occurring examples of long-form content, where the entire index can be held within memory in a forward-pass. we evaluate rpt on four language modeling tasks and find that it improves perplexity across all tasks, outperforming prior work @xcite as well as strong baselines @xcite . moreover, we show that rpt retrieves high-quality chunks compared to retrievers that rely on lexical information. based on our empirical findings, we argue rpt can pave the way toward a next generation of pre-trained lms, where large corpora are used during pre-training, resulting in a language models where retrieval is a strongly embedded component. our code is publicly available at https://github.com/ohadrubin /rpt . ## background to situate our contribution, we review relevant recent ralm work. we extend this to more related work in §6. early work on ralms, such as knn-lm @xcite , used retrieval to improve language modeling by interpolating the next-word distribution produced by the lm with a distribution proposed through a test-time-only retrieval mechanism. @xcite later proposed chunked cross-attention (cca), where retrieval is performed also at training time, and retrieval results are deeply fused into the representations produced by a transformer decoder through attention. however, the retriever was trained separately and kept fixed during training, which prevented it from adapting to the lm over the course of training. trime @xcite , like this work, trained a retrieval-augmented lm from scratch where the retriever component and the decoder lm are trained jointly. our work differs from trime in two aspects: first, trime, like knn-lm, incorporates information from the retriever in a shallow manner through distribution interpolation, while we adopt cca as a deeper fusion mechanism. second, trime takes advantage of lexical clues for supervising the retrieverthat is, given a query, the trime retriever learns to retrieve contexts that will lead to generating the same token as the query. we, on the other hand, use a scoring lm to evaluate what text chunks are relevant for increasing the probability of the chunk being generated, which leads to more semantic retrieval. this is similar to epr (rubin layers also include chunked cross-attention layers that fuse information from retrieved chunks. the right side shows the retriever, which takes a chunk and retrieves the highest-scoring k chunks that appeared earlier in the document. et al., 2022), which used this idea for learning to retrieve prompts for in-context learning, and perplexity distillation in atlas @xcite . however, atlas does not train the retriever and lm from scratch and is an encoder-decoder model, more suitable for knowledge-intensive tasks. we, conversely, train from scratch and use a decoder model, more suitable for modeling long texts. we present our method in two parts. first, our architecture ( §3.1), which leverages cca to fuse retrieved representations into the lm, but adds a learned retriever component. second, we present the training method ( §3.2- §3.3), where the retriever is trained to retrieve chunks useful for generating a future chunk according to a reference lm. ## long-range lm datasets we evaluate rpt on four datasets, covering domains such as books, code, and mathematical writing, which require the ability to recall information over long distances. provide statistics on dataset size and the distribution over document length, showing that documents are long across all datasets and in particular pg19 and books3, where documents typically contain 10 foot_4 tokens or more. we briefly review the datasets. pg19 introduced in @xcite , pg19 is a widely used long-range language modeling benchmark containing books from project gutenberg, and covering a wide range of literary genres, styles, and topics. we adopt the exact setup and data split from prior work @xcite @xcite . books3 is a corpus of books released as part of the pile @xcite , containing a vast collection of literary works from different domains. to our knowledge, we are the first to use this corpus as a long-range language modeling benchmark. 4 codeparrot @xcite is a corpus of clean, nearly deduplicated python code from various github repositories. modeling code requires understanding patterns and contextualizing information over long distances, making it a natural candidate for testing long-range lms. in our experiments, we follow the approach of wu et al. arxiv is a corpus of preprint papers extracted from arxiv. it consists of mathematical texts that require maintaining coherence and referring to previously mentioned information over extended text. prior work evaluated long-range lms on this corpus @xcite @xcite , but did not release their corpus. thus, we use the preprocessed corpus and data splits made available by @xcite . ## experiments we now turn to experiments for comparing rpt to prior work across our four datasets. ## conclusion in this work, we present the retrieval-pretrained transformer (rpt), a retrieval-augmented lm where the retriever is trained as a native component of the lm to retrieve semantically relevant chunks for future text prediction. we evaluate rpt on four long-range language modeling tasks, including books, code, and mathematical writing. we demonstrate that by seamlessly integrating the retriever into the architecture and training process, rpt benefits from the fusion of retrieved context, improving over strong retrieval-augmented baselines. while this work focuses on retrieval from long texts, we argue our empirical findings show that adapting our procedure for general web-based corpora retrieval is an exciting future direction. this will require overcoming technical difficulties related to scaling and pretraining corpus construction. we envision rpt will pave the way for a new generation of pretrained language models with retrieval deeply integrated throughout their architecture and training process.