Unify retrieval and generation in RAG systems by compressing documents into shared continuous embeddings that serve both retrieval and generation: implement joint training with differentiable selection, achieving up to 16× context compression while improving generation quality.
Unify retrieval and generation in RAG systems by compressing documents into shared continuous embeddings that serve both retrieval and generation: implement joint training with differentiable selection, achieving up to 16× context compression while improving generation quality.
CLaRa: Bridging Retrieval and Generation
Traditional RAG systems treat retrieval and generation as separate modules: retrievers select full documents while generators consume raw text, creating redundancy and representation mismatch. This skill demonstrates how to unify both components through shared continuous document representations that simultaneously serve retrieval ranking and generation context—all optimized jointly with gradients flowing bidirectionally.
The core innovation is using compact compressed representations for both ranking and generation, eliminating the gap between what retrievers optimize for and what generators need.
Core Concept
CLaRa (Continuous Latent Reasoning for Augmented generation) implements:
Salient Compressor Pretraining (SCP): Learns to compress documents into meaningful vectors by synthesizing QA pairs and paraphrases
Joint CLaRa Training: Query reasoner and answer generator train end-to-end with differentiable ranking of compressed documents
Bidirectional Gradient Flow: Retrieval relevance aligns with downstream generation quality through unified objective
Architecture Overview
Document Compressor: Encodes full documents into compact semantic vectors
Query Reasoner: Maps questions to embedding space for retrieval
Differentiable Selector: Learns to rank documents using Straight-Through estimation
Answer Generator: Consumes compressed documents for generation
Joint Loss: Single objective optimizing retrieval + generation together
Implementation Steps
The system trains through pretraining and joint optimization phases.
1. Build Document Compressor with Salient Content Focus
Train encoder to compress documents while preserving information-rich features.
defcreate_synthetic_qa_pairs(documents, qa_generation_model, num_pairs_per_doc=5):
"""
Generate synthetic QA pairs from documents.
Uses LLM to create questions about key information.
Args:
documents: List of document texts
qa_generation_model: LLM for QA generation
num_pairs_per_doc: How many QA pairs per document
Returns:
qa_pairs: List of (question, answer, document) tuples
"""
qa_pairs = []
for doc in documents:
# Generate multiple QA pairs from document
prompt = f"""
Given this document:
{doc}
Generate {num_pairs_per_doc} QA pairs testing understanding of KEY information.
Format: Question: ... Answer: ...
"""
response = qa_generation_model.generate(prompt)
pairs = parse_qa_pairs(response)
for q, a in pairs:
qa_pairs.append((q, a, doc))
return qa_pairs
3. Pretraining Phase: SCP
Pretrain compressor to reconstruct document content from compressed vectors.
defpretrain_salient_compressor(compressor, documents, qa_pairs, num_epochs=10):
"""
Pretrain compressor to preserve salient information.
Uses reconstruction loss + QA answering auxiliary loss.
"""
optimizer = torch.optim.Adam(compressor.parameters(), lr=1e-4)
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
for epoch inrange(num_epochs):
for doc_idx, document inenumerate(documents):
# Tokenize document
tokens = tokenizer(document, return_tensors='pt')['input_ids']
doc_embeddings = get_embedding_layer()(tokens)
# Compress document
compressed, reconstruction = compressor(doc_embeddings)
# Loss 1: Reconstruction (preserve content)
reconstruction_loss = torch.nn.functional.mse_loss(
reconstruction, doc_embeddings
)
# Loss 2: QA auxiliary (answer questions about document)
doc_qa_pairs = [qa for qa in qa_pairs if qa[2] == document]
qa_loss = 0.0for question, answer, _ in doc_qa_pairs[:3]: # Use top 3 QA pairs# Can question be answered from compressed representation?# Simplified: use LLM to verify
can_answer = verify_qa_from_compressed(
compressed, question, answer
)
qa_loss += (1.0 - can_answer)
# Combined loss
total_loss = reconstruction_loss + 0.5 * (qa_loss / max(len(doc_qa_pairs), 1))
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
return compressor
4. Implement Differentiable Document Selector
Create ranking mechanism with gradient flow for joint training.
classDifferentiableDocumentSelector(torch.nn.Module):
"""
Selects top-k documents for generation with differentiable ranking.
Uses Straight-Through estimation to enable gradient flow through discrete selection.
"""def__init__(self, num_docs_to_select=3):
super().__init__()
self.k = num_docs_to_select
defforward(self, query_embedding, document_embeddings, temperature=0.1):
"""
Rank documents and select top-k with gradient flow.
Args:
query_embedding: (batch, embed_dim)
document_embeddings: (batch, num_docs, embed_dim)
temperature: Controls softness of selection
Returns:
selected_docs: (batch, k, embed_dim)
selection_weights: (batch, num_docs) for gradient computation
"""
batch_size, num_docs, embed_dim = document_embeddings.shape
# Compute relevance scores (cosine similarity)# Normalize embeddings
query_norm = query_embedding / (torch.norm(query_embedding, dim=-1, keepdim=True) + 1e-8)
docs_norm = document_embeddings / (torch.norm(document_embeddings, dim=-1, keepdim=True) + 1e-8)
# Similarity: (batch, num_docs)
scores = torch.bmm(docs_norm, query_norm.unsqueeze(-1)).squeeze(-1)
# Soft selection weights using temperature-scaled softmax
soft_weights = torch.softmax(scores / temperature, dim=-1)
# Hard selection: top-k (discrete operation)
_, top_k_indices = torch.topk(scores, k=self.k, dim=-1)
# Straight-Through Estimator: use soft weights for backward, hard for forward
hard_selection = torch.zeros_like(soft_weights)
hard_selection.scatter_(1, top_k_indices, 1.0)
# Straight-through: forward with hard, backward with soft
selection_weights = hard_selection - soft_weights.detach() + soft_weights
# Select documents using soft weights (for gradient flow)
weighted_docs = torch.einsum('bnd,bn->bnd', document_embeddings, selection_weights)
# Sum to get final selected context
selected_docs = weighted_docs.sum(dim=1) # (batch, embed_dim)return selected_docs, selection_weights
5. Implement Joint CLaRa Training
Train query reasoner, selector, and generator end-to-end.
defclara_joint_training(
compressor,
query_reasoner,
selector,
answer_generator,
batch,
optimizer,
learning_rate=1e-4):
"""
Single joint training step for CLaRa.
Optimizes retrieval (via selector) and generation together.
Args:
compressor: Document compression model
query_reasoner: Maps questions to embedding space
selector: Differentiable document ranking
answer_generator: Text generator using selected documents
batch: Dict with 'questions', 'documents', 'answers'
optimizer: PyTorch optimizer
learning_rate: Gradient step size
Returns:
loss: Combined retrieval + generation loss
"""
questions = batch['questions']
documents = batch['documents']
answers = batch['answers']
# Step 1: Compress all documents
doc_embeddings = []
for doc in documents:
doc_tokens = tokenizer(doc, return_tensors='pt')['input_ids']
doc_emb = embedding_layer(doc_tokens)
compressed, _ = compressor(doc_emb)
doc_embeddings.append(compressed)
doc_embeddings = torch.stack(doc_embeddings) # (batch, num_docs, compressed_dim)# Step 2: Encode questions
query_embeddings = []
for q in questions:
q_tokens = tokenizer(q, return_tensors='pt')['input_ids']
q_emb = embedding_layer(q_tokens)
q_encoded = query_reasoner(q_emb)
query_embeddings.append(q_encoded)
query_embeddings = torch.stack(query_embeddings) # (batch, embed_dim)# Step 3: Select documents with gradient flow
selected_docs, selection_weights = selector(
query_embeddings, doc_embeddings, temperature=0.1
)
# Step 4: Generate answers using selected documents
decoder_input = torch.cat([query_embeddings, selected_docs], dim=-1)
logits = answer_generator(decoder_input)
# Step 5: Compute losses# Loss 1: Generation loss (standard next-token prediction)
generation_loss = torch.nn.functional.cross_entropy(
logits[:, :-1].reshape(-1, logits.shape[-1]),
answers.reshape(-1)
)
# Loss 2: Retrieval loss (encourage selecting correct documents)# Simplified: documents that enable better generation get higher "relevance" signal
retrieval_loss = -generation_loss.detach() # Gradient signal flows backward through selector# Combined loss
total_loss = generation_loss + 0.5 * retrieval_loss
# Backward and update
optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(
list(compressor.parameters()) + list(query_reasoner.parameters()) + list(selector.parameters()) + list(answer_generator.parameters()),
1.0
)
optimizer.step()
return {
'total_loss': total_loss.item(),
'generation_loss': generation_loss.item(),
'retrieval_loss': retrieval_loss.item()
}
6. Inference with Compression
Use trained system for efficient retrieval-augmented generation.
defclara_inference(
compressor,
query_reasoner,
selector,
answer_generator,
question,
document_corpus,
tokenizer,
max_answer_length=128):
"""
Generate answer using CLaRa with compressed document retrieval.
Args:
compressor, query_reasoner, selector, answer_generator: Trained modules
question: Input question string
document_corpus: List of available documents
tokenizer: Text tokenizer
max_answer_length: Maximum generation length
Returns:
answer: Generated text
used_documents: Indices of retrieved documents
"""# Compress document corpus
doc_embeddings = []
for doc in document_corpus:
doc_tokens = tokenizer(doc, return_tensors='pt')['input_ids']
doc_emb = embedding_layer(doc_tokens)
compressed, _ = compressor(doc_emb)
doc_embeddings.append(compressed)
doc_embeddings = torch.stack(doc_embeddings)
# Encode question
q_tokens = tokenizer(question, return_tensors='pt')['input_ids']
q_emb = embedding_layer(q_tokens)
query_embedding = query_reasoner(q_emb)
# Select documents
selected_docs, selection_weights = selector(query_embedding, doc_embeddings)
# Track which documents were selected
used_doc_indices = torch.topk(selection_weights.squeeze(), k=3).indices.tolist()
# Generate answer
decoder_input = torch.cat([query_embedding, selected_docs], dim=-1)
answer_tokens = answer_generator.generate(
inputs_embeds=decoder_input,
max_length=max_answer_length,
num_beams=3
)
answer = tokenizer.decode(answer_tokens[0], skip_special_tokens=True)
return answer, used_doc_indices
Practical Guidance
When to Use CLaRa:
RAG systems where document count is large (>1000) and context compression is valuable
Tasks benefiting from unified optimization of retrieval and generation
Scenarios where reducing context length is important for latency
Systems where document relevance varies significantly
When NOT to Use:
Tasks with small document corpora (<100 documents)
Scenarios requiring full-document context for quality
Systems where interpretability of retrieval decisions is critical
Key Hyperparameters:
compressed_dim: Document vector size (64-256; larger retains more info but increases memory)
num_docs_to_select: Top-k documents for generation (2-5 typical)
temperature: Softness of document selection (0.05-0.2)
scp_epochs: Pretraining iterations (5-15)
lambda_retrieval: Weight of retrieval loss (0.3-1.0)
Compression Effectiveness:
Average compression ratio: 10-16× (1000-token docs → 64-128-dim vectors)
Pretraining on synthetic QA improves downstream performance by 15-20%
