# Long-range Language Modeling with Self-retrieval

**Authors:** Ohad Rubin, Jonathan Berant  
**arXiv:** [2306.13421](https://arxiv.org/abs/2306.13421)  
**Date:** 2023-06-23  
**Summary page:** https://www.summarizepaper.com/en/paper/2306.13421v1/

## Key points

- Retrieval-augmented language models (LMs) are gaining attention
- Retrieval-Pretrained Transformer (RPT) is proposed to jointly train a retrieval-augmented LM from scratch for long-range language modeling tasks
- RPT model computes query representations for recently generated text chunks in a long document and uses them to retrieve earlier chunks located potentially tens of thousands of tokens before
- Information from retrieved chunks is fused into the LM representations to predict the next target chunk
- The retriever component is trained with a semantic objective that aims to retrieve chunks that increase the probability of the next chunk according to a reference LM
- RPT model evaluated on four long-range language modeling tasks spanning books, code, and mathematical writing, where documents are generally long across all datasets
- RPT model has 12 layers with hidden dimension d=1024 and eight attention heads with head dimension 128; CCA is applied every two layers, and two neighbors are used unless mentioned otherwise.
- RPT demonstrates promising results in long range language modeling tasks by improving retrieval quality and perplexity compared to strong baselines.

## Summary

Retrieval-augmented language models (LMs) have gained significant attention in recent times. To address the issue of limited ability to adapt between the retriever and LM components when added separately, this work proposes the Retrieval-Pretrained Transformer (RPT), an architecture and training procedure for jointly training a retrieval-augmented LM from scratch for long-range language modeling tasks. The RPT model computes query representations for recently generated text chunks in a long document and uses them to retrieve earlier chunks located potentially tens of thousands of tokens before. Information from retrieved chunks is fused into the LM representations to predict the next target chunk. The retriever component is trained with a semantic objective that aims to retrieve chunks that increase the probability of the next chunk according to a reference LM. The RPT model is evaluated on four long-range language modeling tasks spanning books, code, and mathematical writing, where documents are generally long across all datasets. At training time, sequences of length L=16384 tokens are split into four devices each consuming 4096 tokens. The decoder stack takes 2048 tokens as input in a sliding window approach containing ℓ=32 chunks of length m=64. Rotary Positional embedding is employed while all models are trained for 500K steps on TPUv4-64 with an effective batch size of 217 tokens. The GPT-NeoX tokenizer is used for all models trained and Pythia serves as our scoring language model with deduplicated 1.4B parameter version used for scoring top-20 BM25 candidates. The RPT model has 12 layers with hidden dimension d=1024 and eight attention heads with head dimension 128; CCA is applied every two layers, and two neighbors are used unless mentioned otherwise. The Books3 dataset is a corpus of books released as part of the Pile, containing a vast collection of literary works from different domains which we use as a long range language modeling benchmark for the first time. Overall, RPT demonstrates promising results in long range language modeling tasks by improving retrieval quality and perplexity compared to strong baselines.