Retrieval-Augmented Training › Code Implementation
Michael BrenndoerferMarch 11, 202611 min read
Language AI HandbookAll chapters Part 66 of 67Frontier Methods of 2025
Retrieval-Augmented Training › Code ImplementationOne section · Sections in this chapter8 / 13 ↓
Sections in this chapter13
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us build a simplified but complete implementation of the RETRO components to develop concrete intuition for the architecture. We will implement the core building blocks: the chunked cross-attention mechanism, a toy retrieval system, and a miniature RETRO-style transformer.
Setting Up: Imports and Utilities
In[4]:
Codeimport numpy as np
import torch
# Set reproducibility
torch.manual_seed(42)
np.random.seed(42)Building a Toy Retrieval Database
We will create a simple retrieval system using cosine similarity over dense embeddings. In a full RETRO implementation, this would be approximate nearest-neighbor search over trillions of chunks, but the interface is identical.
In[5]:
Codeimport torch.nn.functional as F
class ToyRetrievalDatabase:
"""Simple retrieval database using cosine similarity over embeddings."""
def __init__(
self, n_chunks: int = 200, embed_dim: int = 64, vocab_size: int = 100
):
self.n_chunks = n_chunks
self.embed_dim = embed_dim
self.vocab_size = vocab_size
# Simulate chunk embeddings (in practice, computed by a frozen encoder)
self.embeddings = F.normalize(torch.randn(n_chunks, embed_dim), dim=-1)
def retrieve(self, query_embedding: torch.Tensor, k: int = 2) -> list:
"""Retrieve k nearest neighbors by cosine similarity."""
query_norm = F.normalize(query_embedding.unsqueeze(0), dim=-1)
similarities = torch.matmul(query_norm, self.embeddings.T).squeeze(0)
top_k_indices = torch.topk(similarities, k=k).indices.tolist()
return top_k_indices
def get_chunk_tokens(self, idx: int, chunk_len: int = 8) -> torch.Tensor:
"""Return a token tensor for the chunk at index idx."""
# In practice, these would be the actual tokenized chunks
gen = torch.Generator()
gen.manual_seed(idx)
return torch.randint(0, self.vocab_size, (chunk_len,), generator=gen)Chunked Cross-Attention
The core RETRO mechanism is chunked cross-attention. This module computes attention from main-sequence tokens to encoded neighbor tokens:
In[6]:
Codeimport math
import torch.nn as nn
import torch.nn.functional as F
class ChunkedCrossAttention(nn.Module):
"""Chunked cross-attention as in RETRO."""
def __init__(self, d_model: int, n_heads: int):
super().__init__()
self.d_model = d_model
self.n_heads = n_heads
self.d_k = d_model // n_heads
self.q_proj = nn.Linear(d_model, d_model, bias=False)
self.k_proj = nn.Linear(d_model, d_model, bias=False)
self.v_proj = nn.Linear(d_model, d_model, bias=False)
self.out_proj = nn.Linear(d_model, d_model, bias=False)
def forward(
self,
chunk_hidden: torch.Tensor, # (batch, chunk_len, d_model)
neighbor_hidden: torch.Tensor, # (batch, k * neighbor_len, d_model)
) -> torch.Tensor:
B, L, _ = chunk_hidden.shape
M = neighbor_hidden.shape[1]
Q = (
self.q_proj(chunk_hidden)
.view(B, L, self.n_heads, self.d_k)
.transpose(1, 2)
)
K = (
self.k_proj(neighbor_hidden)
.view(B, M, self.n_heads, self.d_k)
.transpose(1, 2)
)
V = (
self.v_proj(neighbor_hidden)
.view(B, M, self.n_heads, self.d_k)
.transpose(1, 2)
)
scale = math.sqrt(self.d_k)
attn_weights = torch.matmul(Q, K.transpose(-2, -1)) / scale
attn_probs = F.softmax(attn_weights, dim=-1)
attended = torch.matmul(attn_probs, V)
attended = attended.transpose(1, 2).contiguous().view(B, L, -1)
return self.out_proj(attended)The Neighbor Encoder
The encoder processes retrieved neighbors bidirectionally. In RETRO, this is a smaller transformer that produces hidden states for the retrieved text:
In[7]:
Codeclass NeighborEncoder(nn.Module):
"""Bidirectional transformer encoder for retrieved neighbors."""
def __init__(
self, vocab_size: int, d_model: int, n_heads: int, n_layers: int
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=d_model * 4,
batch_first=True,
norm_first=True,
)
self.transformer = nn.TransformerEncoder(
encoder_layer, num_layers=n_layers
)
self.d_model = d_model
def forward(self, neighbor_tokens: torch.Tensor) -> torch.Tensor:
"""
neighbor_tokens: (batch, k, neighbor_len)
Returns: (batch, k * neighbor_len, d_model)
"""
B, k, L = neighbor_tokens.shape
flat_tokens = neighbor_tokens.reshape(B * k, L)
embeds = self.embedding(flat_tokens) * math.sqrt(self.d_model)
encoded = self.transformer(embeds)
return encoded.reshape(B, k * L, self.d_model)A Minimal RETRO Layer
Now we put it together: a RETRO transformer layer that combines self-attention with chunked cross-attention:
In[8]:
Codeclass RETROLayer(nn.Module):
"""A single RETRO transformer layer: self-attention + CCA + FFN."""
def __init__(self, d_model: int, n_heads: int, is_retro_layer: bool = True):
super().__init__()
self.is_retro_layer = is_retro_layer
self.d_model = d_model
self.self_attn = nn.MultiheadAttention(
d_model, n_heads, batch_first=True
)
self.norm1 = nn.LayerNorm(d_model)
if is_retro_layer:
self.cca = ChunkedCrossAttention(d_model, n_heads)
self.norm_cca = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.GELU(),
nn.Linear(d_model * 4, d_model),
)
self.norm2 = nn.LayerNorm(d_model)
def forward(
self,
x: torch.Tensor,
attn_mask: torch.Tensor,
neighbor_hidden: torch.Tensor,
chunk_size: int,
) -> torch.Tensor:
B, S, D = x.shape
# 1. Self-attention
x_normed = self.norm1(x)
sa_out, _ = self.self_attn(
x_normed,
x_normed,
x_normed,
attn_mask=attn_mask,
need_weights=False,
)
x = x + sa_out
# 2. Chunked cross-attention (only in RETRO layers)
if self.is_retro_layer and neighbor_hidden is not None:
n_chunks = S // chunk_size
cca_outputs = []
for u in range(n_chunks):
start = u * chunk_size
end = start + chunk_size
chunk_normed = self.norm_cca(x[:, start:end, :])
cca_out = self.cca(chunk_normed, neighbor_hidden)
cca_outputs.append(cca_out)
x = x + torch.cat(cca_outputs, dim=1)
# 3. Feed-forward
x = x + self.ffn(self.norm2(x))
return xA Mini RETRO Model
In[9]:
Codeclass MiniRETRO(nn.Module):
"""A small RETRO-style language model. Every other layer includes CCA."""
def __init__(
self,
vocab_size: int = 100,
d_model: int = 64,
n_heads: int = 4,
n_layers: int = 4,
chunk_size: int = 8,
k: int = 2,
neighbor_len: int = 8,
):
super().__init__()
self.vocab_size = vocab_size
self.d_model = d_model
self.chunk_size = chunk_size
self.k = k
self.neighbor_len = neighbor_len
self.token_emb = nn.Embedding(vocab_size, d_model)
self.pos_emb = nn.Embedding(512, d_model)
# Neighbor encoder: smaller than main model
self.encoder = NeighborEncoder(
vocab_size, d_model // 2, n_heads // 2, n_layers=1
)
self.enc_proj = nn.Linear(d_model // 2, d_model)
# Alternate standard and RETRO layers
self.layers = nn.ModuleList(
[
RETROLayer(d_model, n_heads, is_retro_layer=(i % 2 == 1))
for i in range(n_layers)
]
)
self.norm_final = nn.LayerNorm(d_model)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
def forward(
self,
input_ids: torch.Tensor, # (batch, seq_len)
neighbor_ids: torch.Tensor, # (batch, n_chunks, k, neighbor_len)
) -> torch.Tensor:
B, S = input_ids.shape
device = input_ids.device
positions = torch.arange(S, device=device).unsqueeze(0)
x = self.token_emb(input_ids) + self.pos_emb(positions)
causal_mask = torch.triu(
torch.full((S, S), float("-inf"), device=device), diagonal=1
)
# Encode neighbors for the first chunk (demo simplification)
chunk_neighbors = neighbor_ids[:, 0, :, :] # (B, k, neighbor_len)
enc_out = self.encoder(chunk_neighbors) # (B, k * neighbor_len, d//2)
neighbor_h = self.enc_proj(enc_out) # (B, k * neighbor_len, d)
for layer in self.layers:
if layer.is_retro_layer:
x = layer(x, causal_mask, neighbor_h, self.chunk_size)
else:
x = layer(x, causal_mask, None, self.chunk_size)
return self.lm_head(self.norm_final(x))Training a Mini RETRO Model
Let us train this model on a simple synthetic language modeling task and compare it against a standard transformer baseline:
In[10]:
Codedef make_causal_lm_data(
vocab_size: int = 100,
n_samples: int = 500,
seq_len: int = 32,
chunk_size: int = 8,
k: int = 2,
neighbor_len: int = 8,
):
input_ids = torch.randint(0, vocab_size, (n_samples, seq_len))
n_chunks = seq_len // chunk_size
neighbor_ids = torch.randint(
0, vocab_size, (n_samples, n_chunks, k, neighbor_len)
)
labels = torch.cat(
[input_ids[:, 1:], torch.zeros(n_samples, 1, dtype=torch.long)], dim=1
)
return input_ids, neighbor_ids, labels
class StandardLM(nn.Module):
"""Standard GPT-like transformer for comparison."""
def __init__(
self,
vocab_size: int = 100,
d_model: int = 64,
n_heads: int = 4,
n_layers: int = 4,
):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, d_model)
self.pos_emb = nn.Embedding(512, d_model)
layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=d_model * 4,
batch_first=True,
norm_first=True,
)
self.transformer = nn.TransformerEncoder(layer, num_layers=n_layers)
self.norm = nn.LayerNorm(d_model)
self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
B, S = input_ids.shape
pos = torch.arange(S, device=input_ids.device).unsqueeze(0)
x = self.token_emb(input_ids) + self.pos_emb(pos)
mask = torch.triu(
torch.full((S, S), float("-inf"), device=input_ids.device),
diagonal=1,
)
x = self.transformer(x, mask=mask, is_causal=True)
return self.lm_head(self.norm(x))
VOCAB_SIZE = 100
SEQ_LEN = 32
CHUNK_SIZE = 8
K = 2
NEIGHBOR_LEN = 8
BATCH_SIZE = 32
N_EPOCHS = 40
input_ids, neighbor_ids, labels = make_causal_lm_data(
vocab_size=VOCAB_SIZE,
n_samples=500,
seq_len=SEQ_LEN,
chunk_size=CHUNK_SIZE,
k=K,
neighbor_len=NEIGHBOR_LEN,
)
retro_model = MiniRETRO(
vocab_size=VOCAB_SIZE,
d_model=64,
n_heads=4,
n_layers=4,
chunk_size=CHUNK_SIZE,
k=K,
neighbor_len=NEIGHBOR_LEN,
)
standard_model = StandardLM(VOCAB_SIZE)
retro_opt = torch.optim.Adam(retro_model.parameters(), lr=3e-3)
standard_opt = torch.optim.Adam(standard_model.parameters(), lr=3e-3)
retro_losses = []
standard_losses = []
n_batches = len(input_ids) // BATCH_SIZE
for epoch in range(N_EPOCHS):
perm = torch.randperm(len(input_ids))
epoch_retro_loss = 0.0
epoch_standard_loss = 0.0
for b in range(n_batches):
idx = perm[b * BATCH_SIZE : (b + 1) * BATCH_SIZE]
x_batch = input_ids[idx]
n_batch = neighbor_ids[idx]
y_batch = labels[idx]
retro_opt.zero_grad()
retro_logits = retro_model(x_batch, n_batch)
retro_loss = F.cross_entropy(
retro_logits.view(-1, VOCAB_SIZE), y_batch.view(-1)
)
retro_loss.backward()
retro_opt.step()
epoch_retro_loss += retro_loss.item()
standard_opt.zero_grad()
std_logits = standard_model(x_batch)
std_loss = F.cross_entropy(
std_logits.view(-1, VOCAB_SIZE), y_batch.view(-1)
)
std_loss.backward()
standard_opt.step()
epoch_standard_loss += std_loss.item()
retro_losses.append(epoch_retro_loss / n_batches)
standard_losses.append(epoch_standard_loss / n_batches)
final_retro_loss = retro_losses[-1]
final_standard_loss = standard_losses[-1]Out[11]:
ConsoleFinal training loss (RETRO): 0.0034 Final training loss (Standard): 1.9122 RETRO parameter count: 296,672 Standard parameter count: 245,632
On this synthetic task with random neighbors, both models learn to reduce training loss, but they do not form a controlled quality comparison. The RETRO model has 296,672 parameters versus 245,632 for the standard model, and its training loss falls much faster. Because the neighbors carry no semantic signal and the chart shows training rather than held-out loss, that gap should be read as extra capacity fitting the synthetic examples, not as evidence that retrieval improves generalization. The demo captures the architectural mechanics but not RETRO's representational benefit on meaningful retrievals.
Out[12]:
Visualization
Visualizing Attention Patterns in CCA
One of the most instructive things we can examine is what the chunked cross-attention attends to. If the cross-attention is working correctly, tokens in a chunk should develop non-uniform attention patterns over neighbor tokens:
In[13]:
Codedef get_cca_attention_weights(model, input_ids, neighbor_ids):
"""Extract CCA attention weights from the first RETRO layer."""
model.eval()
with torch.no_grad():
B, S = input_ids.shape
positions = torch.arange(S).unsqueeze(0)
x = model.token_emb(input_ids) + model.pos_emb(positions)
chunk_neighbors = neighbor_ids[:, 0, :, :]
enc_out = model.encoder(chunk_neighbors)
neighbor_h = model.enc_proj(enc_out)
causal_mask = torch.triu(torch.full((S, S), float("-inf")), diagonal=1)
# Forward through standard layer first (layer 0)
x = model.layers[0](x, causal_mask, None, model.chunk_size)
# Manually extract CCA weights from layer 1 (first RETRO layer)
retro_layer = model.layers[1]
x_normed = retro_layer.norm1(x)
sa_out, _ = retro_layer.self_attn(
x_normed,
x_normed,
x_normed,
attn_mask=causal_mask,
need_weights=False,
)
x = x + sa_out
chunk_h = x[:, : model.chunk_size, :]
chunk_normed = retro_layer.norm_cca(chunk_h)
cca = retro_layer.cca
chunk_len = chunk_h.shape[1]
neighbor_len = neighbor_h.shape[1]
Q = (
cca.q_proj(chunk_normed)
.view(B, chunk_len, cca.n_heads, cca.d_k)
.transpose(1, 2)
)
K = (
cca.k_proj(neighbor_h)
.view(B, neighbor_len, cca.n_heads, cca.d_k)
.transpose(1, 2)
)
scale = math.sqrt(cca.d_k)
attn = F.softmax(torch.matmul(Q, K.transpose(-2, -1)) / scale, dim=-1)
attn_weights = attn.mean(dim=(0, 1)).cpu().numpy()
return attn_weights
sample_input = input_ids[:1]
sample_neighbors = neighbor_ids[:1]
attn_weights = get_cca_attention_weights(
retro_model, sample_input, sample_neighbors
)Out[14]:
ConsoleCCA attention weight matrix shape: (8, 16) rows = chunk tokens (8) cols = neighbor tokens (16): 2 neighbors x 8 tokens each Attention weight statistics: Mean: 0.0625 (uniform baseline = 0.0625) Max: 0.3121 Std: 0.0725 Attention entropy (mean): 2.2222 / max 2.7726 Entropy ratio: 0.801 (1.0 = fully uniform, lower = more focused)
The attention entropy ratio tells us how selective the CCA is. A ratio near 1.0 means the model attends roughly equally to all neighbor tokens. Here the ratio is about 0.80, so the learned attention is clearly non-uniform. Because this demonstration uses random neighbors and measures training behavior only, however, individual bright cells cannot be interpreted as semantically relevant retrievals; they show that the CCA mechanism can form selective patterns on the synthetic task.
Out[15]:
Visualization