Retrieval-Augmented Generation (RAG) Systems

1. The RAG Paradigm

Retrieval-Augmented Generation (RAG) augments a language model's generation process by retrieving relevant information from an external knowledge base before producing a response. This addresses fundamental limitations of parametric knowledge in LLMs:

• Knowledge cutoff: LLMs are limited to training data up to a certain date
• Hallucination: LLMs can generate plausible but incorrect information
• Domain specificity: General LLMs may lack expertise in specialized domains
• Verifiability: RAG provides source attribution for generated claims

1.1 Formal Definition

Given a query , a retrieval corpus , and a language model , RAG produces:

where retrieves the most relevant documents, and is the retrieval distribution (often uniform over top-K or based on relevance scores).

2. RAG Architecture

3. Document Chunking Strategies

3.1 Why Chunking?

LLMs have context windows limited to tokens. Documents in the corpus often exceed this limit. Chunking decomposes each document into smaller segments:

where each chunk has length .

3.2 Chunking Methods

Fixed-size chunking: Split at every tokens with overlap :

Typical values: , tokens.

Recursive chunking: Split hierarchically by structural boundaries (paragraphs → sentences → tokens):

Semantic chunking: Use embeddings to detect topic boundaries:

Split when (cosine similarity threshold).

Agentic chunking: Use an LLM to determine optimal chunk boundaries based on semantic coherence.

3.3 Chunk Overlap

Overlap prevents information loss at boundaries. For chunks with overlap :

where is the stride (chunk size minus overlap).

4. Dense Retrieval

4.1 Bi-Encoder Architecture

Dense retrieval encodes queries and documents into a shared embedding space:

The relevance score is the cosine similarity (or dot product):

4.2 Training Objectives

In-batch negatives: For a batch of query-document pairs , the loss uses all other documents as negatives:

where is the temperature parameter.

Hard negative mining: Include documents that are similar but not relevant:

4.3 Contrastive Retrieval Loss

The full contrastive loss combines multiple negatives:

4.4 Implementation with Sentence Transformers

from sentence_transformers import SentenceTransformer, losses, InputExample
from torch.utils.data import DataLoader

model = SentenceTransformer("BAAI/bge-large-en-v1.5")

# Training data: (query, positive_doc, hard_negative)
train_examples = [
    InputExample(texts=["What is RLHF?", "RLHF is a technique...", "RLHF stands for..."]),
    InputExample(texts=["How does backprop work?", "Backpropagation computes gradients...", "The loss function is..."]),
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32)

# MultipleNegativesRankingLoss uses in-batch negatives
train_loss = losses.MultipleNegativesRankingLoss(
    model=model,
    scale=20.0,  # 1/temperature
    name="cosine"
)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    output_path="retriever-model"
)

5. Vector Databases and Indexing

5.1 Approximate Nearest Neighbor (ANN) Search

Exact nearest neighbor search has complexity where is the number of vectors and is the dimension. ANN algorithms achieve sub-linear time at the cost of approximate results.

HNSW (Hierarchical Navigable Small World):

• Builds a multi-layer graph where each node connects to neighbors
• Search starts at the top layer and navigates down
• Time complexity: per query

IVF (Inverted File Index):

• Partitions vectors into Voronoi cells using k-means
• At query time, searches only the closest clusters
• Time complexity:

5.2 Quantization

Product Quantization (PQ) reduces memory by encoding vectors as compact codes:

where is a sub-vector and maps to one of centroids.

Memory reduction: from bytes (float32) to bits.

6. Re-ranking

6.1 Cross-Encoder Re-ranking

After initial retrieval (bi-encoder), a cross-encoder re-ranks the top- results by jointly encoding the query and document:

The cross-encoder achieves higher accuracy than bi-encoders due to full cross-attention between query and document, but is times slower (must score each pair independently).

6.2 Re-ranking Formula

Given initial retrieval scores and cross-encoder scores , the final score combines both:

where balances the two signals.

6.3 ColBERT: Late Interaction

ColBERT encodes queries and documents independently but uses MaxSim for scoring:

