Retrieval-augmented generation

Problem

A decoder-only model stores what it learned in pretraining in its weights. That knowledge stops at a cutoff date, never included your private documents, and comes back without a source, so a confident wrong answer looks the same as a right one. Retraining or fine-tuning every time a document changes is slow and expensive, and it still does not give citations.

Retrieval-augmented generation (RAG) splits the job in two. A retriever finds the few passages in a corpus most likely to answer the question; a generator reads the question together with those passages and writes the answer. Updating knowledge becomes updating an index, and every answer can point at the passages it used. The cost is a second system that can fail on its own: if retrieval misses the passage, the generator cannot use it.

Intuition

It is an open-book exam. A student who has to answer from memory guesses when unsure; a student allowed to bring the textbook does better, but only if they can find the right page quickly, and only if they read the page instead of writing what they already believed. RAG's quality is the product of those two skills. Searching well is a similarity search problem, and reading well is the generator's job; most of the engineering goes into the first.

Mechanics

Chunking

Retrieval returns units, and the unit is the chunk. Chunks must be small enough that a vector or a term count can represent one topic and that several fit in the prompt, and large enough to hold a complete answer with its context. The common recipes:

• Fixed-size windows. Every n words or tokens. Simple and predictable, but a boundary can cut a sentence, or an answer, in half.
• Overlapping windows. Each window starts n − o words after the previous one, so any span up to o words long lies wholly inside some chunk. The cost is a bigger index and duplicate text when neighbouring chunks are both retrieved.
• Structure-aware splits. Pack whole sentences, paragraphs or sections up to a budget, so boundaries fall where the author put them. Headings and the document title are often prepended to each chunk so it still makes sense out of context.

The RAG paper reused the DPR index: Wikipedia split into disjoint 100-word passages, about 21 million of them. Chunk size is a hyperparameter to tune against retrieval metrics on your own questions, not a constant.

Dense, lexical and hybrid retrieval

A dense retriever is a two-tower model: a query encoder q(·) and a document encoder d(·) map text to vectors, and relevance is their inner product. The RAG paper uses DPR's BERT encoders and writes the retrieval distribution as

p_η(z | x) ∝ exp(d(z)^⊤ q(x))

Document vectors are computed once and stored in a vector index, so a query costs one encoder pass plus an approximate maximum-inner-product search; with unit-length vectors this is cosine similarity. Dense retrieval matches meaning ("charging" finds "charged"), but it can miss rare exact strings such as part numbers, names and error codes, and it only knows what its training pairs taught it.

A lexical retriever, usually BM25, scores exact term overlap weighted by rarity (see similarity search). It handles exact strings and needs no training, but it scores zero when the question and the passage use different words for the same thing. Hybrid retrieval runs both and merges the rankings. Their scores are on unrelated scales, so the simplest robust merge ignores scores and uses ranks, reciprocal rank fusion:

RRF(d) = Σ_{rankings r} 1 / (k + rank_r(d)), k = 60

Re-ranking

First-stage retrievers encode query and document separately, which is what makes them fast and indexable. A re-ranker sees them together. A cross-encoder (monoBERT is the classic) feeds "[query] [SEP] [passage]" through one transformer and outputs a relevance score, so attention can compare every query token with every passage token. It is far more accurate and far too slow to run over the corpus, so it only reorders the first stage's top 20 to 100 candidates. ColBERT sits in between: it stores one vector per document token and scores with "MaxSim", the sum over query tokens of the best-matching document token similarity, which keeps documents indexable while matching at the token level.

Re-ranking cannot recover a passage the first stage missed. That is why first-stage recall is measured on its own: it is the ceiling for everything after it.

Prompt assembly

The retrieved chunks, the question and the instructions must fit in the context window, with room left for the answer. Assembly decides how many chunks, in what order, and how they are marked up: each chunk with an id and source so the model can cite it, a clear separation between instructions and retrieved text, and an instruction to say so when the passages do not contain the answer.

More is not always better. Liu et al. found that models used information at the start or end of a long context much better than information in the middle ("lost in the middle"), and irrelevant passages distract the generator. Common responses are a small k after a good re-ranker, a relevance-score cutoff, and placing the strongest chunks at the edges of the context.

Combining documents: RAG-Sequence and RAG-Token

Putting chunks in the prompt treats retrieval as a fixed preprocessing step. The RAG paper instead treats the retrieved document z as a latent variable and marginalizes over the top k. RAG-Sequence uses one document for the whole answer:

p(y | x) ≈ Σ_{z ∈ top-k} p_η(z | x) Π_i p_θ(y_i | x, z, y_<i)

RAG-Token may draw each token from a different document:

p(y | x) ≈ Π_i Σ_{z ∈ top-k} p_η(z | x) p_θ(y_i | x, z, y_<i)

Both are differentiable in the retrieval scores, so the loss on the answer trains the query encoder without any labels saying which document was relevant. The paper keeps the document encoder, and so the index, fixed and fine-tunes the query encoder and the BART generator. Fusion-in-Decoder takes a third route: encode each (question, passage) pair separately and let the decoder attend over all of them at once.

Evaluating retrieval and generation separately

An end-to-end score says the system is wrong but not where. Evaluate each half against its own labels:

• Retrieval. For each question, label which chunks contain the answer. Recall@k is the fraction of questions with a relevant chunk in the top k; MRR averages 1/rank of the first relevant chunk (0 when none is found); nDCG handles graded relevance (see learning to rank).
• Generation given good context. Feed the labelled relevant chunks directly and score the answer. If answers are wrong even then, the generator or prompt is the problem, not retrieval.
• Answer quality. For short answers, exact match and token F1 after normalization (lowercase, strip punctuation and articles), as in SQuAD. For long answers, human or model judges, and two separate questions: is it correct, and is it faithful, meaning every claim is supported by the retrieved chunks.

