Improving language models by retrieving from trillions of tokens

Improving language models by retrieving from trillions of tokens

Sebastian Borgeaud$^{\dagger}$, Arthur Mensch$^{\dagger}$, Jordan Hoffmann$^{\dagger}$, 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$^{\ddagger}$, Erich Elsen$^{\ddagger}$ and Laurent Sifre$^{\dagger}$ $^{\ddagger}$
All authors from DeepMind

Abstract

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

$^{\dagger}$ Equal contributions
$^{\ddagger}$ Equal senior authorship
Corresponding authors: {sborgeaud|amensch|jordanhoffmann|sifre}@deepmind.com

Executive Summary: Executive Summary

The paper addresses a core challenge in scaling language models: performance gains have come mainly from increasing parameters, compute, and training data, but this couples memorization and computation in ways that are expensive and hard to update. As models grow beyond hundreds of billions of parameters, diminishing returns, high energy costs, and difficulties incorporating new knowledge or mitigating bias become acute. The authors explore retrieval from an external database as a complementary, semi-parametric route that effectively multiplies the data available at inference without a proportional increase in model size.

The work set out to demonstrate that conditioning an autoregressive transformer on chunks retrieved from a database of trillions of tokens can match or exceed the performance of much larger parametric models on language modelling and downstream tasks, while keeping the model size modest (up to 7 B non-embedding parameters).

RETRO augments standard transformers by splitting sequences into 64-token chunks, retrieving the nearest-neighbour chunks (and their continuations) from a frozen BERT-indexed key-value store using approximate nearest-neighbour search (SCaNN), and integrating them via a lightweight encoder plus chunked cross-attention layers inserted every third decoder layer. The approach was trained from scratch on a 5-trillion-token multilingual MassiveText corpus (with a 600 B–1.75 T retrieval subset) and also tested by “RETRO-fitting” existing pre-trained baselines. Evaluation used standard benchmarks (C4, Wikitext103, the Pile, Lambada, Curation Corpus) plus a held-out “future” Wikipedia set to probe leakage.

Key findings are that RETRO yields a roughly constant gain across model sizes from 150 M to 7 B parameters, equivalent to a ~10\times increase in parametric model size on several datasets; the 7.5 B RETRO model matches or beats GPT-3, Jurassic-1 (178 B), and Gopher (280 B) on most subsets of the Pile while using 25–40\times fewer parameters; performance improves steadily as the retrieval database scales from billions to two trillion tokens and as the number of neighbours rises to ~10–40; fine-tuning on Natural Questions reaches 45.5 % exact-match accuracy, competitive with prior retrieval-augmented QA systems; and filtered bit-per-byte analysis shows that roughly half the gains persist even on chunks with negligible overlap to the training set, indicating genuine generalization beyond simple copying.

These results matter because they demonstrate an orthogonal scaling axis—database size and retrieval integration—that is far cheaper to update than re-training model weights. Knowledge can be refreshed, biased or toxic content filtered, or private data removed simply by editing the database. The approach also improves sample interpretability, since retrieved neighbours can be inspected directly.

The authors recommend that future model development routinely incorporate retrieval at trillion-token scale, explore end-to-end learned retrievers, and investigate database filtering for safety and privacy. They caution that careful deduplication and leakage-aware evaluation remain essential, and note that their largest experiments still leave open questions about behaviour at even greater scales and about how retrieval interacts with instruction tuning or multimodal data. Overall, the work establishes retrieval as a practical, high-leverage complement to parameter scaling for building more capable and maintainable language models.

1. Introduction

Section Summary: Recent advances in language modeling have come from scaling up neural networks like Transformers, which predict text by processing ever-larger amounts of data and parameters. This paper explores an alternative by building a Retrieval Transformer, called Retro, that augments a standard model with direct access to a huge external database of text, retrieving similar passages to improve predictions without dramatically raising computation. The method scales across different model sizes, achieves strong results on benchmarks, and includes careful evaluation steps to avoid overestimating performance due to data overlap.

Language modelling (LM) is an unsupervised task that consists of modelling the probability of text, usually by factorising it into conditional next-token predictions $p(x_1, \dotsc, x_n) = \prod_i p(x_i | x_{<i})$. Neural networks have proven to be powerful language models, first in the form of recurrent architectures ([1, 2, 3]) and more recently in the form of Transformers ([4]), that use attention to contextualise the past. Large performance improvements have come from increasing the amount of data, training compute, or model parameters. Transformers have been scaled from $100$ million parameter models in seminal work to over hundred billion parameters ([5, 6]) in the last two years which has led to models that do very well on a wide array of tasks in a zero or few-shot formulation. Increasing model size predictably improves performance on a wide range of downstream tasks ([7]). The benefits of increasing the number of parameters come from two factors: additional computations at training and inference time, and increased memorization of the training data.

In this work, we endeavor to decouple these, by exploring efficient means of augmenting language models with a massive-scale memory without significantly increasing computations. Specifically, we suggest retrieval from a large text database as a complementary path to scaling language models. Instead of increasing the size of the model and training on more data, we equip models with the ability to directly access a large database to perform predictions—a semi-parametric approach. At a high level, our Retrieval Transformer ($\textsc{Retro}$) model splits the input sequence into chunks and retrieves text similar to the previous chunk to improve the predictions in the current chunk. Existing retrieval for language modelling work only considers small transformers ($100$ millions parameters) and databases of limited size (up to billions of tokens) ([8, 9, 10, 11]). To our knowledge, our work is the first to show the benefits of scaling the retrieval database to trillions of tokens for large parametric language models. Our main contributions are the following.

  • We introduce $\textsc{Retro}$, a retrieval-enhanced autoregressive language model (Section 2.2). We use a chunked cross-attention module to incorporate the retrieved text (Section 2.4), with time complexity linear in the amount of retrieved data. We show that retrieving based on a pre-trained frozen $\textsc{Bert}$ model (Section 2.3) works at scale, removing the need for training and updating a retriever network.
  • We show that our method scales well with model size and database size (Figure 1): $\textsc{Retro}$ provides a constant gain for models ranging from 150M to 7B parameters, and $\textsc{Retro}$ can be improved at evaluation time by increasing the database size and the number of retrieved neighbours. Our largest model obtains state-of-the-art results on a range of downstream evaluation datasets including Wikitext103 ([12]) and the Pile ([13]) (Section 4). We show that $\textsc{Retro}$ can be fine-tuned to achieve competitive performance on downstream tasks such as question answering (Section 4.3).
  • We propose an evaluation aware of proximity of test documents with the training set (Section 2.6), addressing the problem of test set leakage ([14]). This is relevant for all language models, and especially for retrieval-enhanced models since they have direct access to the training dataset during evaluation. Using this methodology, we show that the performance of $\textsc{Retro}$ comes from both explicit neighbour copying and general knowledge extraction (Section 4.4).

**Figure 1:** **Scaling of $\textsc{Retro}$.** The performance gain of our retrieval models remains constant with model scale (left), and is comparable to multiplying the parameteric model size by $\sim 10 \times$. The gain increases with the size of the retrieval database (middle) and the number of retrieved neighbours (right) on the C4 validation set, when using up to 40 neighbours. Past this, performance begins to degrade, perhaps due to the reduced quality. At evaluation $\textsc{Retro}$ can be used without retrieval data ($\textsc{Retro}$ [OFF]), bringing limited performance degradation compared to baseline transformers.

2. Method

Section Summary: The method introduces a retrieval-enhanced transformer called Retro that augments standard language modeling by pulling in relevant text passages from a massive database during both training and generation. It works at the level of fixed-size token chunks rather than single tokens: a frozen BERT model turns each chunk into an embedding used to find nearest neighbors in a key-value store holding trillions of tokens, and these neighbors are then fed into an encoder-decoder stack that combines them with the original sequence while preserving autoregressive order. The same large multilingual corpus supplies both the training data and the retrieval database, with careful filtering to limit overlap between train and test sets.

We design our retrieval-enhanced architecture to be capable of retrieving from a database with trillions of tokens. For this purpose, we retrieve at the level of contiguous token chunks instead of individual tokens which reduces storage and computation requirements by a large linear factor. Our method first constructs a key-value database, where values store raw chunks of text tokens and keys are frozen $\textsc{Bert}$ embedddings ([15]). We use a frozen model to avoid having to periodically re-compute embeddings over the entire database during training. Each training sequence is then split into chunks, which are augmented with their $k$-nearest neighbour retrieved from the database. An encoder-decoder architecture integrates retrieval chunks into the model's predictions. We summarize the $\textsc{Retro}$ architecture in Figure 2, and detail it in this section. We end the section by introducing a new methodology to evaluate language models when an evaluation set is partially present in the training set.

**Figure 2:** **$\textsc{Retro}$ architecture.** *Left:* simplified version where a sequence of length $n=12$ is split into $l=3$ chunks of size $m=4$. For each chunk, we retrieve $k=2$ neighbours of $r=5$ tokens each. The retrieval pathway is shown on top. *Right:* Details of the interactions in the $\textsc{Cca}$ operator. Causality is maintained as neighbours of the first chunk only affect the last token of the first chunk and tokens from the second chunk.

2.1 Training dataset

We use a multi-lingual version of MassiveText ([16]) for both training and retrieval data. The dataset consists of text documents from multiple sources and multiple languages totalling over 5 trillion tokens (detailed in Table 1). Sequences are sampled from subsets of the training data, with sampling weights given in the right-most column of Table 1. We tokenize the dataset using SentencePiece ([17]) with a vocabulary of 128, 000 tokens. During training (unless otherwise specified), we retrieve from 600B tokens from the training data. The training retrieval database is made of the same subsets as the training data, in proportion that matches the training sampling frequencies. During evaluation the retrieval database consists in the full union of these datasets, with the exception of books for which we use a sub-sample of 4%. The evaluation retrieval database thus contains 1.75T tokens. To limit test set leakage, we compute the $13$-gram Jaccard similarity between train and test documents using the MinHash scheme and remove all training documents with high similarity (0.8 or higher) to a validation or test set document. Additionally, we remove all validation and test articles from Wikitext103 ([12]) from our Wikipedia training data.

: Table 1: MassiveText. The last column indicates the sampling weight during training. The multilingual subsets include documents in 10 languages. The full breakdown is given in Appendix A.1.

Source Token count (M) Documents (M) Multilingual Sampling frequency
Web 977, 563 1, 208 Yes 55%
Books 3, 423, 740 20 No 25%
News 236, 918 398 No 10%
Wikipedia 13, 288 23 Yes 5%
GitHub 374, 952 143 No 5%

2.2 Retrieval-enhanced autoregressive token models

Our approach uses retrieval as a way to augment input examples at the granularity of small chunks of tokens. Formally, we consider sequences of integer tokens in $\mathbb{V} = [1, v]$, obtained using a text tokenizer[^1]. We split each $n$-token-long example $X = (x_1, \dots, x_n)$ into a sequence of $l$ chunks $(C_1, \dots, C_l)$ of size $m=\frac{n}{l}$, i.e. $C_1 \triangleq {(x_1, \dots, x_m), , \dots, , C_{l} \triangleq} (x_{n - m + 1}, \dots, x_n) \in \mathbb{V}^{m}$. We use $n=2048$ and $m=64$. We augment each chunk $C_u$ with a set $\textsc{Ret}\mathcal{D}(C_u)$ of $k$ neighbours from the database $\mathcal{D}$. $\textsc{Ret}\mathcal{D}$ (or $\textsc{Ret}$ for brevity) is a non-trainable operator specified in Section 2.3. Token likelihoods are provided by a model, parameterized by $\theta$, that takes as input both previous tokens and their retrieved neighbours. This defines the following retrieval-enhanced sequence log-likelihood:

[^1]: We use the notation $[1, v] \triangleq {1, \dots, v }$ throughout the text.

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

