Improving language models by retrieving from trillions of tokens

RETRO (Retrieval-Enhanced Transformer) adds a lookup step to an autoregressive language model. The input is split into chunks of 64 tokens. For each chunk, the model retrieves its k nearest neighbors from a database, where each entry is a 64-token chunk plus the 64 tokens that followed it in the original document. Similarity is the distance between frozen BERT embeddings of the chunks, searched approximately with ScaNN, which the authors say takes about 10 ms over a 2-trillion-token database.

The retrieved text goes through a small transformer encoder, and the main model attends to it with chunked cross-attention: tokens in chunk u+1 can attend to the neighbors retrieved for chunk u, which keeps the model autoregressive. Cross-attention layers are inserted every third block. Because BERT is frozen, all neighbors for the training data can be computed in advance.

Training uses DeepMind's multilingual MassiveText data; during training the model retrieves from a 600-billion-token subset, and at evaluation from 1.75 trillion tokens. Models range from 172M to 7.5B parameters, about 8% to 30% more than matching baselines.

The gains are large and consistent. Retrieval improves every dataset in the main scaling comparison at every model size, by roughly what a 10 times larger model would give, and keeps improving as the database grows and as more neighbors are retrieved (up to about 40). On the Pile, the 7.5B RETRO beats Jurassic-1 (178B) and Gopher (280B) on most subsets. An existing transformer can be "retrofitted" by training only the new cross-attention and encoder weights on 3% of the original data.

The authors also take test leakage seriously. Because the model can copy from its database, they measure how much evaluation text overlaps database chunks and show results as a function of that overlap. Retrieval helps even on chunks with little overlap, but much of the most dramatic result, on Wikitext103, comes from overlap.

Retrieval can stand in for model size. The central result is that a small model with a large, relevant database can match much larger models on knowledge-heavy text. For customers, this is the argument for retrieval-augmented systems over ever-larger models: knowledge lives in a database that can be updated, filtered and audited.

Index at the chunk level. RETRO stores one key per 64-token chunk, not per token. That is why it could index trillions of tokens; token-level kNN-LM would need about 4 KB per token. The same arithmetic applies to customer vector indexes: chunk size sets the index size, the retrieval granularity and the cost.

Check evaluation sets against the index. A retrieval model can copy evaluation text if it is in the database. RETRO's Wikitext103 perplexity drops from 22.96 to 3.92 with the full database, and the authors attribute this partly to overlap that their deduplication missed. When a customer's RAG system scores very well on a test set, check whether the test documents, or near copies, are in the index.

Measure with a leakage filter. The paper's method is reusable: for each evaluation chunk, find the longest run of tokens it shares with its retrieved neighbors, and report results separately for chunks with low overlap. That separates "found the answer" from "found the test".

Know what RETRO is not. It is retrieval built into pretraining, with a custom architecture. Most production RAG puts retrieved text into the prompt of an off-the-shelf model. RETRO's lessons transfer (scaling the database, chunk granularity, leakage), the architecture usually does not.