Failure modes

Prompt injection deserves its own line: Greshake et al. showed that text planted in a web page or document a RAG system retrieves can steer the model, because the model sees retrieved text and instructions in the same context. Retrieval turns everyone who can write to your corpus into someone who can write to your prompt.

Worked example

The corpus is one invented product manual of 121 words, for the "Trailhead 2" headlamp, and five questions, each with the phrase that answers it. The retrievers are stand-ins small enough to compute by hand, labelled for what they are: BM25 is real BM25 (k₁ = 1.2, b = 0.75); the "vector" retriever uses character-trigram count vectors and cosine similarity, a crude substitute for a learned encoder that still matches "charging" with "charged"; the re-ranker is a hand-built scorer, not a trained cross-encoder. Questions drop a short stop-word list before scoring.

The rest of the example uses the 5 sentence chunks:

Implementation

import torch
import torch.nn.functional as F

def build_index(encode, chunks):                      # encode: list[str] -> (N, d)
    with torch.no_grad():
        return F.normalize(encode(chunks), dim=-1)

def dense_search(encode, index, queries, k=20):
    with torch.no_grad():
        q = F.normalize(encode(queries), dim=-1)      # (B, d)
    scores = q @ index.T                              # cosine, (B, N)
    return torch.topk(scores, k, dim=-1)              # values, indices: (B, k)

def rrf(rankings, k=60):                              # rankings: lists of chunk ids
    fused = {}
    for ranking in rankings:
        for rank, chunk_id in enumerate(ranking, start=1):
            fused[chunk_id] = fused.get(chunk_id, 0.0) + 1 / (k + rank)
    return sorted(fused, key=fused.get, reverse=True)

def rerank(cross_encoder, query, chunks, candidates, keep=5):
    with torch.no_grad():
        scores = cross_encoder([(query, chunks[i]) for i in candidates])   # (C,)
    order = torch.argsort(scores, descending=True)[:keep]
    return [candidates[i] for i in order.tolist()], scores[order]

def assemble(question, chunks, ids, count_tokens, budget, min_score=None, scores=None):
    parts, used = [], count_tokens(question) + 200   # room for instructions
    for j, i in enumerate(ids):
        if min_score is not None and scores[j] < min_score:
            break
        block = f"[{i}] {chunks[i]}"
        if used + count_tokens(block) <= budget:
            parts.append(block)
            used += count_tokens(block)
    context = "\n\n".join(parts)
    return ("Answer using only the numbered passages and cite them as [n]. "
            "If they do not contain the answer, say so.\n\n"
            f"Passages:\n{context}\n\nQuestion: {question}\nAnswer:")

def rag_log_likelihoods(doc_scores, token_logp):     # (K,), (K, T)
    log_pz = F.log_softmax(doc_scores, dim=-1)
    sequence = torch.logsumexp(log_pz + token_logp.sum(-1), dim=0)
    token = torch.logsumexp(log_pz[:, None] + token_logp, dim=0).sum()
    return sequence, token

Brute-force q @ index.T is exact, and its cost grows linearly with the corpus; once that is too slow, use an approximate index such as HNSW or IVF-PQ from the vector indexes page, and measure recall against the exact search. Use the same encoder, normalization and text preprocessing for queries and documents, and re-embed the whole corpus when the encoder changes: vectors from two models are not comparable. Store each chunk's source, position and version with its vector so answers can cite and filter. Build the evaluation set early, a few hundred real questions with labelled relevant chunks, and tune chunk size, k, the fusion and the re-ranker against recall@k and MRR before touching the prompt.

Tradeoffs

Connections

The dense retriever is a two-tower model served from a vector index; BM25 and hybrid fusion are developed on similarity search, and recall, MRR and nDCG on learning to rank. The generator is a decoder-only model, and the alternative way to add knowledge is fine-tuning. RAG's marginal over documents is a softmax-weighted mixture, computed in log space with the log-softmax and logsumexp used on the loss functions page.

References

• Lewis et al. Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS 2020. arXiv:2005.11401
• Karpukhin et al. Dense Passage Retrieval for Open-Domain Question Answering. EMNLP 2020. arXiv:2004.04906
• Guu et al. REALM: Retrieval-Augmented Language Model Pre-Training. ICML 2020. arXiv:2002.08909
• Izacard and Grave. Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering. EACL 2021. arXiv:2007.01282
• Nogueira and Cho. Passage Re-ranking with BERT. 2019. arXiv:1901.04085
• Khattab and Zaharia. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. SIGIR 2020. arXiv:2004.12832
• Cormack, Clarke and Büttcher. Reciprocal Rank Fusion Outperforms Condorcet and Individual Rank Learning Methods. SIGIR 2009. doi:10.1145/1571941.1572114
• Robertson and Zaragoza. The Probabilistic Relevance Framework: BM25 and Beyond. Foundations and Trends in IR 2009. doi:10.1561/1500000019
• Rajpurkar et al. SQuAD: 100,000+ Questions for Machine Comprehension of Text. EMNLP 2016. arXiv:1606.05250
• Liu et al. Lost in the Middle: How Language Models Use Long Contexts. TACL 2024. arXiv:2307.03172
• Greshake et al. Not What You've Signed Up For: Compromising Real-World LLM-Integrated Applications with Indirect Prompt Injection. 2023. arXiv:2302.12173
• PyTorch v2.15.0: F.normalize, torch.matmul, torch.topk, torch.argsort, F.cosine_similarity, F.log_softmax, torch.logsumexp.