We set $\textsc{Ret}(C_1) = \emptyset$, namely the likelihood of tokens from the first chunk does not depend on any retrieval data. This likelihood definition preserves autoregressivity: the probability of the $i$-th token of the $u$-th chunk, $x_{(u-1) m + i}$, only depends on previously seen tokens $(x_j){1 \leq j < (u-1)m + i}$ and on the data retrieved from the previous chunks $(\textsc{Ret}(C{u'})){u'<u}$. We can therefore directly sample with log-probability $\ell$, where sampling within the chunk $C_u$ is conditioned on the neighbours $(\textsc{Ret}(C{u'}))_{u'<u}$. This makes retrieval-enhanced models directly comparable with the largest language models that are evaluated by sampling.

2.3 Nearest neighbour retrieval

Retrieval neighbours.

Our database consists of a key-value memory. Each value consists of two contiguous chunks of tokens which we denote $[N, F]$ where $N$ is the neighbour chunk which is used to compute the key, and $F$ is its continuation in the original document. The corresponding key is the $\textsc{Bert}$ embedding of $N$, averaged over time, that we denote $\textsc{Bert}(N)$. For each chunk $C$, we retrieve its approximate $k$-nearest neighbours from our key-value database using the $L_2$ distance on BERT embeddings $d(C, N) = || \textsc{Bert}(C) - \textsc{Bert}(N)||2^2$. The model receives the corresponding values $\textsc{Ret}(C) \triangleq ([N^1, F^1], \dots, [N^k, F^k])$. Both neighbour chunks and their continuations provide meaningful improvements, as illustrated in our ablation study (Appendix D). We use a length $64$ for both $N^j$ and $F^j$, thus $\textsc{Ret}(C)$ has a shape of $k \times r$ with $r=128$. To avoid retrieving the chunk $C{u+1}$ in the retrieval set $\textsc{Ret}(C_u)$, which would break causality during training, we filter out neighbours originating from the same document as the training sequence $X$.

For a database of $T$ elements, we can query the approximate nearest neighbours in $\mathcal{O}(\log T)$ time. We use the SCaNN library ([18]) to achieve this. This means that we can query our $2$ trillion token database in $10, \textrm{ms}$ whilst evaluating or sampling from the model; this expense is amortized over a chunk length. Performing retrieval on-the-fly is too slow to keep up with the training calculations—we leverage the frozen aspect of the embedding operator $\textsc{Bert}$ to precompute all approximate nearest neighbours and save the results as part of the data. In Figure 9 in the Appendix, we show results where we only retrieve neighbours within Wikipedia. We find that neighbours tend to come from 2-3 links away from a given article whereas random articles are more than 5 links apart.

2.4 $\textsc{Retro}$ model architecture

Our model relies on an encoder-decoder transformer architecture, integrating the retrieved data through a cross-attention mechanism as introduced in [4]. First, the retrieved tokens $\textsc{Ret}(C)$ are fed into an encoder Transformer, which computes the encoded neighbours set $E$. Denoting the intermediate activations by $H$, our transformer decoder then interleaves $\textsc{Retro}$-blocks $\textsc{Retro} (H, E)$ and standard Transformer blocks $\textsc{LM}(H)$ (the hyperparameter $P \subseteq [1, L]$ determines at which layers we use a $\textsc{Retro}$-block). These blocks are built from three different residual operators with signature $\mathbb{R}^{n \times d} \to \mathbb{R}^{n \times d}$: a fully-connected layer $\textsc{Ffw}$, the standard sequence-level self-attention layer $\textsc{Attn}$, and a chunked cross-attention layer $\textsc{Cca}(\cdot, E)$ that incorporates information from the retrieval encoder:

$ \textsc{Retro} \left(H, E\right) \triangleq \textsc{Ffw} \left(\textsc{Cca}\left(\textsc{Attn}\left(H\right), E\right) \right), \quad\text{and}\quad \textsc{Lm}(H) \triangleq \textsc{Ffw}(\textsc{Attn}(H))\tag{2} $

Since $\textsc{Ffw}$, $\textsc{Attn}$ and $\textsc{Cca}$ are all autoregressive operators whose output at position $i$ only depends on $(h_{j})_{j \leq i}$, any succession of $\textsc{Retro}$ and $\textsc{lm}$ layers, followed by a token classification head defines an autoregressive log-likelihood Equation 1. An overview of the model architecture is given in Algorithm 1 and in Figure 2. We next describe the retrieval encoder and the chunked cross-attention layer in more detail, and explain how to sample from $\textsc{Retro}$.

**Hyperparam:** P and P(enc), indices of layers with cross-attention in the decoder and encoder respectively
 **Hyperparam:** L and L(enc), number of decoder layers and number of encoder layers.
 **Input:** X ∈ V(n): sequence of tokens. (RET(C(u)))_{1 ≤ u ≤ l}: the retrieved neighbours
 **Output:** O ∈ R(n) × | V |}: the output logits
 def ENCODER(RET(C(u))_{1 ≤ u ≤ l}, H)
  (H(u))_{u ∈ [1, l]} ← SPLIT(H)
  **for** j ∈ [1,k], u ∈ [1, l]: // Encoder shared across neighbours and chunks
      E(u)(j) = EMB(enc)(RET(C(u))^j) // May be shared with the decoder EMB
    **for** p' ∈ [1, L(enc)]:
          E(u)(j) ← ATTN(enc)(E(u)(j)) // Bi-directional attention
      **if** p' ∈ P(enc):
        E(u)(j) ← CA(enc)(E(u)(j), H(u))
          E(u)(j) ← FFW(enc)(E(u)(j))
  **return** E
 H ← EMB(X)
**for** p ∈ [1, L]:
          H ← ATTN(H) // Causal attention
  **if** p = min(P):
    // The neighbour ENCODER is conditioned with the decoder activations of the last layer before the first cross-attention

               E = ENCODER(RET(C(u))_{1 ≤ u ≤ l}, H)
  **if** p ∈ P:
    H ← CCA(H, E)
          H ← FFW(H)
  O ← READ(H)

Encoding retrieval neighbours.

For each chunk $C_u$, the $k$ retrieval neighbours $\textsc{Ret}(C_u)$ are fed into a bi-directional transformer $\textsc{Encoder}$, yielding the outputs $E_u^j \triangleq \textsc{Encoder}(\textsc{Ret}(C_u)^j, H_u) \in \mathbb{R}^{r \times d'}$, where $j \in [1, k]$ indexes each neighbour. The retrieval encoder is a non-causal transformer. It is conditioned on $H_u$, the activations of chunk $C_u$, through cross-attention layers; this allows the representations of the retrieval encoder to be modulated by the retrieving chunk in a differentiable way. More precisely, the encoding of the $j^{\text{th}}$ neighbour of the $u^{\text{th}}$ chunk, $\textsc{Ret}(C_u)^j$, depends on the attended activation $H_u \triangleq {(h_{(u-1) m + i})}{i \in [1, m]}\in \mathbb{R}^{m \times d}$ of chunk $C_u$ at layer $\min(P)$. All neighbours for all chunks are encoded in parallel, yielding a full encoded set $E \triangleq (E_u^j){u \in [1, l], j \in [1, k]} \in \mathbb{R}^{l \times k \times r \times d'}$. We denote $E_u \in \mathbb{R}^{k \times r \times d'}$ as the encoded neighbours for chunk $u \in [1, l]$.

Chunked cross-attention.

To perform the $\textsc{Cca}$ operation, we first split a given intermediate activation $H \in \mathbb{R}^{n \times d}$ into $l{-}1$ attending chunks ${\left(H_u^+ \triangleq {(h_{u, m +i - 1})}{i \in [1, m]} \in \mathbb{R}^{m \times d} \right)}{u \in [1, l-1]}$, as depicted on the right of Figure 2. $H_u^+$ holds the intermediary embeddings of the last token in chunk $C_u$ and of the first $m-1$ tokens in $C_{u+1}$ [^2]. We compute the cross-attention between $H_u^+$ and $E_u$ —the encoded retrieval set obtained from chunk $C_u$. Attention is computed across time and across neighbours simultaneously, as we merge the neighbour and time dimensions of $E_u$ before applying cross-attention. Since there is a notion of alignment between data chunks and retrieval neighbours, we use relative positional encodings as described in Appendix B.1.2.

[^2]: The last token of chunk $C_u$ is the first to be able to access the retrieved content $E_u$ while maintaining autoregressivity in Equation 1. Hence, there is a one token overlap between chunk $C_u = { \left(x_{(u-1) m + i} \right) }{i \in [1, m]}$ and the corresponding attending chunk $C_u^+ \triangleq { \left (x{u , m + i -1} \right)}_{i \in [1, m]}$.

We concatenate the $l{-}1$ outputs of the per-chunk cross-attentions (each of shape $m \times d$) across time, and properly pad the result; we thus form the output activation $\textsc{Cca}(H, E) \in \mathbb{R}^{n \times d}$. Formally, for each chunk $C_u$ and for each token $i \in [1, m]$ we set

$ \textsc{Cca}(H, E){u, m + i - 1} \triangleq \textsc{Ca}(h{u , m + i - 1}, E_{u}),\tag{3} $

where $\textsc{Ca}$ is the cross-attention residual operator over time-concatenated encoded neighbours. We recall that this operator is defined in its simplest version by three parameter matrices $K \in \mathbb{R}^{d \times c}, Q \in \mathbb{R}^{d \times c}$ and $V \in \mathbb{R}^{d \times d}$. For all $h \in \mathbb{R}^d$ and $Y \in \mathbb{R}^{T \times d}$, we define

$ \textsc{Ca}(h, Y) \triangleq \textrm{softmax}(Y K Q^T h) Y V,\tag{4} $

where the softmax is performed on the second dimension and all products are matrix products. We use multi-head cross-attention, and add positional encodings to the softmax (see Appendix B.1.2).

The first $m-1$ tokens cannot attend to any neighbour of a previous chunk; at these positions, we define $\textsc{Cca}$ as the identity, setting ${\textsc{Cca}(H, E)}j \triangleq h_j$ for all tokens $j \in [1, m-1]$. Finally, the last token $h{l m}$ attends to the last retrieval set $E_l$ and we set $h_{l, m} \triangleq \textsc{Ca}(h_{l, m}, E_l)$ (not shown in Figure 2). Algorithm 1 contains a simplified implementation of $\textsc{Cca}$. Note that chunked cross-attention is autoregressive: the output of $\textsc{Cca}$ at position $i$ depends on the sequence from tokens from $0$ to $i$ that is input to $\textsc{Cca}$.

With $\textsc{Retro}$ models, even though each $\textsc{Cca}$ cross-attention attends only to the neighbours of the preceding chunk $\textsc{Ret}(C_{u-1})$, the dependencies over previous neighbours are propagated via the self-attention operations. The activations of the $i$ $^{\text{th}}$ token in the $u$ $^{\text{th}}$ chunk therefore potentially depend upon the set of all previous neighbours $\textsc{Ret}(C_{u'})_{u'<u}$, without incurring the quadratic cost of cross attending to that set.

Sampling.

When sampling, at the end of a chunk $C_u$, we use SCaNN to retrieve neighbours $\textsc{Ret}(C_u)$, based on the embedding $\textsc{Bert}(C_u)$. The encoded neighbours $E_u = \textsc{Encoder}(\textsc{Ret}(C_u))$ are then used to condition the generation of the next chunk $C_{u+1}$, which we do incrementally: overall the cost of sampling is thus quadratic in the size of the sampled sequence, as when sampling from regular Transformers; the added cost of retrieval is linear in the number of chunks $l$, and is negligible compared to the token sampling cost in practice.

2.5 Baseline Transformer architecture

We use a transformer ([4]) similar to the one described in ([5]), with some minimal changes: we replace LayerNorm with RMSNorm ([19]) and use relative position encodings ([20]). As baselines, we train retrieval-free transformers with 132M, 368M, 1.3B and 7.0B parameters (embedding matrices are excluded from parameter counts). The hyperparameters we used are detailed in Table 2. All retrieval models use the same size encoder for the retrieval data, with $d'=896$ and 2 layers, which roughly adds $19M$ parameters. The encoder uses relative positional encodings. The retrieval models contain one $\textsc{Retro}$-block every 3 blocks, starting from layer 6. For our smallest model, $\textsc{Cca}$ is applied in layers 6, 9 and 12 of the main pathway and also once for query conditioning in the encoder, which adds an additional $12M$ parameters. The relative number of extra parameters reduces as we increase the baseline model size. All models are implemented using JAX ([21]) and Haiku ([22]).

\begin{tabular}{r r r r c c c}
  Baseline parameters & $\textsc{Retro}$ & $d$ & $d_{\textrm{ffw}}$ & # heads & Head size & # layers \\ 
  \toprule
  132M & 172M (+30\%) & 896 & 3, 584 & 16 & 64 & 12 \\
  368M & 425M (+15\%) & 1, 536 & 6, 144 & 12 & 128 & 12 \\
  1, 309M & 1, 451M (+11\%) & 2, 048 & 8, 192 & 16 & 128 & 24 \\
  6, 982M & 7, 532M $ $ (+8\%) & 4, 096 & 16, 384 & 32 & 128 & 32\\
  \end{tabular}

2.6 Quantifying dataset leakage exploitation

$\textsc{Retro}$ models may arguably benefit more easily from evaluation dataset leakage, i.e. the fact that we evaluate on data that were also present in the training set. To better understand how retrieval improves language modelling performance, we therefore quantify evaluation likelihood as a function of the overlap between the evaluation and training datasets.

The following approach can be used with any language model, and depends only on the frozen retriever system presented in Section 2.3. We split the evaluation sequences ${(X_i)}_i$ into chunks of length $m \le 64$, and we see the training data as a set of chunks $\mathcal{C}$. For each evaluation chunk $C \in \mathcal{C}$, we retrieve the 10 closest neighbours (of length up to 128) in the training data. We then compute the longest token substring common to both the evaluation chunk and its neighbours. This gives a number $s \in [0, m]$. The value $r(C) = \frac{s}{m}$, ranging from $0$ (chunk never seen) to $1$ (chunk entirely seen), gives a reliable indication of how much overlap there is between the evaluation chunk and the training data. For a given model, we then obtain the log-likelihood $\ell(C)$ of each chunk $C$, and the number of bytes $N(C)$ it encodes. We then consider the filtered bits-per-bytes of the model:

$ \forall, \alpha \in [0, 1], \quad \mathcal{C}\alpha \triangleq { C \in \mathcal{C}, r(C) \leq \alpha }, \quad\textrm{bpb}(\alpha) \triangleq \frac{\sum{C \in \mathcal{C}\alpha} \ell(C)}{ \sum{C \in \mathcal{C}_\alpha} N(C) },\tag{5} $

which correspond to the bits-per-bytes on the set of chunks that overlap less than $\alpha:%$ with the training chunks. Note that the full evaluation bit-per-bytes performance is recovered by $\mathrm{bpb}(1)$. The function $\textrm{bpb}(\cdot)$ allows us to evaluate the impact of evaluation leakage over predictive performance: for low $\alpha$, $\textrm{bpb}(\alpha)$ gives an indication on how the model performs on chunks that are entirely new; the slope of $\textrm{bpb}(\cdot)$ shows how much the model exploits evaluation leakage.

3. Related Work

Section Summary: The section reviews prior approaches that augment language models with external retrieval, ranging from early cache-based methods and nearest-neighbor interpolation at the token level to more recent systems such as REALM, RAG, and FiD that retrieve passages or answers for question-answering tasks. It contrasts these with Retro, which uses a far larger retrieval corpus, frozen BERT embeddings over text chunks, and a distinctive chunked cross-attention mechanism for integration. The text also notes that training on a substantial portion of the internet raises important questions around privacy, safety, and fairness.

We first review existing work on using retrieval for language modelling, and compare $\textsc{Retro}$ to these works (see Table 3). As we train $\textsc{Retro}$ models on a large dataset containing a substantial section of the internet, our work raises potential privacy, safety, and fairness issues that we then review.

\begin{tabular}{ l c c c c } 

 & # Retrieval tokens & Granularity & Retriever training & Retrieval integration \\
\toprule

Continuous Cache & $\mathcal{O} \left(10^3 \right)$ & Token & Frozen ($\textsc{LSTM}$) & Add to probs \\ 

$k\textrm{NN-LM}$ & $\mathcal{O} \left(10^9 \right)$ & Token & Frozen (Transformer) & Add to probs \\

$\textsc{Spalm}$ & $\mathcal{O} \left(10^9 \right)$ & Token & Frozen (Transformer) & Gated logits \\

$\textsc{Dpr}$ & $\mathcal{O} \left(10^9 \right)$ & Prompt & Contrastive proxy & Extractive QA \\ 

$\textsc{Realm}$ & $\mathcal{O} \left(10^9 \right)$ & Prompt & End-to-End & Prepend to prompt \\ 

$\textsc{RAG}$ & $\mathcal{O} \left(10^9 \right)$ & Prompt & Fine-tuned $\textsc{Dpr}$ & Cross-attention \\ 

$\textsc{FiD}$ & $\mathcal{O} \left(10^9 \right)$ & Prompt & Frozen $\textsc{Dpr}$ & Cross-attention \\ 

$\textsc{Emdr}^2$ & $\mathcal{O} \left(10^9 \right)$ & Prompt & End-to-End (EM) & Cross-attention \\ 

\textbf{\textsc{Retro} (ours)} & $\mathbf{\mathcal{O} \left ({10^{12}} \right)}$ & \textbf{Chunk} & \textbf{Frozen ($\textsc{Bert}$)} & \textbf{Chunked cross-attention} \\ 

\end{tabular}

3.1 Retrieval for language modelling

[23] show that scaling the training data to trillions of tokens improves the machine translation performance of $n$-gram models. More recently, GPT-2 ([5]), GPT-3 ([6]), and Jurassic-1 ([24]) show that scaling up language models leads to massive improvements on many downstream tasks. At the same time, [25] demonstrate that large-scale language models can perfectly memorise parts of their training data, suggesting that enhancing models with retrieval may lead to further improvements. However, significant leakage between train and test datasets ([14, 26]) makes comparing and evaluating large models trained on large datasets difficult, especially once retrieval capabilities over the training dataset are added.

Historically, information retrieval for text relies on inverted index matching such as TF-IDF and BM25 ([27]). Foundational work use latent topic modelling approaches like LDA ([28]) to identify relevant neighbours ([29]). Work in machine translation such as [30] and [31] retrieve translation pairs based on edit distance between source sentences and guide the translation output using the closest retrieved target sentences. The retrieval database may also be structured — for example, [32] use a symbolic knowledge graph to improve an RNN language model.

With the success of deep learning, retrieving systems have partly switched to dense learned representations based on a neural network's activations. Continuous cache ([33]) adds probability mass to tokens for which previous activations resemble the current activation vector, extending the model's context to the local history. $k\textrm{NN-LM}$ ([8]) applies this idea to transformers and extends the retrieval database to English Wikipedia, resulting in substantial improvements on Wikitext103 evaluation. Continuous cache and $k\textrm{NN-LM}$ do not modify the underlying neural-network models, but interpolate at inference between the language model's output and distributions computed from retrieved tokens. These methods can therefore be plugged into any model without additional training, although this limits the model's ability to reason about the retrieved text. $\textsc{Spalm}$ ([9]) addresses this limitation by adding an extra gating network to post-process the retrieved data; yet most of the network is unaffected by the retrieval during inference.

The retrieval representations may be trained directly instead of relying on a pre-trained model—retriever systems have been developed for this purpose, primarily on open-domain question answering. For example, $\textsc{Dpr}$ ([34]) trains two $\textsc{Bert}$ models (for queries and keys respectively) using a contrastive loss to align the representations of a question and of its answers. [35] use an inverse cloze task to find semantic representations of passages for retrieval. These works differs from continuous cache and $k\textrm{NN-LM}$ in that they embeds passages (or chunks) of text together, as opposed to each token individually. The retriever network is trained in isolation of the downstream task that uses the retrieval data. This potential issue is specifically addressed by $\textsc{Realm}$ ([10]), which trains the retrieval system end-to-end to maximize the final training cross-entropy. This comes with the extra complexity of searching the database during training and periodically updating the embedding table, severely limiting the scale at which it can operate. $\textsc{RAG}$ ([11]) and $\textsc{FiD}$ ([36]) build upon $\textsc{Dpr}$ to set the state of the art on question answering benchmarks by training encoder-decoder transformer models. More recently, $\textsc{Emdr}^2$ ([37]) extends $\textsc{FiD}$ by using an expectation-maximization algorithm to train the retriever end-to-end and achieves state of the art results compared to similarly sized models.

In the open-domain dialogue setting, BlenderBot 2.0 ([38]) learns to issue textual internet queries, outperforming dense retrieval methods when evaluated on a task measuring how close model responses are to those of humans. This involves collecting a dataset of human dialogues with associated search queries, which limits the scalability of this approach. [39] introduce the Guided Transformer, a modified Transformer similar to $\textsc{Retro}$, for document retrieval and clarifying question selection. Although effective on question answering and other tasks with strong conditioning, none of these methods are designed to model arbitrary text sequences, in contrast with $\textsc{Retro}$.

$\textsc{Retro}$ shares components with $k\textrm{NN-LM}$ and $\textsc{Dpr}$ in that it uses frozen retrieval representations. $\textsc{Retro}$ models longer sequences than QA examples; this requires to reason at a sub-sequence level, and to retrieve different documents for the different chunks of a sequence. Similar to $\textsc{FiD}$, $\textsc{Retro}$ processes the retrieved neighbours separately in the encoder, and assemble them in the chunked cross-attention. This differs from e.g. $\textsc{Realm}$, that prepends retrieved documents to the prompt. Using chunks allows for repeated retrieval whilst generating a sequence as opposed to retrieving only once based on the prompt alone. Furthermore, retrieval is done during the whole pre-training process in $\textsc{Retro}$, and is not simply plugged-in to solve a certain downstream task. Finally, previous methods based on dense query vectors use small models and retrieval datasets with less than 3B tokens (English Wikipedia). Table 3 summarizes the difference of $\textsc{Retro}$ with existing approaches.

3.2 Privacy, safety and fairness

[40, 41] highlight several dangers of large language models. Those stem from their ability to memorise training data, their high training cost, the static nature of their training data ([42]), their tendency of amplifying inherent biases in the training data, and their ability to generate toxic language ([43]). In this section we inspect these dangers, focusing on how retrieval augmented language models may exacerbate or mitigate them.

Large language models can perfectly memorise parts of their training data ([25]). When coupled with large training datasets gathered from the web or other sources, this has clear privacy and safety implications. Retrieval models such as $\textsc{Retro}$ that have access to the entire training dataset during inference exacerbate these privacy issues by being able to directly copy training data. However, retrieval systems offer a path towards mitigating these concerns via obliteration of the retrievable data at inference time. In addition, differential privacy training ([44]) of retrieval models could guarantee that no private information is stored in the model weights, while individualisation on private data could be made by updating the retrieval database at inference time.

Due to their high training cost, re-training large language model regularly to incorporate new data, languages, and norms is prohibitively expensive. To keep retrieval models up-to-date, it may be sufficient to update the retrieval database, which is orders of magnitude cheaper than re-training a model from scratch. In addition to the benefits of updating models in terms of fairness and bias, simply training large language models has a significant energy cost ([45, 46]). Retrieval mechanisms offer a path to reducing the compute requirements needed to train and update language models that reach a certain performance.

Large language models are prone to generating toxic outputs, as shown in [43]. [40, 47] advocate for the importance of better training data curation and documentation. Additionally, if portions of the training data are found to be eliciting biased or toxic outputs after training, retrieval allows for some correction, as the offending retrieval data can be retroactively filtered. However, it is also the case that without careful analysis and intervention, retrieval models may exacerbate biases that are present in the training data. Retrieval models can also add a further source of bias through the selection mechanism for retrieval documents. Further work in this area is required to better understand how retrieval affects the bias and toxicity of the model outputs.

Finally, samples from large models are difficult to interpret, making mitigating these issues all the more challenging ([48, 49]). Retrieval provides more insights in to the outputs of a model, as one can directly visualise or modify the neighbours that are being used. The examples in Table 6, Table 7, Table 20 and Table 21 illustrate how retrieval makes language models more factual and interpretable by providing more transparent outputs.

4. Results

Section Summary: The results show that Retro models outperform standard Transformer baselines on a range of language modeling benchmarks such as C4, Wikitext103, and The Pile, with gains that hold steady or grow as models scale from 150 million to 7 billion parameters. Larger retrieval databases and more retrieved neighbors further boost performance, enabling the 7.5 billion parameter Retro model to beat much larger systems like Jurassic-1 and Gopher on most tests. Similar advantages appear in question-answering evaluations and remain after checks for data leakage.

We first report results on language modelling benchmarks. Second, we show how to $\textsc{Retro}$ fit pre-trained Transformer language models into retrieval models with few additional FLOPs. Next, we report $\textsc{Retro}$ results on question answering. Finally, we report evaluation metrics with leakage filtering, to better understand the source of the gains with retrieval.

4.1 Language modelling

Datasets.

We evaluate our models on C4 ([50]), Wikitext103 ([12]), Curation Corpus ([51]), Lambada ([52]) and the Pile ([13]). We also evaluate on a set of manually selected Wikipedia articles that were added or heavily edited in September 2021, months after our pre-training and retrieval dataset was collected (details are given in Appendix A.2). We construct the dataset with articles from the "future" and manually remove new articles that strongly overlap documents in our training data. This guarantees that the evaluation documents are not leaked in our training data.

For C4, Wikitext103, the Pile, and our Wikipedia dataset we evaluate the language modelling performance on entire documents and measure the bits-per-byte (bpb). We favour bits-per-byte over loss as it is tokenizer agnostic. We evaluate with a sequence length of 2048 tokens but use a stride of 1024 within documents to mitigate boundary effects. On Curation Corpus we concatenate the article, the "TL;DR:" string, and the summary, but only evaluate the bpb on the summary. For Lambada we evaluate the accuracy on the last word, using greedy generation.

Model scaling.

In Figure 1(left) and Figure 3 we show the language modelling performance as we scale models from 150 million to 7 billion (non-embedding) parameters. We see that on all datasets, $\textsc{Retro}$ outperforms the baseline at all model sizes. Furthermore, we observe that improvements do not diminish as we scale the models. The performance is dataset dependent, with the largest gains on Wikitext103 and C4. Wikipedia articles and other web pages are similar to Wikitext103 documents, even if not exact copies (Section 4.4), we thus obtain dramatic improvements on Wikitext103 as our retrieval model is able to directly exploit these overlaps. The smallest gains are for Curation Corpus, where $\textsc{Retro}$ only slightly outperforms the baseline. This is expected as Curation Corpus summaries are designed to only contain information from the source article and are not included in our retrieval database. On our "future" Wikipedia September 2021 dataset, we also observe consistent gains for all model sizes.

**Figure 3:** **Scaling with respect to model size.** (a) LAMBADA top-1 accuracy. (b) Evaluation loss on curation corpus. (c) Perplexity on Wikitext103 valid. (d) Bits-per-byte on selected Wikipedia articles from September 2021.

Data scaling.

Figure 1 (middle) shows how scaling the retrieval database at evaluation improves the language modelling performance. We observe dramatic gains as the retrieval data is increased from Wikipedia (4 billion tokens) to all of Massive text (1.7T tokens). Figure 1(right) shows how performance scales as we increase the number of retrieved chunks. Despite being only trained with 2 neighbours, we see consistent improvements for all models when the number of neighbours is increased from 1 to 10. Furthermore, we observe that larger models are able to better utilise more neighbours: the 172M model improves with up to 10 neighbours, whereas the 7B model improves with up to 40 neighbours.

The Pile.

We evaluate our 7B models on the Pile test sets[^3] and compare against the 178B parameter Jurrasic-1 ([24]) model and the 280B parameter Gopher ([16]) model. We do not compare against GPT-3 as it is outperformed by Jurassic-1 and Gopher on almost all subsets. Figure 4 shows the relative improvements in bits-per-byte over our 7B transformer baseline for our 7.5B $\textsc{Retro}$ model, Jurassic-1 and Gopher. Jurassic-1 outperforms the baseline on all datasets except for books, likely due to the inclusion of books in our training data. Gopher and $\textsc{Retro}$ outperform the baseline on all test sets. Overall, $\textsc{Retro}$ 7.5B outperforms Jurassic-1 and Gopher on a majority of the test sets. On the dm_mathematics and ubuntu_irc subsets, our $\textsc{Retro}$ model does not outperform our 7B baseline and underperforms Jurassic-1. We hypothesise that the retrieved neighbours on these datasets are not helpful, due to a combination of what is in our retrieval dataset and the efficacy of the nearest-neighbour search.

[^3]: Due to legal and ethical concerns relating to their use, we exclude the Enron Emails and the Youtube Subtitles datasets.

**Figure 4:** **The Pile: Comparison of our 7B baseline against Jurassic-1, Gopher, and $\textsc{Retro}$.** We observe that the retrieval model outperforms the baseline on all test sets and outperforms Jurassic-1 on a majority of them, despite being over an order of magnitude smaller.

\begin{tabular}{l l r r r r r}
  Model & Retrieval Set & #Database tokens & #Database keys & Valid & Test \\
  \toprule
  Adaptive Inputs ([53]) & - & - & - & 17.96 & 18.65 \\
  $\textsc{Spalm}$ ([9]) & Wikipedia & 3B & 3B & 17.20 & 17.60\\
  $k\textrm{NN-LM}$ ([8]) & Wikipedia & 3B & 3B & 16.06 & 16.12 \\
  Megatron ([54]) & - & - & - & - & 10.81 \\ 

  \midrule
  Baseline transformer (ours) & - & - & - & 21.53 & 22.96 \\
  $k\textrm{NN-LM}$ (ours) & Wikipedia & 4B & 4B & 18.52 & 19.54\\
  $\textsc{Retro}$ & Wikipedia & 4B & 0.06B & 18.46 & 18.97\\
  $\textsc{Retro}$ & C4 & 174B & 2.9B & 12.87 & 10.23 \\
  $\textsc{Retro}$ & MassiveText (1\%) & 18B & 0.8B & 18.92 & 20.33\\
  $\textsc{Retro}$ & MassiveText (10\%) & 179B & 4B & 13.54 & 14.95\\
  $\textsc{Retro}$ & MassiveText (100\%) & 1792B& 28B & \textbf{3.21} & \textbf{3.92}\\

  \end{tabular}

Wikitext103.

To validate our approach in a controlled setting, we compare our method with $k\textrm{NN-LM}$ ([8]) on the Wikitext103 dataset in Table 4. We train a baseline transformer on the training set of Wikitext103. This transformer has 24 layers, 1024 hidden units, 16 heads and a key size of 64, as in [53]. Our baseline does not have adaptive input, and our tokenizer has an open vocabulary, unlike [53], which makes our baseline perplexities a bit higher. The full experiment details and hyperparameters are given in Appendix C.2 and Table 11.

We re-implement $k\textrm{NN-LM}$ with our tokenizer and baseline transformer to produce embeddings of size 1024 for every token in Wikitext103. $k\textrm{NN-LM}$ has probabilities $p_{k\textrm{NN-LM}}= \lambda p_{k\textrm{NN}} + (1-\lambda) p_{\textsc{Lm}} $ with $p_{k\textrm{NN}} \left (n_k \right) \propto \exp \left (-\alpha d_k \right)$. We tune $\lambda=0.118$ and $\alpha=0.00785$ on the validation set (Figure 7) and report performance for these hyperparameters on both the validation and test set.

We fine-tune our baseline transformer into a $\textsc{Retro}$ model (Figure 7), using the Wikitext103 training data and retrieving from Wikipedia with 2 neighbours. We only train the new weights, as explained in Section 4.2, and share the embedding weights between the encoder and the main pathway. This is necessary for Wikitext103 which is quite small, as training $\textsc{Retro}$ from scratch in this setting leads to over-fitting.

We evaluate the fine-tuned $\textsc{Retro}$ model with different retrieval sets. We use 10 neighbours at evaluation for both $\textsc{Retro}$ and $k\textrm{NN-LM}$. When retrieving from Wikipedia, we obtain results comparable to our $k\textrm{NN-LM}$ implementation. Furthermore, scaling the retrieval database to MassiveText yields dramatic improvements, though this is partly due to leakage (see Section 4.4). For reproducibility, we also include results when retrieving from C4, which are close to previous state-of-the-art and comparable to using 10 % of MassiveText.

It is worth noting that $k\textrm{NN-LM}$ requires 1024 floats for every token in the retrieval dataset, totalling 15 terabytes (Tb) for the 4 billion tokens in Wikipedia. $k\textrm{NN-LM}$ and other token-level retrieval approaches therefore don't scale to retrieval databases with trillions of tokens such as MassiveText. In comparison, $\textsc{Retro}$ only requires 215Gb to index our Wikipedia dataset, and 93Tb for MassiveText. Inspecting the number of retrieval database entries in Table 4 makes it clear why retrieving at the chunk level is necessary when scaling to datasets with trillions of tokens.

4.2 $\textsc{Retro}$-fitting baseline models

We extend baseline models into $\textsc{Retro}$ models by freezing the pre-trained weights and training only chunked cross-attention and neighbour encoder parameters (less than 10% of weights for the 7B model) in Figure 5. This offers an efficient alternative path to enhance transformers with retrieval, requiring only 6 million sequences (3% of the pre-training sequences that we used). Additionally, by only training the new weights we ensure that when evaluated without retrieval, the original model performance is exactly maintained. $\textsc{Retro}$ fitting models quickly surpasses the performance of baseline models and even achieves performance close to that of $\textsc{Retro}$ models trained from scratch. The experiment hyperparameters are given in Appendix C.3.

**Figure 5:** **$\textsc{Retro}$-fitting a baseline transformer.** Any transformer can be fine-tuned into a retrieval-enhanced transformer by randomly initializing and training only the chunked cross-attention and retrieval encoder weights. Fine-tuning in this way quickly recovers and surpasses the non-retrieval performance, and almost achieves the same performance as training a retrieval model from scratch (shown by the arrow on the right hand side of each plot). We find good performance $\textsc{Retro}$-fitting our models training on only 3% the number of tokens seen during pre-training.

4.3 Question answering

We fine-tune our retrieval models on the Natural Questions ([55]) dataset to demonstrate that our retrieval pathway can be used to inject information from arbitrary data sources. We use the version^4 provided by [36] which is augmented with the retrieved passages from $\textsc{Dpr}$ ([34]). We fine-tune all the weights of our 7.5B pre-trained $\textsc{Retro}$ model for 25, 000 steps using the top 20 retrieved passages. We format the data as "question: question \n answer: answer" and left pad the data such that "answer:" coincides with the end of the first chunk of 64 tokens and thus aligns with the first retrieving chunk. The model has access to the question via the previous tokens in the sequence as well as the top 20 DPR Wikipedia passages and their titles via the chunked cross-attention mechanism. The exact match scores are shown in Table 5 and the full fine-tuning details are given in Appendix C.4. Our method is competitive with previous approaches such as $\textsc{Realm}$, $\textsc{RAG}$ and $\textsc{Dpr}$, but underperforms the more recent $\textsc{FiD}$. In contrast with this work, we find that increasing the number of neighbours past 20 does not improve $\textsc{Retro}$ performance on this task. We hypothesise that the encoder-decoder structure of T5—the base model in $\textsc{FiD}$ — and the T5 pre-training objective leads to a model that relies more on the encoder output than $\textsc{Retro}$, which is important in the QA setting. To compete with T5-finetuned models, future work should consider ways of forcing $\textsc{Retro}$ to rely further on the retrieval encoder output when producing tokens.

\begin{tabular}{l c}
  Model & Test Accuracy\\  
  \toprule
  $\textsc{Realm}$ ([10]) & 40.4 \\
  $\textsc{Dpr}$ ([34]) & 41.5 \\
  $\textsc{RAG}$ ([11]) & 44.5 \\
  $\textsc{Emdr}^2$ ([37]) & 52.5 \\
  $\textsc{FiD}$ ([36]) & 51.4 \\
  $\textsc{FiD}$ + Distill. ([56]) & \textbf{54.7} \\
  \midrule
  Baseline 7B (closed book) & 30.4 \\
  $\textsc{Retro}$ 7.5B (DPR retrieval) & 45.5 \\
  \end{tabular}

4.4 Relating retrieval performance to dataset leakage.

We report the filtered eval losses as detailed in Section 2.6 on C4, Curation Corpus and Wikitext103 in Figure 6. On C4 and Wikitext103, for which there is leakage into the training set, the slope is negative for both baseline models and $\textsc{Retro}$ models. $\textsc{Retro}$ models exploit leakage more strongly than baseline models, as indicated by the more negative slope. This is due to its explicit ability to copy-paste existing training chunks to predict leaked evaluation chunks (see a qualitative example of this model behavior on a Wikitext103 article in Table 19). On Curation Corpus, retrieval provides a constant offset, which is expected as there is by design no leakage between Curation Corpus and the training dataset.

On the other hand, $\textsc{Retro}$ outperforms baseline models at all leakage levels, down to $\alpha = 12.5%$. At this level, the loss is computed on chunks with less than $8$ contiguous tokens shared with the closest matching chunk in the training dataset—this is a reasonable level of overlap at which we consider that there is no local leakage. Retrieval thus improves predictions on both chunks that are syntactically similar to chunks in the training set, and on chunks that are syntactically different from all training chunks. This points toward a non trivial $\textsc{Retro}$ capacity of generalizing based on both model parameters and retrieval database. Similar results are found on the Pile dataset (see Figure 12, Appendix F.3).

4.5 Using $\textsc{Retro}$ for sampling

We show examples of samples obtained using the 7.5B $\textsc{Retro}$ model in Table 6, Table 7 and Appendix E. For each chunk (the first one being the prompt), we juxtapose sampled chunks $C_u$ with retrieved neighbours $\textsc{Ret}(C_{u})$. To give an indication of local overlap, we colour each sampled token in chunk $C_u$ based on the length of the longest common prefix (LCP) found in the retrieved chunks $\textsc{Ret}(C_{u-1})$. Similarly, we colour the retrieved chunks based on the LCP in the sampled chunk. For the sample in Table 6, for which we chose the prompt, we observe that the retrieved chunks influence the sample as there are overlaps between the sampled tokens and neighbour tokens. Overall, retrieval reduces hallucinations (in line with the findings of [57]) and makes the model more knowledgeable, when comparing with samples produced with retrieval disabled. In the sample in Table 7, the model recognises that the prompt is the beginning of the first scene of Hamlet and leverages retrieval data to continue it with only a few mistakes. We provide further examples in Appendix E, including examples from the evaluation sets, as well as the detailed procedure used for colouring the tables.

**Figure 6:** **Performance vs. longest common retrieval substring.** Evaluation loss as a function of allowed longest common substring between evaluation data chunks and their nearest neighbours. Retrieval still helps when considering chunks with no more than 8 contiguous tokens overlapping with training dataset chunks.

5. Conclusion

Section Summary: Retrieval-Enhanced Transformers, or Retro, is a new method that lets language models draw on enormous external databases containing trillions of tokens, effectively multiplying the data they can use during training by about ten times. The approach delivers gains that hold up even for models with billions of parameters, often matching the performance of much larger non-retrieval systems on benchmarks like Wikitext103 while remaining competitive on tasks such as question answering. It is also flexible—existing models can be quickly adapted to use retrieval, and the improvements are not mainly the result of test-data leakage—showing that retrieval offers a practical, efficient alternative to simply making models bigger.

We present Retrieval-Enhanced Transformers ($\textsc{Retro}$), a method for modelling arbitrary text sequences whilst retrieving from databases with trillions of tokens—scaling the data available to models by an order of magnitude compared to what is typically consumed during training. $\textsc{Retro}$ models gains do not diminish for models with up to at least 7B parameters, and correspond to non-retrieval models with 10 $\times$ more parameters on certain datasets. On Wikitext103 and the Pile, $\textsc{Retro}$ outperforms previous models trained on large scale datasets. We also show that $\textsc{Retro}$ is competitive on retrieval-intensive downstream tasks such as question answering.

$\textsc{Retro}$ models are flexible and can be used without retrieval at evaluation and still achieve comparable performance to baseline models. Conversely, baseline models can be rapidly fine-tuned into $\textsc{Retro}$ models to obtain nearly the same performance as if trained from scratch. Careful analysis shows that only a modest fraction of the gains obtained by $\textsc{Retro}$ are due to test set leakage. In general, we caution for such leakage in large-scale language datasets and suggest further work in better understanding the role of test set leakage in the performance of large-scale language models.

Overall, our work demonstrates at an unprecedented scale that semi-parametric approaches can provide an orthogonal, more efficient approach than raw parameter scaling as we seek to build more powerful language models.

Acknowledgements

Section Summary: The authors express their gratitude to a wide range of colleagues and collaborators for providing valuable help, advice, and feedback during the project. They specifically acknowledge contributions from individuals such as Nikolai Grigorev, Marc'aurelio Ranzato, and many others listed by name. Additional thanks are given to Zonglin Li, David Simcha, and the developers of the ScaNN tool for their targeted support.

We would like to thank Nikolai Grigorev, Marc'aurelio Ranzato, Cyprien de Masson d'Autume, Po-Sen Huang, Johannes Welbl, Lisa Anne Hendricks, Ethan Perez, Jeff Stanway, Eric Noland, Gregory Wayne, John Jumper, Julian Schrittwieser, Lorrayne Bennett, Devang Agrawal, Dani Yogatama, Susannah Young, Nando de Freitas, Demis Hassabis, and Koray Kavukcuoglu for their help, advice and reviews. Additionally, we would like to thank Zonglin Li, David Simcha, and the ScaNN developers for their help.

::: {caption="Table 6: Sample - Beavers are interesting animals. The Retro[Off] sample quickly diverges to other animals while the Retro[On] sample tends to stay focused on the beaver topic due to neighbour conditioning."}

:::

::: {caption="Table 7: Sample - Hamlet, Act 1, Scene 1. The Retro[Off] sample has correct syntax but is hallucinated, and ends with repetition of one character (FRANCISCO Approach me not). The Retro[On] sample is the correct continuation of the original text, and is robust to formatting differences between our prompt and the retrieved data."}

:::

Appendix

Section Summary: The appendix first describes the MassiveText training dataset in detail, breaking down its sources such as web pages, books, news, Wikipedia, and code across multiple languages along with their sizes and sampling weights during training. It also explains the creation of a small evaluation set of recent Wikipedia articles from September 2021, including how the articles were filtered, parsed, and cleaned to minimize overlap with training data. The section then outlines technical aspects of the RETRO retrieval architecture, focusing on its fully feed-forward encoder-decoder structure, chunked cross-attention mechanisms, and relative positional encodings used to integrate retrieved neighbors.

A. Datasets

We provide a full description of MassiveText and of our extract of recent Wikipedia articles.

A.1 Full description of MassiveText

The full break down of MassiveText by source and languages is given in Table 8. For a full description and analysis of MassiveText, see [16].

\begin{tabular}{c c c r r r}
Source & Language & Token count (M) & Documents & Sampling weight\\  
\toprule
\multirow{ 10}{*}{Web} & En & 483, 002 & 604, 938, 816 & 0.314 \\ 
 & Ru & 103, 954 & 93, 004, 882 & 0.033 \\ 
 & Es & 95, 762 & 126, 893, 286 & 0.033 \\ 
 & Zh & 95, 152 & 121, 813, 451 & 0.033 \\ 
 & Fr & 59, 450 & 76, 612, 205 & 0.033 \\ 
 & De & 57, 546 & 77, 242, 640 & 0.033 \\ 
 & Pt & 44, 561 & 62, 524, 362 & 0.033 \\ 
 & It & 35, 255 & 42, 565, 093 & 0.033 \\ 
 & Sw & 2, 246 & 1, 971, 234 & 0.0044 \\ 
 & Ur & 631 & 455, 429 & 0.0011 \\ 
\midrule
Books & En & 3, 423, 740 & 20, 472, 632 & 0.25 \\
\midrule
News & En & 236, 918 & 397, 852, 713 & 0.1 \\
\midrule
\multirow{10}{*}{Wikipedia} & En & 3, 977 & 6, 267, 214 & 0.0285 \\ 
 & De & 2, 155 & 3, 307, 818 & 0.003 \\ 
 & Fr & 1, 783 & 2, 310, 040 & 0.003 \\ 
 & Ru & 1, 411 & 2, 767, 039 & 0.003 \\ 
 & Es & 1, 270 & 2, 885, 013 & 0.003 \\ 
 & It & 1, 071 & 2, 014, 291 & 0.003 \\ 
 & Zh & 927 & 1, 654, 772 & 0.003 \\ 
 & Pt & 614 & 1, 423, 335 & 0.003 \\ 
 & Ur & 61 & 344, 811 & 0.0001 \\ 
 & Sw & 15 & 58, 090 & 0.0004 \\ 
\midrule
Github & - &374, 952 & 142, 881, 832 & 0.05 \\
\midrule
Total & - & 5, 026, 463 & 1, 792, 260, 998 & 1
\end{tabular}

A.2 Wikipedia September 2021

: Table 9: Full set of articles included in our Wikipedia Sept. 2021 evaluation dataset.

Megan Rohrer Aakashavaani
Emma Raducanu Junior Eurovision Song Contest 2021
Ambra Sabatini Pavilion Bukit Jalil
WhyDonate Blake Desjarlais
The Juggernaut (company) 2021 All-Ireland Senior Football Championship Final
Angela Diaz Drift-barrier hypothesis
2020 Summer Paralympics Venomics
2021 Afghan protests Great Circle (novel)
Rexh Xhakli Hurricane Ida
Julia Laskin 2021 Montenegrin episcopal enthronement protests
Cuijk At War With the Silverfish
Ghoubet Wind Power Station

We create an evaluation dataset consisting of 23 Wikipedia articles that were added or heavily edited in September 2021, after we collected our training dataset. In addition, we filter out articles that rely too heavily on templated content, using the method detailed in Section 2.6 to identify articles with chunks that have a high overlap with their neighbours. Figure 10 show that little overlap remains between our test dataset and the retrieved neighbours from the training dataset. The full list of included articles is given in Table 9.

We first parse articles using mwparserfromhell^5. We then remove sections with the following titles: "references", "external links", "sources", "further reading", "see also", "citations", and "note". In the remaining sections, we remove Wikilinks and remove the following templates: "reflist", "notelist", "notelist-ua", "notelist-lr", "notelist-ur", and "notelist-lg". We also exclude objects with the "ref" or "table" tag and clean the remaining text with the strip_code function. Finally, we concatenate the title and all the sections and use nn to delimitate them.

B. Details on the retrieval architecture

We give details on the $\textsc{Retro}$ architecture, and on the fine-tuning procedure we use for $\textsc{Retro}$ fitting existing language models.

B.1 $\textsc{Retro}$ architecture and implementation

B.1.1 Feed-forward architecture

As mentioned in the main text, the overall encoder-decoder architecture is fully feed-forward. We start with a sequence $X \in \mathbb{V}^n = (C_u){1 \leq u \leq l}$, and its pre-computed neighbours ${(\textsc{Ret}(C_u))}{1 \leq u \leq l}$ and returns logits in $\mathbb{R}^{n \times | \mathbb{V} |}$. Along with $\textsc{Attn}$, $\textsc{Ffw}$, $\textsc{Cca}$ and $\textsc{Ca}$ operators introduced in the main text, we define the decoder embedding layer $\textsc{Emb}: \mathbb{V}^n \to \mathbb{R}^{n \times d}$, the $\textsc{Split}$ operator that extracts chunked intermediary embeddings $\textsc{Split}(H) \triangleq (H_u){1 \leq u \leq l} \in \mathbb{R}^{l \times m \times d}$ and the read-out layer $\textsc{Read}: \mathbb{R}^{n \times d} \to \mathbb{R}^{n \times | \mathbb{V} |}$. We then describe the forward pass in Algorithm 1. In addition to the usual Transformer ones, $\textsc{Retro}$ architecture hyperparameters involves the layer indices $P{\textrm{enc}}$ and $P$, at which the encoder and the decoder perform cross-attention.

B.1.2 Relative positional encoding in the chunked cross-attention layer

The $\textsc{Ca}$ operator uses relative positional logits, that are computed from a specific relative distance separating data tokens from retrieval tokens. Indeed, we expect any retrieval neighbour $\textsc{Ret}(C_u)^j$ and the chunk $C_u$ to be relatively well aligned, and assume that they start at the same position. Therefore, when computing $\textsc{Ca}(H_u^+, E_u)$, we set the distance between the data token $i \in [1, l]$ of chunk $C_u^+$ and the retrieval token $i' \in [1, 2 l]$ of $\textsc{Ret}(C_u)^j$ to be

$ d(i, i') \triangleq i - i' + l - 1.\tag{6} $

When computing the encoder cross-attentions $\textsc{Ca}(\textsc{Ret}(C_u)^j, H_u)$, we set the distance between the retrieval token $i' \in [1, 2 l]$ and the data token $i \in [1, l]$ to be

$ d_{\textrm{enc}}(i', i) \triangleq i' - i.\tag{7} $

Positional logits are obtained as a linear transform of a cosine vector computed from $(d(i, i'))_{i, i'}$, and are added to content logits, as in a regular self-attention block.

B.1.3 Chunked cross-attention implementation

Our implementation of the $\textsc{Cca}$ operator, shown in Algorithm 1, is based on a vectorized application of a cross-attention layer. For simplicity, we omit the multi-head attention logic and use the simplest Q, K, V attention. We omit relative positional logits computation, described above.


n = 128  # Sequence length
m = 16  # Chunk length
r = 32  # Retrieval length
k = 4  # Number of neighbours
d = 16  # Embedding size
l = n // m  # Number of chunks

## Parameters

> **Section Summary**: The provided code and tables detail model dimensions, cross-attention placement, and training hyperparameters for several Retro language models ranging from 247M to 7.5B parameters. Encoder and decoder embeddings are kept separate by default, with cross-attention layers inserted every third layer starting at layer six; all weights except the freshly initialized retrieval encoder and cross-attention parameters are frozen during fine-tuning. Pre-training runs for a fixed 419 billion tokens using AdamW, cosine learning-rate schedules with linear warm-up, and batch sizes of 256 or 1024, while Wikitext-103 runs use a smaller 18-layer transformer with its own learning-rate schedule.

Q = jnp.zeros((d, d))
K = jnp.zeros((d, d))
V = jnp.zeros((d, d))

def relative_positional_encodings(attending_length, attended_length):
  # Classical relative positional encodings
  ...

def cross_attention(chunk, neighbour):
  m, d = chunk.shape
  r, d = neighbour.shape
  queries = chunk @ Q
  keys = neighbour @ K
  logits = queries @ keys.T
  values = neighbour @ V
  return logits, values

def multi_neighbour_cross_attention(chunk, neighbours):
  m, d = chunk.shape
  k, r, d = neighbours.shape

  logits, values = jnp.vectorize(cross_attention, 
                                 signature='(m,d),(r,d)->(m,r),(r,d)')(
                                     chunk, neighbours)
  assert logits.shape == (k, m, r)
  assert values.shape == (k, r, d)
  logits += relative_positional_encodings(m, r)[None, :, :]
  logits = jnp.moveaxis(logits, 0, -1).reshape((m, r * k))
  values = jnp.moveaxis(values, 0, 1).reshape((r * k, d))
  return jax.nn.softmax(logits) @ values

def multi_chunk_cross_attention(observation, neighbours):
  attending_chunks = jnp.pad(observation[m-1:], 
                             ((0, m - 1), (0, 0)),
                             mode='constant').reshape(l, m, d)
  chunked_output = jnp.vectorize(multi_neighbour_cross_attention,
                                 signature='(m,d),(k,r,d)->(m,d)')(
                                     attending_chunks, neighbours)
  assert chunked_output.shape == (l, m, d)
  output = jnp.pad(chunked_output.reshape(n, d), 
                   ((m - 1, 0), (0, 0)),
                   mode='constant')[:n]
  return output

observation = jnp.zeros((n, d))  # Input
neighbours = jnp.zeros((l, k, r, d))

h = multi_chunk_cross_attention(observation, neighbours)

assert h.shape == (n, d) # Output
B.1.4 Optional sharing of embedding matrices

We use disjoint embeddings for the encoder and decoder by default, which allows us to use a different dimensionality for the encoder (typically kept at $d_{\textsc{Enc}} = 896)$ and for the decoder (that we scale up to $d=8192$). It is possible to share the embeddings, with little difference in training, as we show in the ablation section.

B.2 Baseline to $\textsc{Retro}$ model fine-tuning

As shown in Figure 5, we found that we were able to take a pre-trained baseline transformer and add $\textsc{Retro}$ through fine-tuning. In all cases, we froze all weights from pre-training and freshly initialised the retrieval encoder and cross-attention weights. In all cases, the cross-attention is added every third layer starting at layer six. The learning rate for the three smaller models was set to 2 x 10^-4 and half that for the larger model. We experimented with allowing the entire model to resume training during fine-tuning but consistently found that the best approach was to freeze the pre-trained model. This kept the retrieval-off performance frozen whereas when all weights were tuned the retrieval off performance would degrade.

C. Training details and hyperparameters

We provide the hyperparameters used in the various experiments of Section 4.

C.1 Language model pre-training

In Table 10, we show the hyperparameters of the different models we train. In all cases, we train for 419, 430, 400, 000 training tokens. The three smaller models are trained with a batch size of 256 and the largest model is trained with a batch size of 1024. The minimum learning rate is set to 0.1 times the maximum learning rate, which is shown in Table 10. The learning rate is decayed using a cosine cycle length that matches the total number of training tokens. All models are trained using AdamW ([58]) with a weight decay parameter of 0.1. The learning rate linearly increases from $10^{-7}$ to the maximum learning rate over the first 750 steps of training.

\begin{tabular}{c c c c c c c c c c}
  Baseline & $d_{model}$ & $d_{ffw}$ & $#$ heads & Head size & $#$ layers & $P$ & $P_{\textsc{Enc}}$ & Max LR \\ 
  \toprule
  247M & 896 & 3584 & 16 & 64 & 12 & $[6, 9, 12]$ & $[1]$ & 2 $\times10^{-4}$ \\
  564M & 1536 & 6144 & 12 & 128 & 12 & $[6, 9, 12]$ & $[1]$ & 2 $\times10^{-4}$ \\
  1, 574M & 2048 & 8192 & 16 & 128 & 24 & $[9, 12, \dots, 24]$ & $[1]$ & 2 $\times10^{-4}$ \\
  7, 505M & 4096 & 16384 & 32 & 128 & 32 & $[9, 12, \dots, 32]$ & $[1]$ & 1 $\times10^{-4}$ \\
  \end{tabular}

All models use ZeRO to shard the optimiser state ([59]). Additional infrastructure details can be found in [16].

\begin{tabular}{l l r} 
  \toprule
  Model & Number of layers & 18 \\
  & $d$ & 1024 \\
  & $d_\textsc{Ffw}$ & 4096 \\ 
  & Key size & 64 \\
  & Value size & 64 \\
  & Number of heads & 16 \\ 

  \midrule
  Training data & Dataset & Wikitext103train \\
  & Sequence length & 3072 \\
  & Batch size & 128 \\
  & Tokenizer vocabulary size & 128, 000 \\ \midrule
  Optimisation & optimiser & Adam \\
  & Adam's $\beta_1 $ & 0.9 \\
  & Adam's $\beta_2$ & 0.95 \\
  & Adam's $\varepsilon$ & 1e-8 \\
  & Dropout rate & 0.25 \\ \hline
  Schedule & Learning rate start & 1e-7 \\
  & Learning rate max & 2.5e-4 \\
  & Learning rate min & 2e-5 \\
  & Warmup steps & 4, 000 \\
  & Cosine cycle steps & 100, 000 \\ \hline
  Evaluation & Overlapping proportion & 87.5 \% \\ 
  \bottomrule

  \end{tabular}

C.2 Wikitext103 comparison

We provide more details on our Wikitext103 results presented in Section 4.1 and Table 4. We train a baseline transformer on the Wikitext103 training set with the hyperparameters presented in Table 11. The learning rate ramps linearly from 1 x 10^-7 to 2.5 x 10^-4 in the first 4, 000 steps, then decays to 2 x 10^-5 at 100, 000 steps using a cosine schedule. The baseline checkpoint at step 35, 000 has the lowest perplexity on Wikitext103 valid, of $21.58$, for overlapping proportion of 75% (sliding window evaluation that only uses probabilities for tokens that have at least 75% of the sequence length of context, when available). We use this checkpoint for all our baseline and $k\textrm{NN-LM}$ numbers reported in Table 4, except that Table 4 reports for an overlapping proportion of 87.5 %, which slightly lowers the perplexity of our baseline to 21.53 on Wikitext103 valid.

We also use the 35, 000 step baseline checkpoint as initialization for a $\textsc{Retro}$ fit, which otherwise uses the same optimiser and schedule hyperparameters but only trains the new retrieval weights, as explained in Section 4.2. Our best $\textsc{Retro}$ fit checkpoint has a Wikitext103 valid perplexity $18.46$, when retrieving from Wikipedia. We use this $\textsc{Retro}$ checkpoint in Table 4 for all other retrieval sets. The evaluation curves for our baseline and $\textsc{Retro}$ fit is shown if Figure 7 (left). In this particular case, because Wikitext103 is quite small, training a $\textsc{Retro}$ model from scratch led to weaker results than the baseline, at least when retrieving from Wikipedia, as we couldn't find an effective way to mitigate the increased over-fitting due to the additional weights of $\textsc{Retro}$.

We also re-implement $k\textrm{NN-LM}$ using the same tokenizer and dataset that we use for our baseline and $\textsc{Retro}$ fitting experiments. $k\textrm{NN-LM}$ has probabilities $p_{k\textrm{NN-LM}}= \lambda p_{LM} + (1-\lambda) p_{kNN}$ with $p_{kNN}(n_k) \propto \exp(-\alpha d_k)$. To tune $\lambda$ and $\alpha$, we begin with $\alpha=0.0012$, which corresponds to the inverse of the standard deviation of the norm of the embeddings that we use as keys and queries for $k\textrm{NN-LM}$. We find the best $\lambda=0.118$. We then find the best $\alpha=0.00785$ for that value of $\lambda$. Figure 7 center and right respectively show the perplexity of $k\textrm{NN-LM}$ as a function of $\lambda$ and $\alpha$.

**Figure 7:** **Wikitext103valid perplexities.** *Left:* Baseline and $\textsc{Retro}$ fit (initialized from baseline's checkpoint at 35, 000 steps) perplexities as a function of training steps. *Center and right:* $k\textrm{NN-LM}$ perplexity as a function of $\lambda$ (for $\alpha = 0.0012$) and $\alpha$ (for $\lambda = 0.12$) respectively.

C.3 $\textsc{Retro}$ fitting baseline models experiments

In Table 12, we give the hyperparameters used for $\textsc{Retro}$ fitting the models on Massive Text.

::: caption="Table 12: Hyperparameters for the Retro fitting experiments"

:::

C.4 Question answering experiments

We fine-tune our 7.5B $\textsc{Retro}$ model for 25, 000 steps, using a batch size of 128, a learning rate cosine scheduled from $10^{-6}$ to $10^{-7}$, with a linear ramp of 750 steps. We use dropout in the decoder only, as it performs better than using dropout in both the encoder and the decoder. Each neighbour is formatted as title: title, source: source. We use the top 20 neighbours from $\textsc{Dpr}$ when training and evaluating.

\begin{tabular}{llr}
\toprule
  Ablation group & Ablation & C4 eval bpb \\
\midrule
  Model & \textsc{Retro} & 0.822 \\
  & No query conditioning & 0.829 \\
  & No CA positional encodings & 0.826 \\
  & Shared embeddings & 0.823 \\
  & 6-layer encoder & 0.821 \\
\midrule  
  Retrieval values & Neighbours N & 0.950 \\
  & Continuations F & 0.895 \\
  & No retrieval & 0.987 \\
\midrule
  Training neighbours & 1 training neighbours & 0.858 \\
  & 4 training neighbours & 0.847 \\
\midrule
 Cross attention position & CA top layer (1/12) & 0.827 \\
  & CA mid layer (6/12) & 0.823 \\
  & CA top layer (12/12) & 0.831 \\
  & CA all layers & 0.860 \\
  & CA every 3 from 1 & 0.823 \\
\bottomrule
\end{tabular}

D. Model ablations

We validate important design choices by evaluating what happens when we do not include them. We use the 247M parameter model for all experiments and we train on a compressed 157 billion token schedule for all ablation experiments. We describe results relative to the default settings presented in the main text and recalled here. We report C4 evaluation loss at the end of the training process, and also compares how the evaluation loss decrease versus the training time, measured relatively to the baseline training time. Results are reported in Figure 8 and Table 13.

**Figure 8:** **Computational efficiency for different variants.** We report the training curves plotting C4 evaluation bytes per bits against time, relative to the time taken to train the baseline $\textsc{Retro}$ model. Overall, our design choices are optimal in term of computational efficiency.

Using relative encodings in cross-attention.

Using relative encodings in cross-attention, as described in Appendix B.1.2, provides a pure improvement both in the number of steps to reach a given performance and computational efficiency.

Conditioning the encoder on the previous chunk.

Conditioning the encoder on the previous chunk's intermediate embeddings, as described in Appendix B.1.1, provides a pure improvement both in term of number of steps and computational efficiency.

Sharing embeddings.

Sharing embeddings across the encoder and the decoder does not affect performance. This motivates us using separate embeddings, as it allows to have a narrower encoder than decoder as we scale up the decoder size.

Attending neighbours and their continuation.

$\textsc{Retro}$ models are trained by attending, for a given chunk, to both the neighbours of the preceding chunk and their continuation in time. We measure how training and evaluating $\textsc{Retro}$ models on neighbours only and their continuation only affects performance. Overall, attending to neighbours only provides $22%$ of the performance improvement due to retrieval in $\textsc{Retro}$, while attending the future of the neighbours gives $56%$ of the performance. Attending to both neighbours and their continuation is the most efficient choice both in term of final performance and training efficiency.

Training a deeper encoder.

All models in the text use a relatively small $\textsc{Retro}$ encoder. We experimented with a $3\times$ deeper encoder. We found that this resulted in a tiny decrease in loss– 0.15% at the cost of a larger training time ($+20%$). Overall, using a shallow encoder is the best choice in term of training efficiency.

Training with multiple neighbours.

We measure the effect of training on a single retrieved neighbour, as well as training on 4 neighbours ($\textsc{Retro}$ uses 2 neighbours in training). Training on a single neighbour results in a large decrease in performance, while training on 4 neighbours does not give substantial performance improvement at the end of training, but induces a large computational overhead. Overall, we find that using 2 neighbours is the best choice in term of training efficiency. Furthermore, evaluation can be done with additional neighbours.

Frequency of cross-attention.

We measure how the frequency of cross-attention in the decoder affects performance. Overall, attending only once at the top or the bottom layer is a bad choice, while attending once on a mid-depth layer is relatively sound. We choose to have cross-attention every 3 layer as this provides a good trade-off between performance and run-time.

E. Qualitative experiments

We illustrate the usage of $\textsc{Retro}$ models by looking at the perplexity of evaluation samples and by producing samples autoregressively.

E.1 Inspecting neighbours and perplexities on evaluation data

To build an intuition of what kind of information is leveraged by $\textsc{Retro}$ models, we suggest to have a closer look at a few evaluation documents and the corresponding retrieved data in Table 16, Table 17, Table 18 and Table 19. In these tables, the 4 rows corresponds to the first 4 chunks of the documents. The left-most column shows the chunk $C_u$ from the document being evaluated, where each token is coloured by the negative cross entropy loss difference $L_{\textsc{Retro}\textsc{[Off]}} - L_{\textsc{Retro}}$, a positive value, coloured in yellow, indicates that $\textsc{Retro}$ performs better when it has access to neighbours data. The second columns also shows the evaluated chunk $C_u$ but where each token $i$ is coloured by the length of the longest common prefix (LCP) with the preceding neighbours, i.e. the largest integer $j$ such that the prefix $(x_{i-j -1}, \dotsc, x_{i})$ also appears in $\textsc{Ret}(C_{u-1})$. Conversely, columns three and four show the first two neighbours and their continuation, respectively $[N^1_u, F^1_u]$ and $[N^2_u, F^2_u]$ coloured by LCP with subsequent chunk $C_{u+1}$. LCP colouring helps to visually identify where the evaluated document overlaps the retrieved data. Note that the first chunk, $C_1$, in the second column is not coloured as it does not have any preceding neighbours to compute LCP with. Similarly, we do not show the neighbours of the fourth chunk, as these are not used to condition any of the first four chunks.

Our qualitative analysis exhibits two major behaviors.

Firstly, we observe that sometimes, specific facts in $C_u$ can be extracted from the preceding neighbours $\textsc{Ret}(C_{u-1})$ and that this can correspond to significant reduction in loss from the $\textsc{Retro}$ model for the corresponding tokens. Some examples of such behavior include the journal name Publishers Weekly in Table 16, the football team name Tyrone in Table 17 or the event dates 25 August to 6 September 2020 in Table 18. In these three examples, the evaluated data consists of recent Wikipedia articles written in September 2021, after we built our retrieval dataset (see Appendix A.2). Yet, relevant information to predict this new data was available in the pre-existing retrieval data and the $\textsc{Retro}$ model seems to be able to correctly leverage it.

On the other hand, we also observe that some of the evaluation data can partially leak in our training and retrieval data, despite the use of deduplication. $\textsc{Retro}$ can dramatically exploit such leakage. Table 19 illustrates this behavior, where the chunks $C_2$ and $C_3$ largely overlaps $\textsc{Ret}(C_{1})$ and $\textsc{Ret}(C_{2})$ respectively, up to small formatting differences, which leads to much lower $\textsc{Retro}$ loss for all the corresponding tokens. Figure 6 shows that it is possible to quantify how much of the $\textsc{Retro}$ loss reduction is due to each of these two behaviors, by filtering out evaluation chunks that overlaps with the retrieval set.

E.2 Inspecting samples

We can follow the same procedure as above on samples generated using $\textsc{Retro}$ models, in order to better understand where retrieval data had an influence on sampling. We show examples of samples obtained using the 7.5B $\textsc{Retro}$ model in Table 6, Table 7, Table 20 and Table 21.

E.3 Neighbour quantification

**Figure 9:** **Wikipedia link-distance between retrieved articles.** For each sequences, chunk combination we compute the link distance between the target and the top-5 neighbours using only Wikipedia. The rank shows the relative neighbour distance, where rank-1 is the first neighbour and rank 5 is the fifth. The different colours represent link distance. Because we do not retrieve from the same document, 1 is the smallest value. We find, on average, the distance between random articles with a path between them is over 5.0

To quantify a notion of distance between the source document and the retrieved chunks, we can ask the distance between source articles when retrieving only from Wikipedia. [60] provides a Wikipedia link dataset which, for each article, contains a list of neighbouring articles. Using this, we construct a directed graph and compute the distance from one page to another. In Figure 9 we compute the link-distance between training sequences and the retrieved neighbours. We find that retrieved documents tend to be from articles that are quite close to the article containing the target. Furthermore, we find that on average the distance increases with rank, suggesting that our neighbours are both useful and that the order is reasonable. This provides confidence for our larger-scale experiments where document distance is less well defined.

F. Complementary quantitative results

We report tables corresponding to quantitative figures of the main text, as well as further filtered language model results on the Pile.

F.1 Main text datasets

We report the performance of $\textsc{Retro}$ and baseline models, measured in bits-per-bytes on evaluation set, in Table 14.

::: {caption="Table 14: Full results for the main language modelling datasets. First three sets of rows correspond to Figure 1, last set of rows to Figure 3."}

:::

F.2 The Pile

In Figure 4, we compare $\textsc{Retro}$ against Jurassic-1 ([24]). The full bits-per-bytes results are reported in Table 15.

\begin{tabular}{l c c c c c }
  Subset & 7B Baseline (Ours) & GPT-3 & Jurassic-1 & Gopher & 7.5B $\textsc{Retro}$ \\ 
  \toprule
  arxiv & 0.742 & 0.838 & 0.680 & \textbf{0.641} & 0.714 \\
  books3 & 0.792 & 0.802 & 0.835 & 0.706 & \textbf{0.653} \\
  dm\_mathematics & 1.177 & 1.371 & \textbf{1.037} & 1.135 & 1.164 \\
  freelaw & 0.576 & 0.612 & 0.514 & 0.506 & \textbf{0.499} \\
  github & 0.420 & 0.645 & 0.358 & 0.367 & \textbf{0.199} \\
  gutenberg\_pg\_19 & 0.803 & 1.163 & 0.890 & 0.652 & \textbf{0.400} \\
  hackernews & 0.971 & 0.975 & 0.869 & 0.888 & \textbf{0.860} \\
  nih\_exporter & 0.650 & 0.612 & \textbf{0.590} & 0.590 & 0.635 \\
  opensubtitles & 0.974 & 0.932 & \textbf{0.879} & 0.894 & 0.930 \\
  philpapers & 0.760 & 0.723 & 0.742 & \textbf{0.682} & 0.699 \\
  pile\_cc & 0.771 & 0.698 & 0.669 & 0.688 & \textbf{0.626} \\
  pubmed\_abstracts & 0.639 & 0.625 & 0.587 & 0.578 & \textbf{0.542} \\
  pubmed\_central & 0.588 & 0.690 & 0.579 & 0.512 & \textbf{0.419} \\
  stackexchange & 0.714 & 0.773 & 0.655 & 0.638 & \textbf{0.624} \\
  ubuntu\_irc & 1.200 & 0.946 & \textbf{0.857} & 1.081 & 1.178 \\
  uspto\_backgrounds & 0.603 & 0.566 & \textbf{0.537} & 0.545 & 0.583 \\
  \end{tabular}

F.3 Filtered results

Distribution of leaked chunks in our main evaluation sets.

We evaluate leakage between the evaluation sets and the training set by measuring the proportion of evaluation chunks with a certain overlap $r(C)$. We show histograms in Figure 10. We can see that $C4$ has some slight overlaps between train and evaluation. Similarly, chunks of Wikitext103 appear in the training set despite having removed the actual Wikitext103 evaluation documents from the training set. On the other hand, our Wikipedia September 21 dataset shows almost no leakage (data being original documents that did not exist at training data creation), and neither does Curation Corpus.

Filtered results on the Pile.

We report chunk overlap distribution and filtered performance curves on the Pile in Figure 12 and Figure 11, respectively. The qualitative interpretation of the filtered curves is the same: $\textsc{Retro}$ models exploit leakage more, but the performance improvement they provide remains significant even on original chunks that haven't been observed in the training set.

**Figure 10:** **Distribution of the overlap between evaluation and train chunks** for C4, Curation Corpus, Wikitext103 and Wikipedia Sept. 2021.

**Figure 11:** **Filtered evaluation losses on the Pile**, with baseline Transformers and $\textsc{Retro}$.

**Figure 12:** **Distribution of the overlap between evaluation and train chunks** for the Pile evaluation sets.

::: {caption="Table 16: Great Circle (novel), from Wikipedia September 21. The article is about a recent novel and chunks $C_3$ and $C_4$ are specifically about its reception. The name Publishers Weekly of the journal that reviewed the novel appears both in the neighbours [$N^1_3, F^1_3], [N^2_3, F^2_3]$ of chunk $C_3$ and in the subsequent chunk $C_4$, where the loss for those tokens is significantly reduced by Retro."}

:::

::: {caption="Table 17: All-Ireland Senior Football Championship Final, from Wikipedia September 21. The name of the team Tyrone appears both in the second neighbours [$N^2_1, F^2_1]$ of chunk $C_1$ and in the subsequent chunk $C_2$, where the loss for those tokens is significantly reduced by Retro."}

:::

::: {caption="Table 18: 2020 Summer Paralympics, from Wikipedia September 21. The original dates of the event, 25 August to 6 September 2020, appears both in the neighbors [$N^1_1, F^1_1], [N^2_1, F^2_1]$ of chunk $C_1$ and in the subsequent chunk $C_2$, where the loss for those tokens is significantly reduced by Retro. Interestingly, in this case, the neighbors were written at a time when the event hadn't yet been postponed."}

:::

::: {caption="Table 19: Daniel Radcliffe, from Wikitext103Valid, retrieval data from c4. The chunks $C_2$ and $C_3$ are almost entirely retrieved from neighbours [$N_1, F_1]$ and [$N_2, F_2]$ respectively, up to formatting differences, which dramatically reduces the loss for these tokens. This example illustrates that when training data leaks into evaluation sets despite deduplication, our Retro model can directly exploit this leakage."}

:::

::: {caption="Table 20: Sample - Déclaration des droits de l'homme: Article premier. The Retro[Off] sample has correct syntax and is almost plausible but is hallucinated. The Retro[On] sample is correctly copied from neighbour data, and robustly re-formated according to our prompt."}

:::

::: {caption="Table 21: Sample - Decimals of π. The Retro[Off] sample quickly diverges two digits after the end of the prompt whereas Retro[On] correctly outputs a large number of π digits, directly copied from the neighbours data."}

:::

References

Section Summary: This section lists academic papers and preprints that provide the technical foundations for modern language modeling research. The works span early recurrent neural networks, the transformer architecture, large-scale training methods, retrieval-augmented approaches, and data preparation techniques. They include both seminal studies from the 2010s and more recent advances through 2021.

[1] T. Mikolov, M. Karafiát, L. Burget, J. Cernockỳ, and S. Khudanpur. Recurrent neural network based language model. Interspeech, 2(3):1045–1048, 2010.

[2] A. Graves. Generating sequences with recurrent neural networks. arXiv preprint arXiv:1308.0850, 2013.

[3] R. Jozefowicz, O. Vinyals, M. Schuster, N. Shazeer, and Y. Wu. Exploring the limits of language modeling. arXiv preprint arXiv:1602.02410, 2016.

[4] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017. URL [https://proceedings.neurips.cc/paper/2017/file/ 3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf](https://proceedings.neurips.cc/paper/2017/file/ 3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf).

[5] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. Preprint, 2019.

[6] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020. URL [https://proceedings.neurips.cc/paper/2020/file/ 1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf](https://proceedings.neurips.cc/paper/2020/file/ 1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf).

[7] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. CoRR, 2020. URL https://arxiv.org/abs/2001.08361.

[8] U. Khandelwal, O. Levy, D. Jurafsky, L. Zettlemoyer, and M. Lewis. Generalization through memorization: Nearest neighbor language models. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=HklBjCEKvH.

[9] D. Yogatama, C. de Masson d’Autume, and L. Kong. Adaptive semiparametric language models. Transactions of the Association for Computational Linguistics, 9:362–373, 2021.

[10] K. Guu, K. Lee, Z. Tung, P. Pasupat, and M. Chang. Retrieval augmented language model pre-training. In International Conference on Machine Learning, 2020.

[11] P. Lewis, E. Perez, A. Piktus, F. Petroni, V. Karpukhin, N. Goyal, H. Küttler, M. Lewis, W.-t. Yih, T. Rocktäschel, S. Riedel, and D. Kiela. Retrieval-augmented generation for knowledge-intensive NLP tasks. In Advances in Neural Information Processing Systems, 2020. URL [https://proceedings.neurips.cc/paper/2020/file/ 6b493230205f780e1bc26945df7481e5-Paper.pdf](https://proceedings.neurips.cc/paper/2020/file/ 6b493230205f780e1bc26945df7481e5-Paper.pdf).

[12] S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=Byj72udxe.

[13] L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.

[14] K. Lee, D. Ippolito, A. Nystrom, C. Zhang, D. Eck, C. Callison-Burch, and N. Carlini. Deduplicating training data makes language models better. arXiv preprint arXiv:2107.06499, 2021.

[15] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Chapter of the Association for Computational Linguistics, June 2019. URL https://aclanthology.org/N19-1423.

[16] J. Rae, S. Borgeaud, T. Cai, K. Millican, J. Hoffmann, F. Song, J. Aslanides, S. Henderson, R. Ring, S. Young, E. Rutherford, T. Hennigan, J. Menick, A. Cassirer, R. Powell, G. van den Driessche, L. A. Hendricks, M. Rauh, P.-S. Huang, A. Glaese, J. Welbl, S. Dathathri, S. Huang, J. Uesato, J. Mellor, I. Higgins, A. Creswell, N. McAleese, A. Wu, E. Elsen, S. Jayakumar, E. Buchatskaya, D. Budden, E. Sutherland, K. Simonyan, M. Paganini, L. Sifre, L. Martens, X. L. Li, A. Kuncoro, A. Nematzadeh, E. Gribovskaya, D. Donato, A. Lazaridou, A. Mensch, J.-B. Lespiau, M. Tsimpoukelli, N. Grigorev, D. Fritz, T. Sottiaux, M. Pajarskas, T. Pohlen, Z. Gong, D. Toyama, C. de Masson d’Autume, Y. Li, T. Terzi, V. Mikulik, I. Babuschkin, A. Clark, D. de Las Casas, A. Guy, J. Bradbury, M. Johnson, B. Hechtman, L. Weidinger, I. Gabriel, W. Isaac, E. Lockhart, S. Osindero, L. Rimell, C. Dyer, O. Vinyals, K. Ayoub, J. Stanway, L. Bennett, D. Hassabis, K. Kavukcuoglu, and G. Irving. Scaling language models: Methods, analysis & insights from training Gopher. arXiv submission, 2021.

[17] T. Kudo and J. Richardson. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv:1808.06226, 2018.

[18] R. Guo, P. Sun, E. Lindgren, Q. Geng, D. Simcha, F. Chern, and S. Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, 2020. URL https://arxiv.org/abs/1908.10396.

[19] B. Zhang and R. Sennrich. Root mean square layer normalization. In Advances in Neural Information Processing Systems, 2019. URL https://proceedings.neurips.cc/paper/2019/file/1e8a19426224ca89e83cef47f1e7f53b-Paper.pdf.

[20] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. Le, and R. Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. In Annual Meeting of the Association for Computational Linguistics, July 2019. URL https://aclanthology.org/P19-1285.

[21] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. V. der Plas, S. Wanderman-Milne, and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.

[22] T. Hennigan, T. Cai, T. Norman, and I. Babuschkin. Haiku: Sonnet for JAX, 2020. URL http://github.com/deepmind/dm-haiku.

[23] T. Brants, A. C. Popat, P. Xu, F. J. Och, and J. Dean. Large Language models in machine translation. In Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning, pages 858–867, 2007.

[24] O. Lieber, O. Sharir, B. Lenz, and Y. Shoham. Jurassic-1: Technical details and evaluation. White Paper. AI21 Labs, 2021.

[25] N. Carlini, F. Tramer, E. Wallace, M. Jagielski, A. Herbert-Voss, K. Lee, A. Roberts, T. Brown, D. Song, U. Erlingsson, A. Oprea, and C. Raffel. Extracting training data from large language models. Preprint, 2021.

[26] P. Lewis, P. Stenetorp, and S. Riedel. Question and answer test-train overlap in open-domain question answering datasets. In Conference of the European Chapter of the Association for Computational Linguistics, Apr. 2021. URL https://aclanthology.org/2021.eacl-main.86.

[27] S. Robertson and H. Zaragoza. The probabilistic relevance framework: BM25 and beyond. Foundations and Trends in Information Retrieval, 3:333–389, Jan 2009.

[28] D. M. Blei, A. Y. Ng, and M. I. Jordan. Latent Dirichlet Allocation. Journal of Machine Learning Research, 3(Jan):993–1022, 2003. URL https://jmlr.csail.mit.edu/papers/v3/blei03a.html.

[29] X. Wei and W. B. Croft. LDA-based document models for ad-hoc retrieval. In ACM SIGIR International Conference on Research and Development in Information Retrieval, 2006. URL http://portal.acm.org/citation.cfm?doid=1148170.1148204.

[30] J. Zhang, M. Utiyama, E. Sumita, G. Neubig, and S. Nakamura. Guiding neural machine translation with retrieved translation pieces. In Conference of the North American Chapter of the Association for Computational Linguistics, 2018.

[31] J. Gu, Y. Wang, K. Cho, and V. O. Li. Search engine guided neural machine translation. In AAAI Conference on Artificial Intelligence, 2018.

[32] S. Ahn, H. Choi, T. Pärnamaa, and Y. Bengio. A neural knowledge language model. arXiv preprint arXiv:1608.00318, 2016.

[33] E. Grave, A. Joulin, and N. Usunier. Improving neural language models with a continuous cache. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=B184E5qee.

[34] V. Karpukhin, B. Oguz, S. Min, P. Lewis, L. Wu, S. Edunov, D. Chen, and W.-t. Yih. Dense passage retrieval for open-domain question answering. In Conference on Empirical Methods in Natural Language Processing, Nov. 2020. URL https://aclanthology.org/2020.emnlp-main.550.

[35] K. Lee, M.-W. Chang, and K. Toutanova. Latent Retrieval for Weakly Supervised Open Domain Question Answering. In Annual Meeting of the Association for Computational Linguistic, June 2019. URL http://arxiv.org/abs/1906.00300.

[36] G. Izacard and E. Grave. Leveraging passage retrieval with generative models for open domain question answering. In Conference of the European Chapter of the Association for Computational Linguistics, Apr. 2021. URL https://aclanthology.org/2021.eacl-main.74.

[37] D. S. Sachan, S. Reddy, W. Hamilton, C. Dyer, and D. Yogatama. End-to-end training of multi-document reader and retriever for open-domain question answering. arXiv preprint arXiv:2106.05346, 2021.

[38] M. Komeili, K. Shuster, and J. Weston. Internet-augmented dialogue generation. arXiv preprint arXiv:2107.07566, 2021.

[39] H. Hashemi, H. Zamani, and W. B. Croft. Guided transformer: Leveraging multiple external sources for representation learning in conversational search. In Proceedings of the 43rd International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 1131–1140, 2020.

[40] E. M. Bender, T. Gebru, A. McMillan-Major, and S. Shmitchell. On the dangers of stochastic parrots: Can language models be too big? In ACM Conference on Fairness, Accountability, and Transparency, 2021.

[41] L. Weidinger, I. Gabriel, C. Griffin, M. Rauh, J. Uesato, J. Mellor, W. Isaac, P.-S. Huang, L. A. Hendricks, M. Cheng, B. Balle, J. Haas, C. Biles, L. Rimell, W. Hawkins, M. Glaese, A. Kasirzadeh, Z. Kenton, S. Brown, A. Birhane, T. Stepleton, G. Irving, and S. Legassick. Ethical and social risks of harm from language models. arXiv submission, 2021.

[42] A. Lazaridou, A. Kuncoro, E. Gribovskaya, D. Agrawal, A. Liska, T. Terzi, M. Gimenez, C. de Masson d'Autume, S. Ruder, D. Yogatama, K. Cao, T. Kociský, S. Young, and P. Blunsom. Pitfalls of static language modelling. CoRR, 2021. URL https://arxiv.org/abs/2102.01951.

[43] S. Gehman, S. Gururangan, M. Sap, Y. Choi, and N. A. Smith. RealToxicityPrompts: Evaluating neural toxic degeneration in language models. In Conference on Empirical Methods in Natural Language Processing, Nov. 2020. URL https://aclanthology.org/2020.findings-emnlp.301.

[44] M. Abadi, A. Chu, I. Goodfellow, H. B. McMahan, I. Mironov, K. Talwar, and L. Zhang. Deep learning with differential privacy. In ACM SIGSAC Conference on Computer and Communications Security, 2016.

[45] E. Strubell, A. Ganesh, and A. McCallum. Energy and policy considerations for deep learning in NLP. In Association for Computational Linguistics, July 2019. URL https://aclanthology.org/P19-1355.

[46] R. Schwartz, J. Dodge, N. A. Smith, and O. Etzioni. Green AI. Communications of the Association for Computing Machinery, 63(12):54–63, Nov. 2020.

[47] E. S. Jo and T. Gebru. Lessons from archives: Strategies for collecting sociocultural data in machine learning. In Proceedings of the 2020 Conference on Fairness, Accountability, and Transparency, pages 306–316, 2020.

[48] Y. Belinkov, S. Gehrmann, and E. Pavlick. Interpretability and analysis in neural NLP. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics: Tutorial Abstracts, pages 1–5, Online, July 2020. Association for Computational Linguistics. doi:10.18653/v1/2020.acl-tutorials.1. URL https://aclanthology.org/2020.acl-tutorials.1.

[49] S. Jain and B. C. Wallace. Attention is not Explanation. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 3543–3556, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi:10.18653/v1/N19-1357. URL https://aclanthology.org/N19-1357.

[50] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67, 2020. URL http://jmlr.org/papers/v21/20-074.html.

[51] Curation. Curation corpus base, 2020.

[52] D. Paperno, G. Kruszewski, A. Lazaridou, N. Q. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Annual Meeting of the Association for Computational Linguistics, Aug. 2016. URL https://aclanthology.org/P16-1144.

[53] A. Baevski and M. Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=ByxZX20qFQ.

[54] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. CoRR, 2019. URL http://arxiv.org/abs/1909.08053.

[55] T. Kwiatkowski, J. Palomaki, O. Redfield, M. Collins, A. Parikh, C. Alberti, D. Epstein, I. Polosukhin, M. Kelcey, J. Devlin, K. Lee, K. N. Toutanova, L. Jones, M.-W. Chang, A. Dai, J. Uszkoreit, Q. Le, and S. Petrov. Natural Questions: a benchmark for question answering research. Transactions of the Association of Computational Linguistics, 7:452–466, Mar. 2019. URL https://aclanthology.org/Q19-1026.

[56] G. Izacard, F. Petroni, L. Hosseini, N. De Cao, S. Riedel, and E. Grave. A memory efficient baseline for open domain question answering. arXiv preprint arXiv:2012.15156, 2020.

[57] K. Shuster, S. Poff, M. Chen, D. Kiela, and J. Weston. Retrieval augmentation reduces hallucination in conversation. arXiv:2104.07567 [cs], Apr. 2021. URL http://arxiv.org/abs/2104.07567.

[58] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=Bkg6RiCqY7.

[59] S. Rajbhandari, J. Rasley, O. Ruwase, and Y. He. Zero: Memory optimizations toward training trillion parameter models. In IEEE International Conference for High Performance Computing, Networking, Storage and Analysis, 2020.

[60] C. Consonni, D. Laniado, and A. Montresor. Wikilinkgraphs: a complete, longitudinal and multi-language dataset of the wikipedia link networks. In AAAI International Conference on Web and Social Media, volume 13, 2019.