retrieval-augmented generation rag

# Retrieval-Augmented Generation (RAG)

## Introduction & Motivation

RAG: augment language models with retrieved documents. Combine retrieval and generation. Applications: QA, knowledge-intensive tasks.

Motivation: Ground language model outputs with external knowledge.

Applications: Fact-based question answering, document-based QA.

---

## Core Concepts & Theory

### Dense Retrieval

Retrieve relevant documents via embeddings.

### Augmentation

Incorporate retrieved context into generation.

### Joint Training

End-to-end retrieval-generation optimization.

### Knowledge Integration

External knowledge incorporation.

---

## Mathematical Formulation

Retrieval:
$$p(z|x) = \frac{\exp(d(x, z))}{\sum_{z'} \exp(d(x, z'))}$$

Generation with Context:
$$p(y|x,z) = ext{decoder}(y | ext{encoder}(x, z))$$

Joint Probability:
$$p(y|x) = \sum_z p(z|x) p(y|x,z)$$

---

## Advanced Theory & Extensions

### Dense Passage Retrieval

Efficient document retrieval.

### Iterative Retrieval

Multi-hop reasoning.

### Fusion-in-Decoder

Combine multiple retrieved documents.

---

## Computational Considerations

Retrieval: O(D) (with efficient indexing).

Encoding: O(T·D).

Generation: O(T·D).

---

## Practical Implementation Strategies

### Approximate Nearest Neighbor Search

Fast retrieval via ANN.

### Hard Negative Mining

Difficult negative selection.

### Knowledge Distillation

Compress retriever.

---

## Benchmark Datasets & Evaluation

Natural Questions: Open-domain QA.

SQuAD: Machine reading comprehension.

TriviaQA: Web-based QA.

---

## Key Challenges & Limitations

### Retrieval Quality

Bottleneck for generation.

### Noise in Retrieved Documents

Irrelevant or contradictory information.

### Computational Efficiency

Large-scale retrieval cost.

---

## Hyperparameter Tuning

Num retrieved docs: 5-20.

Retrieval temperature: 0.5-1.0.

Learning rate: 1e-4 to 5e-4.

---

## Real-World Applications & Case Studies

Open-Domain QA: Web-based question answering.

Fact-Checking: Evidence retrieval for verification.

Legal Document Analysis: Context-aware analysis.

---

## Integration with Other Methods

RAG + dense retrieval; + multi-hop reasoning.

---

## Summary & Key Takeaways

RAG augments generation models with retrieved knowledge.

Principles:
1. Dense retrieval: Embedding-based search.
2. Augmentation: Context incorporation.
3. Joint training: End-to-end optimization.
4. Knowledge: External information integration.
5. Efficiency: Approximate search.

---

## Appendix: Practical Labs

### Lab 1: Dense Retrieval

import numpy as np

def dense_retrieval(query_embedding, document_embeddings, top_k=5):
 """Retrieve top-k documents via dense similarity"""
 # Cosine similarity
 similarities = query_embedding @ document_embeddings.T
 
 # Get top-k
 top_indices = np.argsort(similarities)[-top_k:][::-1]
 top_scores = similarities[top_indices]
 
 return top_indices, top_scores

np.random.seed(42)
query = np.random.randn(768)
query /= np.linalg.norm(query)
docs = np.random.randn(1000, 768)
docs /= np.linalg.norm(docs, axis=1, keepdims=True)
indices, scores = dense_retrieval(query, docs, top_k=5)
assert len(indices) == 5, "Correct retrieval count"
print("✓ Dense retrieval working")

### Lab 2: Augmented Generation Loss

import numpy as np

def rag_loss(query_embedding, retrieved_docs, target_ids, doc_embeddings):
 """Compute RAG loss combining retrieval and generation"""
 # Retrieval loss
 doc_scores = query_embedding @ doc_embeddings.T
 doc_probs = np.exp(doc_scores) / np.sum(np.exp(doc_scores))
 retrieval_loss = -np.log(doc_probs[retrieved_docs[0]] + 1e-8)
 
 # Generation loss (simplified)
 generation_loss = np.mean([1.0 for _ in target_ids])
 
 return retrieval_loss + generation_loss

np.random.seed(42)
query = np.random.randn(768)
docs = np.random.randn(1000, 768)
retrieved = np.array([42, 128, 256])
targets = np.array([1, 2, 3])
loss = rag_loss(query, retrieved, targets, docs)
assert loss > 0, "Positive loss"
print(f"✓ RAG loss: {loss:.4f}")

### Lab 3: Document Augmentation

import numpy as np

def augment_context(query, retrieved_documents, max_length=512):
 """Augment query with retrieved document context"""
 augmented = f"Passages: {' '.join(retrieved_documents[:3])} Question: {query}"
 
 # Truncate to max length
 augmented = augmented[:max_length]
 
 return augmented

query = "What is machine learning?"
docs = ["ML is a subset of AI", "Deep learning uses neural networks", "Supervised learning requires labels"]
augmented = augment_context(query, docs)
assert "Passages:" in augmented, "Context augmented"
assert "Question:" in augmented, "Query included"
print(f"✓ Augmentation: {augmented[:50]}...")

### Lab 4: Fusion-in-Decoder

import numpy as np

def fusion_in_decoder(retrieved_docs, query_embedding):
 """Fuse multiple retrieved documents in decoder"""
 batch_size = retrieved_docs.shape[0]
 fused = np.zeros((batch_size, 768))
 
 for i in range(batch_size):
 # Attention to query
 doc_attention = query_embedding @ retrieved_docs[i].T
 doc_attention = np.exp(doc_attention) / np.sum(np.exp(doc_attention), axis=-1, keepdims=True)
 
 # Aggregate
 fused[i] = np.sum(doc_attention[:, np.newaxis] * retrieved_docs[i], axis=0)
 
 return fused

np.random.seed(42)
docs = np.random.randn(4, 5, 768)
query = np.random.randn(4, 768)
fused = fusion_in_decoder(docs, query)
assert fused.shape == (4, 768), "Correct fused shape"
print("✓ Fusion-in-decoder working")

---