Retrieval-Augmented Training › Training Dynamics with Retrieval

Training RETRO from scratch involves several considerations that differ from standard pretraining.

Every training step requires actual retrieval operations. For each training batch, the input sequences must be split into chunks, each chunk must be encoded into a query vector, and the $k$ nearest neighbors must be retrieved from the database. This adds a non-trivial overhead to each training step compared to a standard dense transformer.

In practice, RETRO mitigates this overhead by pre-computing retrievals offline before training begins. A retrieval index is built once from the training corpus, and for each training sequence, the neighbors are retrieved and cached alongside the training examples. During the actual training forward pass, the retrieved neighbors are simply looked up from this cache rather than computed dynamically. This converts the online retrieval problem into a storage problem: you need to store the cached neighbors for all training sequences, which is significant but manageable.

This pre-computation has an interesting implication: the retrievals are static during training. As the main model's parameters change, the optimal neighbors for each query might change, but the cached neighbors do not update. In the original RETRO paper, the retrievals are computed once and held fixed throughout training. This is a further trade-off that likely reduces quality somewhat but is necessary for practical training.

The RETRO training objective is standard autoregressive language modeling loss. The model is trained to predict the next token given all previous tokens, conditioned on the retrieved neighbors for the relevant chunk. Formally:

$$\mathcal{L} = -\sum_{i=1}^{n} \log p_\theta\left(x_i \mid x_{<i}, \text{RET}(C_{\lfloor i/l \rfloor - 1})\right)$$

where:

• $x_{<i}$ denotes all tokens preceding position $i$
• $\text{RET}(C_{u-1})$ denotes the retrieved neighbors computed from the prefix up to chunk $u-1$
• $l$ is the chunk length
• $\theta$ represents all model parameters

The key point is that for tokens in the first chunk ($u = 1$), there are no previous chunks to use as a retrieval query, so the model must predict those tokens without any retrieved context. The retrieval benefit only kicks in starting from the second chunk. This is analogous to how inference-time RAG typically requires a full query before retrieval can occur.

The original RETRO paper trains the full architecture from scratch with retrieval. However, a subsequent line of work has explored "retrofitting" a standard pretrained model to use the RETRO architecture, a process called RETROfit. In this approach, a standard GPT-style model is extended with the CCA layers and encoder, and these new components are trained while the original model weights are frozen or lightly fine-tuned.

RETROfit is appealing because it avoids the enormous cost of pretraining from scratch. A 7 billion parameter model trained from scratch with RETRO requires the same compute as pretraining a 7 billion parameter GPT model, plus the overhead of retrieval during training. RETROfit allows reuse of existing pretrained weights.

The trade-off is quality. When trained from scratch, the model can develop retrieval-aware representations in its self-attention layers from the very beginning, integrating retrieval knowledge throughout the network. In RETROfit, the existing self-attention weights were optimized without any retrieval signal, so they may not be ideal for the CCA integration. Empirically, RETROfit achieves most but not all of the quality benefit of full RETRO training, with substantially lower training cost.