Improving language models by retrieving from trillions of tokens

In plain words

A language model guesses the next word from the words before it. Its knowledge sits in its weights, so the usual way to make it know more is to make it bigger and train it on more text. This paper tries a different route: the model looks up related passages from a very large text store while it writes, and reads them next to its own input.

The store holds trillions of words. Looking up single words would be far too costly, so the model works in blocks of text: it fetches passages that look like the block it has just read, then uses them to help predict the following block. The lookup index is computed ahead of time, so training does not have to search the store on the fly. Extra attention layers let the model blend the retrieved text into its predictions.

The authors report that a comparatively small model with this lookup holds its own against much larger models without it on several text benchmarks. They also show that an ordinary pretrained model can be given the lookup afterwards by training only the new parts, and that the store can be changed without retraining the model.

TL;DR

• The model predicts text in 64-token chunks and, for each chunk, retrieves two neighbour passages from a database of up to about 2 trillion tokens. Keys are frozen BERT embeddings §2.2 §2.3 §4.1.
• Retrieved passages enter through a small bidirectional encoder and a chunked cross-attention step in some decoder layers §2.4 alg. 1.
• Author-reported: the 7.5B model outperforms Jurassic-1 and Gopher on a majority of Pile test sets, and beats the 7B baseline on C4 and Wikitext103 §4.1 app. F.1 app. F.2.
• No official code was found. The paper's hyperparameters are partly inconsistent with each other, so check each value before implementing app. C.
• Implement the data pipeline and retrieval index as well as the model; the retrieval side is where most of the engineering cost sits §2.3 §4.1.

Problem

Scaling a language model usually means more parameters and more training data, and both raise training and inference cost. Parameters also serve as the model's memory, and the paper argues that part of that memory could come from a lookup instead §1. Earlier retrieval-for-language-modelling work used small Transformers (about 100M parameters) and databases of at most billions of tokens, and retrieved per token or per prompt §1 §3.1 tab. 3. The question is whether retrieval can be scaled to trillions of tokens while keeping the lookup cheap enough to run at training and sampling time §1 §2.3.

Method

RETRO splits each input sequence into chunks, retrieves neighbours for each chunk from a fixed database, and makes the retrieved text available to the chunks that follow. Retrieval is done per chunk, not per token. The database is built once from the training corpus with a frozen BERT encoder §1 §2.3.

1. Tokens and chunks. Text is tokenised with SentencePiece, vocabulary 128,000 §2.1 app. C.2. A training sequence has n = 2048 tokens, split into l chunks of m = 64 tokens §2.2.

2. Database. Each value is a pair [N, F] of two contiguous 64-token pieces: N is the neighbour used to compute the key, F is its continuation in the original document. The key is the BERT embedding of N averaged over time §2.3. The paper does not say how pieces are cut from documents or whether pairs overlap §2.3.

3. Retrieval. For each chunk C, take its k nearest keys under squared L2 distance on BERT embeddings. Neighbours from the same document as the training sequence are removed, so that no future chunk leaks into the past §2.3. Neighbours are retrieved in advance for training data §2.3.

