RAG: Retrieval-Augmented Generation

Chapter 1: Dense Passage Retrieval

Before we can build RAG, we need a way to find relevant documents. Given a question like "What causes aurora borealis?", we need to search through millions of Wikipedia passages and return the handful that contain the answer. This is the job of the retriever.

Traditional search engines use sparse retrieval — methods like TF-IDF or BM25 that match exact keywords. If your question says "aurora" and a passage says "aurora," it's a match. This works well for simple queries but fails badly when the question and answer use different words. "What causes the Northern Lights?" should match a passage about "solar wind interacting with Earth's magnetosphere" — but there are no shared keywords.

RAG uses Dense Passage Retrieval (DPR), published by Karpukhin et al. just months before the RAG paper. Instead of matching keywords, DPR encodes both the question and every passage into dense vector embeddings, then finds passages whose vectors are closest to the question vector.

How DPR works

DPR uses two separate BERT encoders:

The two encoders are trained on pairs of (question, relevant passage) using a contrastive loss: push the question vector close to the relevant passage vector and far from irrelevant passage vectors.

Where p⁺ is the correct passage and p_i⁻ are negative passages (wrong answers). This is the standard InfoNCE loss — the same contrastive loss used in CLIP and SimCLR. It says: make the correct (question, passage) pair have a higher dot product than any incorrect pair.

The MIPS trick: searching millions of passages in milliseconds

At inference time, we need to find the top-k passages out of ~21 million Wikipedia passages. A brute-force search would compute 21 million dot products — too slow. DPR uses Maximum Inner Product Search (MIPS) with Facebook's FAISS library, which builds a compressed index that enables approximate nearest-neighbor search in milliseconds.

python
# Building a DPR index with FAISS
import faiss
import numpy as np

# Pre-compute passage embeddings (done ONCE, offline)
passage_embeddings = encode_all_passages(wikipedia)  # shape: [21M, 768]

# Build FAISS index for fast approximate search
index = faiss.IndexFlatIP(768)          # inner product index
index.add(passage_embeddings)              # add all 21M vectors

# At inference: encode question, search index
q_vec = question_encoder("What causes aurora?")  # [1, 768]
scores, indices = index.search(q_vec, k=5)    # top-5 passages
# scores: [1, 5] — similarity scores
# indices: [1, 5] — passage IDs to look up text

Hard negatives matter

The quality of DPR depends heavily on how you choose negative examples during training. Random negatives (random passages from the corpus) are too easy — the model learns to distinguish by topic rather than by relevance. Hard negatives are passages that are topically related but don't answer the question. Karpukhin et al. found that using BM25-retrieved passages as hard negatives (high keyword overlap but not the answer) dramatically improved retrieval quality.

python
# DPR training with hard negatives
for question, positive_passage in training_data:
    # Easy negative: random passage from corpus
    easy_neg = random.choice(corpus)

    # Hard negative: BM25-retrieved but NOT the answer
    bm25_results = bm25.search(question, k=100)
    hard_neg = [p for p in bm25_results if p != positive_passage][0]

    # Contrastive loss: push q closer to pos, farther from negs
    loss = nce_loss(
        q=encode_q(question),
        pos=encode_p(positive_passage),
        negs=[encode_p(easy_neg), encode_p(hard_neg)]
    )

This detail matters for RAG because the retriever quality directly determines the generator's input quality. Garbage in, garbage out — if the retriever returns irrelevant passages, even a perfect generator can't produce correct answers.

Chapter 2: The RAG Architecture

Now we have the two pieces: a retriever that finds relevant documents and a generator that produces text. RAG combines them into a single end-to-end model. The architecture is elegant: query in, answer out, with document retrieval happening in the middle.

The two components

Retriever: p_η(z | x) — Given an input query x, the retriever returns a distribution over documents z. In practice, it retrieves the top-k documents (k=5 or k=10) from a FAISS index using DPR. The parameter η refers to the BERT_Q question encoder weights (the passage encoder is frozen and pre-computed).

