Source-linked AI summary
Improving language models by retrieving from trillions of tokens
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, Laurent Sifre
TL;DR
The paper addresses how to improve autoregressive language models by giving them access to massive external memory without relying solely on parameter scaling. RETRO retrieves locally similar text chunks with a frozen BERT retriever and integrates them through chunked cross-attention. It achieves strong language-modeling results at large scale, remains competitive on question answering, and supports flexible retrieval and rapid retrofitting, while raising privacy, safety, and fairness concerns.
Problem
The paper explores efficient ways to augment language models with massive-scale memory without significantly increasing computation or model size.
Method
RETRO retrieves similar contiguous text chunks using frozen BERT embeddings and incorporates them into autoregressive predictions through chunked cross-attention.
Results
RETRO outperforms previous large-scale language models on Wikitext103 and the Pile and is competitive on retrieval-intensive tasks such as question answering.
Takeaways & Limitations
Semi-parametric retrieval provides an orthogonal and more efficient alternative to raw parameter scaling for building more powerful language models.
Takeaways & Limitations
Retrieval models can directly copy training data and may exacerbate privacy, safety, and fairness issues, including biases in the data and retrieval selection mechanism.
Abstract
from arXiv · showhide
We enhance auto-regressive language models by conditioning on document chunks retrieved from a large corpus, based on local similarity with preceding tokens. With a $2$ trillion token database, our Retrieval-Enhanced Transformer (RETRO) obtains comparable performance to GPT-3 and Jurassic-1 on the Pile, despite using 25$\times$ fewer parameters. After fine-tuning, RETRO performance translates to downstream knowledge-intensive tasks such as question answering. RETRO combines a frozen Bert retriever, a differentiable encoder and a chunked cross-attention mechanism to predict tokens based on an order of magnitude more data than what is typically consumed during training. We typically train RETRO from scratch, yet can also rapidly RETROfit pre-trained transformers with retrieval and still achieve good performance. Our work opens up new avenues for improving language models through explicit memory at unprecedented scale.
1. Introduction
RETRO augments autoregressive language models with retrieval from a massive text database, providing an alternative to scaling model parameters, data, or compute alone. Its gains increase with retrieval resources and remain effective across model scales.
- Motivation: RETRO equips language models to retrieve similar text chunks from a large database and use them to predict tokens.This semi-parametric approach complements parametric scaling by directly accessing external text memory.
- Contributions: RETRO uses chunked cross-attention to incorporate retrieved text and a frozen BERT retriever to avoid updating the retrieval network.The architecture is designed to scale retrieval to very large databases while limiting computational overhead.
- Scaling: RETRO provides a constant gain for models ranging from 150M to 7B parameters and can improve when the retrieval database or neighbour count increases.On C4, gains increase with retrieval database size and with up to 40 retrieved neighbours, after which performance begins to degrade.
- Results: The largest RETRO model obtains state-of-the-art results on Wikitext103 and the Pile and achieves competitive performance on downstream question answering.The evaluation also addresses test-set leakage by accounting for proximity between test documents and training data.
2. Method
RETRO augments autoregressive language models with chunk-level retrieval from a massive key-value database, then integrates retrieved text through an encoder and chunked cross-attention while preserving causality.
- Method overview: The architecture is designed to retrieve at the contiguous-chunk level, reducing storage and computation requirements while supporting trillion-token databases.Training uses a frozen BERT key encoder and augments each sequence chunk with retrieved neighbours.
- Retrieval-enhanced autoregressive token models: RETRO splits each input sequence into chunks and augments them with nearest-neighbour text retrieved from a database of token chunks.The method uses 2048-token sequences divided into 64-token chunks and retrieves k neighbours for each chunk.
- Nearest neighbour retrieval: The retrieval database stores contiguous token chunks as values and frozen BERT embeddings as keys, enabling approximate nearest-neighbour search by embedding distance.Each retrieved value contains a neighbour chunk and its continuation; same-document neighbours are filtered to preserve causality.
- Nearest neighbour retrieval: Approximate retrieval queries a 2 trillion token database in 10 ms, with the cost amortized over a chunk length.The method uses SCaNN and precomputes embeddings because on-the-fly retrieval is too slow for training.
- Retro model architecture: A bidirectional encoder processes retrieved neighbours, while decoder layers interleave standard Transformer blocks with retrieval blocks that use cross-attention.The retrieval encoder is conditioned on the corresponding chunk activations, making neighbour representations differentiable with respect to the retrieving chunk.
- Retro model architecture: Chunked cross-attention preserves autoregressive generation by restricting each position to its current or earlier retrieval context.The first m−1 tokens of a chunk cannot attend to neighbours from previous chunks, while the final token attends to the chunk’s retrieval set.
3. Related Work
Prior retrieval methods range from lexical and latent-topic matching to dense caches, trained retrievers, and retrieval-augmented Transformers. RETRO combines frozen retrieval representations with chunk-level retrieval for long-sequence language modeling, while retrieval also raises privacy, safety, fairness, and interpretability considerations.
- Retrieval for language modelling: Earlier text retrieval used inverted-index methods such as TF-IDF and BM25, while later work adopted latent topic models and neural representations.These approaches identify relevant neighbours using lexical matching, topic structure, or learned activation similarity.
- Retrieval for language modelling: Continuous cache and kNN-LM interpolate retrieved-token distributions at inference, improving evaluation without modifying the underlying neural language model.Their plug-in design avoids additional training but limits the model’s ability to reason about retrieved text.
- Retrieval for language modelling: Dense retrievers such as DPR train representations for question-answer matching, whereas RETRO uses frozen retrieval representations for long-sequence language modeling.RETRO retrieves different documents for different sequence chunks and repeatedly retrieves during generation.
- Retrieval for language modelling: Unlike methods that prepend retrieved documents or retrieve once, RETRO encodes neighbours separately and assembles them through chunked cross-attention.Its chunk-level design supports sub-sequence reasoning and repeated retrieval while generating a sequence.
- Privacy, safety and fairness: Retrieval can support model updating and data filtering more cheaply than retraining, but it may also expose memorized data and amplify bias or toxicity.The paper identifies direct copying of training data, retrieval-selection bias, and insufficiently analyzed effects on toxic outputs as concerns.
- Privacy, safety and fairness: Because retrieved neighbours can be visualized or modified, retrieval may make model outputs more transparent and interpretable.The paper illustrates this claim with examples in several tables.
4. Results
RETRO improves language-modelling performance across model sizes and retrieval-database scales, with gains varying by dataset and partly reflecting leakage. It also compares favorably with larger models, supports rapid retrofitting, and remains effective on retrieval-intensive tasks.
- Data scaling: Retrieval-database expansion from Wikipedia’s 4B tokens to MassiveText’s 1.7T tokens produces dramatic language-modelling gains.Increasing retrieved neighbours from 1 to 10 consistently improves models; larger models can exploit more neighbours.
- The Pile: On the Pile, Retro 7.5B outperforms the baseline on all test sets and surpasses Jurassic-1 and Gopher on a majority of them.Retro underperforms the baseline on dm_mathematics and ubuntu_irc and underperforms Jurassic-1 on those subsets.
- Wikitext103: On Wikitext103, MassiveText retrieval yields much better perplexity than Wikipedia retrieval, although the strongest result is partly attributable to overlap missed by deduplication.Wikipedia retrieval performs similarly to the authors’ kNN-LM implementation.
- Efficiency: Retro requires 215GB to index Wikipedia and 93TB for MassiveText, whereas token-level kNN-LM requires 15TB for Wikipedia, motivating chunk-level retrieval.Chunk-level retrieval is presented as necessary for databases containing trillions of tokens.
- Retrofitting and downstream tasks: Retro can be rapidly fine-tuned by updating only retrieval-related weights, using less than 10% of the 7B model’s weights and 3% of pre-training sequences.This preserves the original model’s performance when retrieval is disabled; on question answering, Retro is competitive with Realm, RAG, and DPR but underperforms FiD.
- Leakage analysis: Retro outperforms baselines even at α=12.5% leakage, indicating gains on both locally similar and syntactically different evaluation chunks.On Curation Corpus, which has no training leakage by design, retrieval provides a constant offset; on leaked datasets, Retro exploits leakage more strongly.
- Sampling: Retrieved neighbours influence sampled text, with reported examples showing reduced hallucinations and improved knowledge compared with retrieval-disabled sampling.One example continues the opening of Hamlet with only a few mistakes.
5. Conclusion
The paper presents Retro as a retrieval-enhanced language model that scales access to trillions of tokens without relying solely on parameter growth. It reports durable gains across model sizes, competitiveness with larger models and retrieval-based tasks, and efficient retrofitting with limited leakage contribution.
- Contribution: Retro models retrieve from databases with trillions of tokens, scaling the data available to language models by an order of magnitude beyond typical training consumption.The method is framed as a semi-parametric alternative to raw parameter scaling.
- Results: Retro gains persist through at least 7B parameters and correspond on some datasets to non-retrieval models with 10× more parameters.On Wikitext103 and the Pile, Retro outperforms previous models trained on large-scale datasets and is competitive on question answering.
- Flexibility: Baseline transformers can be rapidly fine-tuned into Retro models, while Retro can also operate without retrieval and retain comparable baseline performance.Only a modest fraction of the gains is attributed to test-set leakage.
- Conclusion: The work argues that semi-parametric retrieval offers an orthogonal, more efficient route to stronger language models than raw parameter scaling.This conclusion is stated within the paper’s scope of large-scale language modelling.
A. Datasets
The paper uses MassiveText and a curated set of recent Wikipedia articles as large-scale training and retrieval resources. It describes their composition, sampling, cleaning, and safeguards against train–test overlap.
- Documentation: The paper provides a full description of MassiveText, its source and language breakdown, and the complete list of September 2021 evaluation articles.These resources are documented in the associated tables and dataset descriptions.
- MassiveText: The retrieval database uses the full MassiveText union except for a 4% book subsample, with source proportions matching training sampling frequencies.Table 8 identifies the sampling weight for each dataset during training.
- Wikipedia evaluation set: The September 2021 Wikipedia evaluation set contains articles added or heavily edited after training-data collection, with high-overlap articles filtered out.The dataset is designed to limit leakage from training data while retaining recent articles.
- Preprocessing: Wikipedia preprocessing removes specified reference-related sections, wikilinks, templates, and objects tagged as references or tables before concatenating the remaining text.The title and cleaned sections are joined with blank-line delimiters.
B. Details on the retrieval architecture
RETRO combines chunk-level retrieval with an encoder-decoder Transformer, using cross-attention and relative positional information to condition decoding on retrieved neighbours. Its implementation separates encoder and decoder representations and supports both training from scratch and retrofitting pretrained models.
- Cross-attention: Chunked cross-attention merges neighbour and time dimensions before applying attention, using relative positional encodings to reflect chunk alignment.Positional logits are added to content logits in the attention computation.
- Representations: The architecture uses a frozen retrieval representation and can assign different dimensionalities to encoder and decoder embeddings.The encoder is typically narrower than the decoder, with dEnc = 896 and decoder dimensionality scaled up to d = 8192.
- Retrofitting: Retrofit initializes retrieval and cross-attention weights while freezing pretrained model weights, adding cross-attention every third layer from layer six.The retrofit experiments use newly initialized retrieval components while retaining the pretrained transformer parameters.
- Implementation: The simplified implementation uses 128-token sequences, 16-token chunks, four neighbours, and 32-token retrieval lengths.These values define the example dimensions used in the chunked cross-attention listing.
C.2. Wikitext103 comparison
On Wikitext103, retrofitting a pretrained baseline with retrieval improves validation perplexity, while training RETRO from scratch can overfit this small dataset and underperform the baseline.
- Wikitext103 results: 18.46 validation perplexity is achieved by the best Retrofit checkpoint retrieving from Wikipedia, versus 21.58 for the best 35,000-step baseline checkpoint.The Retrofit starts from the 35,000-step baseline checkpoint and trains only the new retrieval weights.
- Wikitext103 results: Training RETRO from scratch produces weaker results than the baseline on Wikitext103 when retrieving from Wikipedia.The authors attribute this result to increased over-fitting from RETRO’s additional weights on the small dataset.
- kNN-LM comparison: The kNN-LM comparison tunes λ to 0.118 and then α to 0.00785, with Figure 7 showing perplexity as each parameter varies.The initial α value is 0.0012 before tuning λ.
D. Model ablations
Ablations identify design choices that improve efficiency or retrieval gains, including relative encodings, previous-chunk conditioning, two-neighbour training, and cross-attention every third layer. Deeper encoders add substantial time for only a tiny loss decrease.
- Cross-attention design: Relative encodings in cross-attention improve both the steps needed to reach a target performance and computational efficiency.The result is described as a pure improvement relative to the default without relative encodings.
- Cross-attention design: Conditioning the encoder on the previous chunk’s intermediate embeddings improves both training steps and computational efficiency.The ablation reports a pure improvement under both measures.
- Retrieved context: Attending to both neighbours and their continuations is the most efficient choice for final performance and training efficiency.Neighbours alone provide 22% of retrieval’s performance improvement, while their future continuation provides 56%.
- Encoder depth: A 3× deeper encoder decreases loss by 0.15% but increases training time by 20%, making a shallow encoder more efficient.The authors therefore select the shallow encoder for training efficiency.
- Number of neighbours: Training with two neighbours is the best efficiency choice: one neighbour substantially reduces performance, while four add computational overhead without substantial final improvement.Additional neighbours can still be used during evaluation.
- Cross-attention frequency: Using cross-attention every third layer provides a good trade-off between performance and runtime, whereas attending only at the top or bottom layer performs poorly.Attending once at a mid-depth layer is relatively sound.
E. Qualitative experiments
The qualitative experiments examine RETRO through evaluation-sample perplexities and autoregressive sample generation.
- Qualitative experiments: RETRO is illustrated using the perplexity of evaluation samples and autoregressively generated samples.The section introduces both evaluation and generation-based views of model usage.
E.1. Inspecting neighbours and perplexities on evaluation data
The qualitative analysis identifies both useful factual transfer from retrieved neighbours and substantial evaluation leakage that RETRO can exploit.
- Specific facts can be extracted from preceding retrieved neighbours, significantly reducing RETRO’s loss for corresponding tokens.Examples include Publishers Weekly, Tyrone, and the dates 25 August to 6 September 2020.
- Evaluation data can partially leak into training and retrieval data despite deduplication, and RETRO can dramatically exploit that leakage.In one example, evaluation chunks largely overlap retrieved chunks, producing much lower loss for corresponding tokens.
E.2. Inspecting samples
The paper examines how retrieved neighbours relate to source documents and reports additional language-model evaluations and generated-sample analyses.
- Inspecting neighbours: Wikipedia retrieved neighbours tend to come from articles close to the source article in link distance.
- Inspecting neighbours: Retrieved-neighbour link distance increases on average with neighbour rank.
- Additional evaluations: The supplementary material reports further filtered language-model results on the Pile.
- Additional evaluations: RETRO and baseline models are evaluated on the Pile using bits-per-bytes.
- Additional evaluations: The paper compares RETRO against Jurassic-1 and reports the full bits-per-bytes results in a supplementary table.
F.3. Filtered results
Filtered analyses measure evaluation–training overlap and show that RETRO benefits more from leakage, while retaining significant improvement on original, unobserved chunks.
- Reported results: Table 14 contains full results for the main language-modeling datasets, with rows corresponding to the paper’s figures.
- Overlap analysis: Evaluation–training overlap is slight for C4 and Wikitext103, while Wikipedia September 2021 and Curation Corpus show almost no leakage.
- Filtered results: Filtered Pile analyses report chunk-overlap distributions and filtered performance curves.
- Filtered results: RETRO models exploit evaluation leakage more, but their performance improvement remains significant on original chunks absent from training.
- Reported results: Table 15 reports full Pile results in bits-per-bytes and includes Jurassic-1, GPT-3, and Gopher comparisons.
- Qualitative examples: Examples show retrieved text reducing loss for specific tokens and enabling copying or reformatting in generated samples.The examples include factual tokens, overlapping evaluation text, a legal declaration, and digits of π.