4. Likelihood. Tokens of chunk u are predicted from the preceding tokens and from the retrieved neighbours of chunks 1 to u-1. The neighbours of chunk 1 are empty, so the first chunk does not depend on retrieval §2.2 eq. 1.

   L(X \mid \theta, \mathcal{D}) = \sum_{u=1}^{l} \sum_{i=1}^{m} \ell_\theta\left(x_{(u-1)m+i} \mid (x_j)_{j<(u-1)m+i},\ (\mathrm{Ret}_{\mathcal{D}}(C_{u'}))_{u'<u}\right) eq. 1

5. Decoder. The decoder interleaves standard blocks LM(H) = FFW(ATTN(H)) with RETRO blocks RETRO(H, E) = FFW(CCA(ATTN(H), E)) eq. 2. Layers that carry a RETRO block form the set P; the paper gives P per model size §2.4 app. C.1. The overall forward pass is listed as Algorithm 1 alg. 1.

6. Neighbour encoder. Each neighbour is encoded by a bidirectional Transformer. It cross-attends to the activations of the chunk that retrieved it, at the layers in P_enc, so the encoding depends on the chunk §2.4. The encoder has 2 layers with width d' = 896 and adds roughly 19M parameters §2.5.

7. Chunked cross-attention (CCA). Positions in chunk u+1 attend to the encoded neighbours of chunk u. Concretely, the last token of chunk u and the first m-1 tokens of chunk u+1 form the attending block for chunk u. The first m-1 tokens of the sequence, and any positions with no preceding chunk, pass through unchanged §2.4 eq. 3. The cross-attention itself is multi-head with relative positional logits added before the softmax §2.4 app. B.2.

   \mathrm{CCA}(H, E)_{um+i-1} = \mathrm{CA}(h_{um+i-1}, E_u) eq. 3

   \mathrm{CA}(h, Y) = \mathrm{softmax}(Y K Q^\top h)\, Y V eq. 4

8. Sampling. After each chunk is generated, its BERT embedding is used to retrieve neighbours, which are encoded and condition the next chunk. Retrieval cost grows linearly with the number of chunks §2.4.

9. Baseline. Retrieval-free baselines are decoder-only Transformers similar to GPT-2, with RMSNorm in place of LayerNorm and relative position encodings. They are trained at 132M, 368M, 1.3B and 7.0B parameters, not counting embeddings §2.5 tab. 2.

10. RETRO-fitting. Freeze a pretrained baseline. Add randomly initialised CCA and retrieval-encoder weights in every third layer from layer 6, and train only those §4.2 app. B.5. Without retrieval, the frozen model's output is unchanged §4.2.

11. Leakage-aware evaluation. For each evaluation chunk, find its 10 nearest training neighbours (up to 128 tokens each) and measure the longest token substring shared with them. Report bits-per-byte only on chunks whose overlap ratio is at most alpha §2.6 eq. 5.

\mathrm{bpb}(\alpha) = \frac{\sum_{C \in \mathcal{C}_\alpha} \ell(C)}{\sum_{C \in \mathcal{C}_\alpha} N(C)}, \quad \mathcal{C}_\alpha = \{C : r(C) \le \alpha\} eq. 5

Lookup pseudocode, built from §2.3:

# Build the retrieval database  [§2.3]
# Each value holds a neighbour N and its continuation F (64 tokens each)
# The pairing of pieces from documents is not stated  [§2.3]

def build_database(corpus):
    database = empty_index()
    for each pair (N, F) taken from the corpus:
        # Key: frozen BERT embedding of N, averaged over time  [§2.3]
        key = mean_over_time(bert(N))                # (d_bert,)
        database.add(key, value=(N, F))
    # Approximate nearest-neighbour index, built once  [§2.3]
    return database

# Retrieve neighbours for one chunk  [§2.3]
# Used in advance for training data, and after each chunk when sampling  [§2.4]

def retrieve(chunk, document_id, k, database):
    query = mean_over_time(bert(chunk))              # (d_bert,)
    # Squared L2 distance, approximate search (SCaNN in the paper)  [§2.3]
    ranked = database.search(query, metric="l2_squared")
    # Drop neighbours from the same document; keep the first k  [§2.3]
    kept = [n for n in ranked if n.document != document_id]
    return kept[:k]

Forward pass, following Algorithm 1 alg. 1:

# RETRO forward pass  [alg. 1] [§2.4]
# Input:  tokens (B, n); neighbours (B, l, k, r) from the lookup above
# Output: logits (B, n, V)

def retro_forward(tokens, neighbours):
    hidden = embed(tokens)                           # (B, n, d)
    for layer in 1..L:
        # Causal self-attention over the whole sequence  [alg. 1]
        hidden = causal_attention(hidden)
        if layer == min(P):
            # Encode neighbours, conditioned on this layer's activations  [§2.4]
            encoded = retrieval_encoder(neighbours, hidden)
        if layer in P:
            # Chunked cross-attention to the neighbours  [§2.4] [eq. 3]
            hidden = chunked_cross_attention(hidden, encoded)
        hidden = feed_forward(hidden)                # (B, n, d)
    return readout(hidden)                           # (B, n, V)

# Chunked cross-attention  [§2.4] [app. B.3]
# Positions are 1-based, as in the paper; m is the chunk length  [§2.2]

def chunked_cross_attention(hidden, encoded):
    output = hidden                                  # (B, n, d)
    for chunk u in 1..l:
        # Positions u*m .. u*m+m-1 see only the neighbours of chunk u,
        # clipped to the sequence end  [§2.4]
        positions = range(u * m, min(u * m + m, n + 1))
        queries = hidden[positions]
        # Neighbour and time axes merged before attention  [§2.4]
        output[positions] = cross_attention(queries, merge(encoded[u]))
        # Positions 1..m-1 keep their input values  [§2.4]
    return output

Hyperparameters

Values are given as the paper states them. Where the paper gives two different values, both are listed and the conflict is described in Pitfalls.

Results

Author-reported. Setup: the models were trained on MassiveText; evaluation numbers are given per dataset and model size §4.1 app. F.1.

Language modelling at 7.5B, from the full results table app. F.1:

Wikitext103 perplexity, validation and test. The RETRO rows are the baseline fine-tuned with retrieval from the stated set tab. 4 app. C.2:

The paper says the MassiveText 100% result is partly due to overlap with Wikitext103 that deduplication did not catch tab. 4.

Question answering, Natural Questions exact match tab. 5:

Other results, stated in the text:

• The Pile: the 7.5B RETRO model beats Jurassic-1 (178B) and Gopher (280B) on a majority of test subsets. Enron Emails and YouTube Subtitles were excluded §4.1 app. F.2.
• Scaling: retrieval gains stay roughly constant from 150M to 7B parameters, and a retrieval model's gain is comparable to a model roughly ten times larger on some datasets fig. 1 §5.
• Database size: moving retrieval from Wikipedia (4B tokens) to all of MassiveText (1.7T tokens) gives large gains, with leakage a partial cause §4.1 §4.4.
• Ablations on a 247M model trained on 157B tokens (C4 bits per byte, lower is better): full RETRO 0.822; no retrieval 0.987; neighbours only 0.950; continuations only 0.895; one training neighbour 0.858 and four 0.847 (default two, 0.822); cross-attention every third layer from layer 1 0.823; cross-attention in all layers 0.860 app. D.
• Leakage: RETRO still beats the baseline at every leakage level down to 12.5% overlap §4.4 fig. 6.
• Storage: a Wikipedia index is about 215 GB and a MassiveText index about 93 TB; a per-token kNN-LM index for Wikipedia would be about 15 TB §4.1.

Paper vs code

No official code.

Searches done: the only code links in the paper are third-party tools (FiD, used for the QA data §4.3; mwparserfromhell, used for Wikipedia parsing app. A.2); there is no project page; a GitHub repository search for retro org:deepmind found nothing, and searches for retro chunked cross-attention and retrieval-enhanced transformer retro deepmind found one third-party reimplementation with no license, which was not used. The paper says the models were implemented in JAX with Haiku §2.5, but no release was found.

Pitfalls

• Parameter counts disagree. Table 2 gives RETRO sizes of 172M, 425M, 1,451M and 7,532M and baselines of 132M, 368M, 1,309M and 6,982M tab. 2. The hyperparameter table in app. C.1 lists 247M, 564M, 1,574M and 7,505M under the heading "Baseline". Neither set matches the other; choose one and count your model's parameters to confirm.
• Where retrieval layers sit disagrees. The text says a RETRO block appears every three blocks from layer 6 §2.5, and the fine-tuning text says the same app. B.5. The pretraining table uses P = [9, 12, ..., 24] and [9, 12, ..., 32] for the two largest models app. C.1.
• Wikitext103 baseline depth. The text says 24 layers §4.1; the table says 18 layers app. C.2. The head count, key size and width agree.
• kNN-LM interpolation is written two ways. The main text gives p = lambda p_kNN + (1 - lambda) p_LM §4.1; the appendix gives p = lambda p_LM + (1 - lambda) p_kNN app. C.2. The lambda is 0.118 in the text and 0.12 in the figure caption fig. 5. Pick one formula and check it on validation data.
• Pile claim against Table 13. The text says RETRO does not beat the 7B baseline on dm_mathematics or ubuntu_irc §4.1. The table gives lower bits per byte for RETRO than for the baseline on both (1.164 vs 1.177; 1.178 vs 1.200) app. F.2. Check the subsets you use against the table.
• Database size. §2.1 gives 1.75T tokens for the evaluation database. Table 4 gives 1792B database tokens for MassiveText at 100% tab. 4. Both are in the paper, but the Table 4 setting is not fully described.
• "RETRO[Off]" is not defined. The paper says the model runs without retrieval data at evaluation time §1 app. F.1. How the neighbour input is removed (empty inputs, zero vectors, or skipping the CCA layers) is not stated.
• Residual and identity details in CCA. Equation 4 has no residual term, but the text calls the operators residual §2.4. The listing returns only the attention output, with no residual, no multi-head split and no relative positional logits app. B.3. Follow the text for the residual form and the listing for the shape handling.
• Neighbour pairing and index settings are missing. The paper does not state how documents are cut into the 64-token N and F pieces, whether they overlap, the BERT size and tokenisation, or the SCaNN index settings. Each of these changes retrieval quality and should be set and reported.
• Causality and leakage. Neighbours from the training document itself are removed §2.3. Any implementation that keeps them lets the model read the future of its own sequence.
• Cost is mostly in the index. Neighbours are precomputed for training; on-the-fly lookup is described as too slow §2.3. Plan storage for the key-value database (about 93 TB for MassiveText in the paper §4.1).
• Training and evaluation databases differ. Training retrieves from 600B tokens, evaluation from 1.75T §2.1. The paper's gains at evaluation time depend partly on this larger database and partly on leakage §4.4.
• Leakage is a real confound. On C4 and Wikitext103 the RETRO model exploits leaked training data more than the baseline does §4.4. Use the leakage-filtered evaluation (chunks below the overlap threshold) when comparing models §2.6.
• Fine-tuning freezes the base model. The base weights stay frozen so that retrieval-off performance is unchanged; unfreezing the base model made that performance degrade app. B.5.
• Sampling cost. Sampling is quadratic in length, as for ordinary Transformers, with retrieval adding cost linear in the number of chunks §2.4.
• Updating the store. The paper says the database can be changed at inference time, or have data removed, without retraining the model §3.2. It does not test this, and the card's results do not measure it.

Lineage

• Builds on: Transformer decoders and encoders §2.5; frozen BERT embeddings for keys §2.3; approximate nearest-neighbour search (SCaNN) §2.3.
• Compared against: kNN-LM and continuous cache for token-level retrieval; SPALM; DPR, REALM, RAG, FiD and EMDR2 for prompt-level retrieval §3.1 tab. 3. The Wikitext103 and QA comparisons use these methods tab. 4 tab. 5.
• Superseded by: not recorded.