Generator: p_θ(y_i | x, z, y_1:i-1) — Given the input x, a retrieved document z, and previously generated tokens y_1:i-1, the generator predicts the next token y_i. This is BART — a pre-trained seq2seq model. The encoder processes [x; z] (query concatenated with retrieved document), and the decoder generates the answer autoregressively.

The marginalization step

This is the mathematical heart of RAG. We don't just pick the best document and generate from it. Instead, we treat the retrieved documents as latent variables and marginalize them out. The probability of generating output y given input x is:

Where p_η(z | x) is the retriever's score for document z (how relevant it is to the query) and p_θ(y | x, z) is the generator's probability of producing answer y given document z.

In plain English: run the generator k times, once per retrieved document. Each run produces a probability for the answer. Weight each probability by how relevant the document was. Sum them up. Documents that are more relevant contribute more to the final answer.

Tensor shapes through the pipeline

python
# RAG forward pass — exact tensor shapes
x = tokenizer("What year was the Eiffel Tower completed?")
# x.input_ids: [1, seq_len]  (batch=1, sequence length)

# Step 1: Encode query
q = question_encoder(x.input_ids)     # [1, 768]

# Step 2: Retrieve top-k documents
scores, doc_ids = faiss_index.search(q, k=5)
# scores: [1, 5]   doc_ids: [1, 5]
doc_texts = [passage_db[i] for i in doc_ids[0]]

# Step 3: Concatenate query with each document
inputs = [x_text + " [SEP] " + doc for doc in doc_texts]
enc_ids = tokenizer(inputs, padding=True)
# enc_ids.input_ids: [5, max_doc_len]  — 5 docs, padded

# Step 4: Run BART encoder + decoder for each document
gen_scores = bart.generate(enc_ids, return_scores=True)
# gen_scores: [5, output_len, vocab_size]

# Step 5: Marginalize over documents
doc_priors = softmax(scores)              # [1, 5] — retrieval weights
final_probs = (doc_priors @ gen_scores)   # weighted sum over docs

Chapter 3: RAG-Token vs RAG-Sequence

The formula we saw in Chapter 2 — marginalizing over documents — has a subtle but critical ambiguity. When do we marginalize? This choice creates two distinct variants of RAG, each with different properties.

RAG-Sequence: marginalize per complete sequence

In RAG-Sequence, we generate the entire output sequence conditioned on each document separately, then marginalize over documents once at the end:

In plain English: for each document z, generate the complete answer using only that document. Then weight each complete answer by how relevant z was. This means each answer candidate is "faithful" to a single document — the model never mixes information across documents within a single generation.

RAG-Token: marginalize per token

In RAG-Token, we marginalize at each token position separately. Every token in the output can attend to a different document:

Notice the difference: in RAG-Sequence, the sum over z is outside the product over tokens. In RAG-Token, the sum over z is inside the product. This means each output token can draw from a different document.

Decoding details

For RAG-Token, decoding is straightforward. At each step, compute the per-token probability from each document, take the weighted sum, and feed it into standard beam search.

For RAG-Sequence, it's trickier. We run beam search separately for each of the k documents, generating k sets of candidate sequences. Then we need to score each unique candidate across all documents using the full marginalization formula. Some candidates might only appear in one document's beam, in which case we need to run an additional forward pass with the other documents to get their probabilities (or approximate them as zero).

python
# RAG-Sequence decoding (simplified)
candidates = {}
for z_i, score_i in zip(top_k_docs, retrieval_scores):
    # Run beam search conditioned on this single document
    beams = bart.beam_search(query, z_i, num_beams=5)
    for seq, log_prob in beams:
        if seq not in candidates:
            candidates[seq] = []
        candidates[seq].append((score_i, log_prob))

