Retrieval-Augmented Training › The RETRO Architecture

RETRO introduces a specific architectural pattern for integrating retrieved sequences into the forward pass of a transformer language model. Understanding it requires first understanding how the retrieval interface is structured.

The core insight in RETRO is that a long sequence can be processed in fixed-length chunks, and each chunk can be enriched with retrieved material that is relevant to the local context up to that point. Rather than trying to process the entire sequence against a global retrieved set, RETRO maintains alignment between local sequence segments and their retrieved neighbors.

Formally, given an input sequence of tokens $x_1, x_2, \ldots, x_n$, RETRO splits it into non-overlapping chunks of size $l$ (typically 64 tokens):

$$C_u = (x_{(u-1)l + 1}, \ldots, x_{ul})$$

where:

• $C_u$ is the $u$-th chunk
• $l$ is the chunk length (64 tokens in the original paper)
• $u$ ranges from 1 to $n / l$

For each chunk $C_u$, RETRO retrieves $k$ neighbor sequences from the database using the chunk as a query. Each neighbor $N_u^j$ consists of two parts: the retrieved passage itself (the "neighbor") plus a continuation (the "neighbor-plus") which is the text that follows the neighbor in the original document. This continuation provides additional context that is informative but was not the direct retrieval target.

During training, the model processes chunks sequentially, but the retrieval for chunk $C_u$ is based on all tokens seen up to and including chunk $C_{u-1}$. This is a causal constraint: you cannot use chunk $C_u$ itself to retrieve its neighbors, because at inference time you would be retrieving before you have generated those tokens. The retrieval query for chunk $C_u$ is formed from the prefix $C_1, \ldots, C_{u-1}$.

The chunked cross-attention (CCA) mechanism is the architectural innovation that implements this integration. At specific layers in the transformer stack, the representation of each token in chunk $C_u$ attends to the encoded representations of the retrieved neighbors $N_u^1, \ldots, N_u^k$. This is different from standard self-attention in several ways:

• The keys and values come from a separate encoder that processes the retrieved sequences
• Only tokens within the same chunk attend to the same retrieved set
• The cross-attention is causally masked: tokens in chunk $C_u$ can only attend to retrieved neighbors computed from the prefix $C_1, \ldots, C_{u-1}$

The CCA operation for a token at position $i$ within chunk $C_u$ computes:

$$\text{CCA}(q_i, K_u, V_u) = \text{softmax}\left(\frac{q_i K_u^T}{\sqrt{d_k}}\right) V_u$$

where:

• $q_i$ is the query vector for token $i$ derived from its hidden state
• $K_u$ and $V_u$ are the key and value matrices derived from the encoded neighbor sequences for chunk $u$
• $d_k$ is the key dimension

The encoder that produces $K_u$ and $V_u$ is a bidirectional transformer (since the neighbors are fully observed and do not need causal masking). It processes each retrieved neighbor independently and produces a sequence of hidden states that the main model can attend to through CCA.

The full RETRO model interleaves two types of transformer layers.

A standard transformer layer applies multi-head self-attention followed by a feed-forward network in the usual way. No retrieved context is used in these layers.

A RETRO layer contains three sub-components applied in sequence. First, standard self-attention over the input sequence. Second, chunked cross-attention over the encoded neighbors. Third, a feed-forward network. The CCA step is inserted between self-attention and the feed-forward network, after the self-attention has already contextualized the token representations within the sequence.

Not every layer is a RETRO layer. Typically, CCA is only applied in every third or fourth layer, allowing the model to first build up contextualized representations through self-attention before integrating the retrieved material. This design reflects the intuition that retrieval integration is most effective when the model already has a rich understanding of the local context.

The encoder used to process neighbors is significantly smaller than the main model. In the original RETRO paper, the main model has up to 7 billion parameters while the encoder uses approximately the same architecture as a standard 600 million parameter transformer. The encoder is frozen during much of the training process in some configurations, though the original paper trains the full system end-to-end.

The retrieval component is a massive key-value store. In the original RETRO experiments, this database contains 2 trillion tokens drawn from MassiveText, which is the same corpus used to train Gopher and other DeepMind models. The database is split into 64-token chunks, and each chunk is associated with a dense embedding vector computed by a frozen sentence encoder (BERT-based in the original paper).

At retrieval time, the query embedding is computed from the prefix tokens and compared to all chunk embeddings using approximate nearest-neighbor search (specifically, SCaNN in the original implementation). The top-$k$ neighbors by embedding similarity are retrieved, along with their continuations.

The choice to use frozen embeddings for retrieval rather than training the retrieval system end-to-end with the model is a practical compromise. Training a retrieval system that indexes trillions of tokens requires recomputing and reindexing embeddings periodically as the encoder changes, which is computationally expensive. Using a frozen retriever eliminates this cost at the expense of some potential quality.