where and are token-level embeddings. This captures fine-grained token interactions while still allowing pre-computation of document embeddings.

7. Hybrid Search

7.1 Combining Sparse and Dense Retrieval

Hybrid search combines the complementary strengths of sparse (BM25) and dense (embedding) retrieval:

Sparse retrieval excels at:

• Exact keyword matching
• Rare/specialized terms
• Interpretability

Dense retrieval excels at:

• Semantic similarity
• Synonym/p paraphrase matching
• Conceptual relevance

7.2 Reciprocal Rank Fusion (RRF)

where is the set of retrieval systems and is a constant (typically 60).

7.3 Learned Hybrid Approaches

Instead of fixed , learn the fusion weights:

where are learned parameters.

8. Advanced RAG Techniques

8.1 Query Transformation

Query rewriting: Use an LLM to reformulate the query for better retrieval:

Query decomposition: Break complex queries into sub-queries:

HyDE (Hypothetical Document Embeddings): Generate a hypothetical answer, embed it, and use that embedding for retrieval:

8.2 Iterative RAG

Multiple rounds of retrieval and generation:

for t in range(max_iterations):
    d_t = retrieve(q, context_so_far)
    y_t = generate(q, d_t, context_so_far)
    if is_sufficient(y_t):
        return y_t
    context_so_far += d_t

8.3 Self-RAG

The model learns to self-reflect on whether retrieval is needed:

and generates special tokens to indicate when retrieval is required and whether retrieved passages are relevant.

9. Evaluation

9.1 Retrieval Metrics

Recall@K: Fraction of relevant documents in top-K:

Mean Reciprocal Rank (MRR):

Normalized Discounted Cumulative Gain (NDCG):

9.2 Generation Quality

Faithfulness: Does the answer accurately reflect the retrieved context?

Relevance: Does the answer address the user's question?

Answer Correctness: Is the answer factually correct?

9.3 RAG-Specific Metrics

Context Recall: Are all relevant documents retrieved?

Context Precision: Are retrieved documents relevant?

Faithfulness Score:

10. Production Considerations

10.1 Latency Budget

10.2 Cost Analysis

10.3 Scalability

For documents with -dimensional embeddings:

• Memory for embeddings:
• Requires distributed index (sharding across machines)
• HNSW index build time: hours to days

11. Implementation: End-to-End RAG Pipeline

from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings
from langchain.retrievers import BM25Retriever, EnsembleRetriever
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

class RAGPipeline:
    def __init__(self, documents, chunk_size=512):
        self.documents = self._chunk(documents, chunk_size)
        self.embeddings = OpenAIEmbeddings()

        # Dense retriever
        self.dense_store = FAISS.from_documents(
            self.documents, self.embeddings
        )
        self.dense_retriever = self.dense_store.as_retriever(
            search_kwargs={"k": 10}
        )

        # Sparse retriever
        self.sparse_retriever = BM25Retriever.from_documents(
            self.documents, k=10
        )

        # Hybrid retriever
        self.hybrid_retriever = EnsembleRetriever(
            retrievers=[self.dense_retriever, self.sparse_retriever],
            weights=[0.6, 0.4]
        )

        self.llm = ChatOpenAI(model="gpt-4", temperature=0)
        self.prompt = ChatPromptTemplate.from_template("""
        Context: {context}
        Question: {question}
        Answer based on the context above. If the context doesn't
        contain enough information, say so.
        """)

    def query(self, question):
        docs = self.hybrid_retriever.get_relevant_documents(question)
        context = "\n\n".join([d.page_content for d in docs])
        chain = self.prompt | self.llm
        return chain.invoke({"context": context, "question": question})

References

1. Lewis et al. (2020). "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks." NeurIPS.
2. Karpukhin et al. (2020). "Dense Passage Retrieval for Open-Domain Question Answering." EMNLP.
3. Khattab & Zaharia (2020). "ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT." SIGIR.
4. Gao et al. (2023). "Retrieval-Augmented Generation for Large Language Models: A Survey." arXiv.
5. Lewis et al. (2021). "REALM: Retrieval-Augmented Language Model Pre-Training." ICML.