# Marginalize: for each unique candidate, sum p(z|x) * p(y|x,z)
final_scores = {}
for seq, contributions in candidates.items():
    final_scores[seq] = sum(s * exp(lp) for s, lp in contributions)

best_answer = max(final_scores, key=final_scores.get)

Chapter 4: Training RAG

Training RAG is tricky because we have two components — a retriever and a generator — and they need to learn together. The retriever needs to learn what documents are useful for the generator, and the generator needs to learn how to use retrieved documents to produce correct answers.

What gets trained

RAG has three sets of parameters:

The training objective

We train to maximize the marginal log-likelihood of the correct answer:

Where D is the training set of (question, answer) pairs. We don't have supervision for which documents should be retrieved — the documents are latent. The model must learn which documents are useful entirely from the end-to-end signal of producing correct answers.

Gradient flow

How do gradients flow through retrieval? The retriever scores p_η(z | x) are differentiable with respect to η (the question encoder weights). The FAISS lookup itself is non-differentiable (it's a discrete nearest-neighbor search), but we approximate the gradient by fixing the set of retrieved documents during each forward pass and only backpropagating through the scores.

python
# RAG training loop (simplified)
for x, y in dataloader:
    # 1. Encode question (DIFFERENTIABLE w.r.t. eta)
    q = question_encoder(x)                # [B, 768]

    # 2. Retrieve top-k documents (NOT differentiable)
    scores, doc_ids = faiss_index.search(q, k=5)

    # 3. Re-score with differentiable dot product
    doc_embeds = passage_embeds[doc_ids]     # [B, 5, 768] — frozen
    retriever_scores = (q.unsqueeze(1) @ doc_embeds.transpose(-1,-2)).squeeze()
    # [B, 5] — these ARE differentiable w.r.t. q (and thus eta)

    doc_priors = softmax(retriever_scores)  # [B, 5]

    # 4. Generate with BART for each document
    gen_log_probs = []
    for i in range(5):
        bart_input = concat(x, docs[i])
        lp = bart.log_prob(bart_input, y)  # log p(y | x, z_i)
        gen_log_probs.append(lp)

    # 5. Marginalize and compute loss
    gen_probs = torch.stack(gen_log_probs, dim=1).exp()  # [B, 5]
    marginal = (doc_priors * gen_probs).sum(dim=1)          # [B]
    loss = -marginal.log().mean()

    loss.backward()   # gradients flow to BOTH eta and theta
    optimizer.step()

The REINFORCE-like gradient

Let's unpack the gradient math one more step. The loss for a single example is:

For the generator parameters θ, the gradient is straightforward — it flows through p_θ(y | x, z) for each document. For the retriever parameters η, the gradient has an interesting structure:

Where w(z) is a weight proportional to how much document z helped produce the correct answer. This looks like a REINFORCE gradient — the retriever gets a "reward signal" for retrieving documents that led to correct answers. Documents that helped the generator are reinforced; irrelevant documents are downweighted.

Training hyperparameters

A critical detail: the FAISS index is not updated during training. This means the set of retrievable documents stays fixed at whatever was indexed from the DPR passage encoder. The question encoder shifts to produce better queries for the fixed index, but the index itself doesn't improve. REALM (a concurrent paper by Guu et al.) tackled this by periodically re-encoding the entire index — asynchronously, in the background — but this added significant engineering complexity.

Chapter 5: Open-Domain QA Results

The headline evaluation for RAG is open-domain question answering — answering factual questions using only a corpus of documents (Wikipedia) as the knowledge source, with no task-specific architecture or pre-defined answer set.

Lewis et al. evaluated on four major QA benchmarks:

Beyond QA: knowledge-intensive generation

RAG's real advantage shows on tasks that require generating text, not just extracting spans:

Jeopardy Question Generation: Given an entity (e.g., "Eiffel Tower"), generate a Jeopardy-style clue. RAG produced more factual, specific, and diverse clues than BART alone. Human evaluators preferred RAG's outputs on factuality and specificity.

FEVER Fact Verification: Given a claim, classify it as SUPPORTED, REFUTED, or NOT ENOUGH INFO using retrieved evidence. RAG achieved 72.5% accuracy, within 2.7 points of pipeline systems with much more complex architectures.

MS-MARCO (abstractive QA): Generate free-form answers to real Bing search queries. RAG produced more specific and factual answers than BART, though this benchmark wasn't the primary focus.

Why RAG outperforms "retrieve then read"

The standard pipeline approach before RAG was: (1) retrieve relevant passages, (2) feed them to a reader model, (3) extract an answer span. RAG improves on this in three ways:

Chapter 6: Knowledge Updates

Here is perhaps the most practically important property of RAG: you can update what the model "knows" without retraining it. Just swap the document index.

Consider a traditional language model trained in January 2020. By March 2020, COVID-19 has changed the world, but the model still thinks it's a minor outbreak in Wuhan. To update the model, you'd need to retrain it on new data — an expensive process requiring massive compute.

With RAG, you simply update the Wikipedia dump in the FAISS index. The model's weights stay the same. The next time someone asks about COVID-19, the retriever finds the new passages, and the generator produces an up-to-date answer.

What updating the index looks like

python
# Updating RAG's knowledge — no model retraining needed

# 1. Get new documents
new_articles = fetch_latest_wikipedia_dump()  # or any corpus

# 2. Chunk into 100-word passages
new_passages = chunk_documents(new_articles, chunk_size=100)

# 3. Encode with frozen passage encoder (BERT_P)
new_embeddings = passage_encoder.encode_batch(new_passages)
# shape: [num_new_passages, 768]

# 4. Replace FAISS index
new_index = faiss.IndexFlatIP(768)
new_index.add(new_embeddings)

# 5. Swap in the new index — zero downtime
rag_model.retriever.index = new_index

# Done! Model now answers with updated knowledge.
# No gradient updates. No retraining. No GPU hours.

Limitations of index swapping

Knowledge update via index swapping isn't a silver bullet. There are real limitations:

Parametric vs non-parametric knowledge

RAG formalized a distinction that has become fundamental to how we think about LLM systems:

This separation is the single most important conceptual contribution of the RAG paper. The specific architecture (DPR + BART) matters less than the pattern: let the model handle language, let the index handle knowledge.

Chapter 7: Connections

RAG sits at a critical junction in the history of NLP — it formalized the retrieve-then-generate pattern that became the backbone of modern LLM applications. Let's map where it came from and where it led.

Lineage

What RAG got right

Separation of concerns. The most lasting contribution: parametric knowledge (model weights) vs non-parametric knowledge (document index) should be separate systems. This insight survived every subsequent architectural change.

End-to-end training. Training the retriever and generator jointly, with the retriever learning from the end-task signal without retrieval labels. This is cleaner than pipeline approaches where the retriever is trained on a proxy task.

Two marginalization variants. RAG-Token and RAG-Sequence provided a principled framework for how to combine information from multiple documents. This distinction influenced subsequent work on multi-document reasoning.

What changed since RAG

Long-context models. GPT-4 with 128K context, Gemini with 1M+ tokens. When you can paste 50 documents into the context, the marginalization math becomes less important — you just concatenate everything and let attention handle it.

Retrieval as tool use. Modern systems treat retrieval as a tool the model can choose to invoke, not a mandatory step. The model decides when it needs external knowledge.

Chunking strategies. RAG used fixed 100-word chunks. Modern systems use semantic chunking, recursive splitting, and hierarchical retrieval. Chunk quality matters enormously.

Re-ranking. Modern RAG pipelines add a cross-encoder re-ranker between retrieval and generation: retrieve 100 candidates with a fast bi-encoder, re-rank to top-10 with a slower but more accurate cross-encoder.

The production RAG stack (2024)

For context, here's what a modern production RAG system looks like compared to the 2020 paper: