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 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 for 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.

翻译：检索增强语言模型近年来备受关注。然而，检索器通常并非作为语言模型的原生组件进行联合训练，而是被添加至预训练后的语言模型之上，这限制了语言模型与检索器相互适应的能力。本文提出检索预训练Transformer（RPT），这是一种从零开始联合训练检索增强语言模型以处理长文本任务的架构与训练流程。对于长文档中新生成的文本片段，语言模型计算查询表示，进而用于检索文档中可能位于数万词元之前的早期片段。来自被检索片段的信息被融合至语言模型表示中，用于预测下一个目标片段。我们采用语义目标训练检索器组件，其目标是根据参考语言模型检索能够提高下一片段概率的片段。我们在涵盖书籍、代码与数学写作的四项长距离语言建模任务上评估RPT，结果表明相较于强基线方法，RPT在检索质量乃至整体困惑度上均有显著提升。