Retrieval-Augmented Training › Integrating Retrieved Context: The Mechanics

Understanding how retrieved text is used in the forward pass requires working through the architecture in detail.

The neighbor encoder is a bidirectional transformer that processes each retrieved sequence independently. Given a retrieved neighbor $N_u^j$ of length $m$ tokens (typically 128 tokens: 64 for the neighbor plus 64 for the continuation), the encoder produces a sequence of hidden states:

$$H_u^j = \text{Encoder}(N_u^j) \in \mathbb{R}^{m \times d_{\text{enc}}}$$

where $d_{\text{enc}}$ is the encoder hidden dimension. These hidden states are the keys and values that the main model attends to through CCA.

Because the encoder is bidirectional, each token in the neighbor sequence attends to all other tokens, including those that come after it. This is appropriate because the neighbor is fully observed at inference time. The model is not generating the neighbor; it is reading it.

The encoder processes all $k$ neighbors in parallel. Their outputs are concatenated along the sequence dimension to form the full key-value matrix that the main model attends to for a given chunk.

When the main model reaches a CCA layer, it performs cross-attention from the token representations in chunk $C_u$ to the encoded neighbor representations. Each token in $C_u$ can attend to all tokens across all $k$ encoded neighbors.

The cross-attention mechanism computes:

$$\begin{aligned} Q_u &= H_u^{\text{main}} W_Q \\ K_u &= H_u^{\text{enc}} W_K \\ V_u &= H_u^{\text{enc}} W_V \\ \text{CCA}(C_u) &= \text{softmax}\left(\frac{Q_u K_u^T}{\sqrt{d_k}}\right) V_u \end{aligned}$$

where:

• $H_u^{\text{main}}$ is the main model's hidden state for tokens in chunk $C_u$
• $H_u^{\text{enc}}$ is the concatenated encoded representation of all $k$ neighbors for chunk $u$
• $W_Q, W_K, W_V$ are learned projection matrices
• $d_k$ is the key dimension

The output of CCA is combined with the main model's self-attention output through a residual connection and layer normalization, before being passed to the feed-forward sublayer.

This design has a natural interpretation. The self-attention over $C_u$ allows tokens to contextualize themselves relative to the current local and global sequence context. The cross-attention to neighbors then allows those contextualized representations to be enriched with retrieved information. The feed-forward network integrates these two sources of information into the final token representations.

The causal constraint in RETRO creates a one-chunk lag in retrieval. When predicting tokens in chunk $C_u$, the model uses retrieved neighbors queried from the prefix ending at chunk $C_{u-1}$. This means the model never retrieves information that is directly relevant to the tokens it is currently predicting based on those same tokens, which preserves autoregressive causality.

In practice, this lag is small (64 tokens) and the retrieved neighbors are typically relevant because they were retrieved using the preceding context, which is closely related to the current context in most natural text. A text that discusses the French Revolution in one chunk will almost certainly continue to discuss it in the next chunk, so neighbors retrieved for the preceding context remain relevant.

The lag does create a theoretical gap at sequence boundaries, where the very beginning of a new topic might not have optimal retrievals because the preceding context established a different topic. This is a known limitation of the chunked approach.