Retrieval-Augmented Generation

A standard autoregressive language model defines a probability distribution $p_\theta(y \mid x)$ over an output sequence $y$ given an input $x$, with all world knowledge implicitly encoded in the parameters $\theta$. Retrieval-Augmented Generation instead defines generation as a two-stage process. First, a retriever $p_\eta(z \mid x)$ scores candidate documents $z$ from a corpus $\mathcal{Z}$ for relevance to $x$. Second, a generator $p_\theta(y \mid x, z)$ produces the output conditioned on both the query and the retrieved document(s). The retriever is typically a dense retrieval model: a bi-encoder that maps queries and documents into a shared embedding space, where relevance is approximated by inner product or cosine similarity, exactly the machinery covered under Dense retrieval and embeddings. The key modeling move is to treat the retrieved document $z$ as a latent variable and marginalize over it: ``` p(y | x) = Σ_z p_η(z | x) · p_θ(y | x, z) ``` Because summing over an entire corpus (millions to billions of documents) is intractable, this sum is approximated by restricting it to the top-$k$ documents returned by nearest-neighbor search — a controlled, tractable approximation whose quality depends on retrieval recall. Two architectural variants matter here. In RAG-Sequence, one document is used to generate the entire output sequence, and the final probability marginalizes complete sequences over the top-k documents. In RAG-Token, a different (weighted) mixture over documents can be used at every generated token, letting the model blend evidence from several documents within a single answer. Fusion-in-Decoder (FiD) is a related and now more common pattern: instead of marginalizing at the probability level, the decoder attends jointly over the encoded representations of all retrieved passages simultaneously, letting cross-attention do the evidence-combination implicitly. RAG is best understood as a retrieve-then-read generalization of open-domain question answering, promoted from a task-specific pipeline into a general-purpose generation mechanism that can be trained end to end or used entirely at inference time with a frozen retriever and frozen or fine-tuned generator.

The system has two components trained (or configured) somewhat independently, and one inference-time procedure that ties them together.

Indexing (offline, done once per corpus update). Every document in the corpus is encoded into a dense vector using a document encoder $E_D$, and these vectors are stored in an approximate nearest-neighbor (ANN) index (e.g., HNSW, IVF-PQ, or a similar structure). This is the expensive, embarrassingly parallel, one-time (or periodically refreshed) cost.

Retrieval (online, per query). The query is encoded with a query encoder $E_Q$ (often sharing weights or architecture with $E_D$, but not always identical), and the ANN index returns the top-$k$ nearest document vectors by similarity score.

Generation (online, per query). The generator — a sequence-to-sequence or decoder-only LLM — receives the query concatenated with (or cross-attending to) the retrieved passages, and produces the answer. If citations are required, retrieved passage identifiers are threaded through so the final answer can point back to its evidence. ``` function RAG_ANSWER(query x, corpus index I, k): q_vec = ENCODE_QUERY(x) # dense embedding, O(d) docs = ANN_SEARCH(I, q_vec, k) # top-k passages, see complexity below scores = SOFTMAX([SIM(q_vec, ENCODE_DOC(d)) for d in docs])

# --- Fusion-in-Decoder style: encode each (query, doc) pair independently --- encoded_reps = [] for d in docs: rep = ENCODER(CONCAT(x, d)) # each passage encoded separately encoded_reps.append(rep) all_reps = CONCAT(encoded_reps) # decoder attends over all of them jointly y = DECODER_GENERATE(all_reps) # single joint decode, cross-attends to all docs return y, docs # return answer AND supporting evidence function RAG_SEQUENCE_MARGINAL(query x, docs, generator):

# Alternative scoring: marginalize whole-sequence probability over documents total_prob = 0 for (d, p_retrieve) in docs_with_scores: p_gen = PROBABILITY(generator, y | x, d) # probability of a candidate y given this doc total_prob += p_retrieve p_gen return total_prob ``` Training has two common regimes: 1. <strong>Frozen retriever, fine-tuned or frozen generator (retrieval-augmented inference).</strong> The retriever and its index are built once (often reusing an existing dense retriever); the generator, which may be a large pretrained LLM, is used with in-context retrieved passages via prompting, requiring no additional training. This is the dominant pattern in production systems today because it needs no gradient updates to the retriever. 2. <strong>End-to-end trained RAG.</strong> The generator's loss is backpropagated through the retrieval scores $p_\eta(z\mid x)$ as well, nudging the query encoder to retrieve documents that make the correct answer more likely. Because the ANN search itself (the discrete top-k selection) is non-differentiable, the standard trick is to keep the retrieved set fixed within a training step (or refresh the index periodically) and only backpropagate through the continuous* relevance scores of that fixed set, treating the reranking as differentiable and the retrieval-and-index-lookup as an occasionally-refreshed subroutine.

Time complexity of retrieval. With an exact brute-force search over $N$ documents each of embedding dimension $d$, retrieval costs $O(N \cdot d)$ per query — linear in corpus size, which is prohibitive for $N$ in the billions. Approximate nearest-neighbor structures (HNSW, IVF, product quantization) reduce this to roughly $O(\log N \cdot d)$ or $O(\sqrt{N}\cdot d)$ depending on the structure, at the cost of recall: the true nearest neighbor is found with high probability but not certainty. This is the standard sublinear-search tradeoff also central to Dense retrieval and embeddings: you trade a small, tunable probability of missing the best document for orders-of-magnitude speedup.

Time complexity of generation. Given $k$ retrieved passages each of length $L$ tokens, a Fusion-in-Decoder style encoder processes each passage independently in $O(k \cdot L^2)$ time (quadratic self-attention per passage, but not quadratic across passages, since they're encoded separately before the decoder fuses them) — this is why FiD scales better with $k$ than naively concatenating all $k$ passages into one long context, which would cost $O((kL)^2)$ under standard full self-attention. The decoder's cross-attention over the concatenated encoder outputs costs $O(T \cdot kL)$ for an output of length $T$. Total generation cost is therefore roughly

compared to $O((kL)^2 + T\cdot kL)$ for naive concatenation — a substantial saving when $k$ is large, since encoding cost grows linearly rather than quadratically in $k$.

Space complexity. The index stores $N$ vectors of dimension $d$, i.e., $O(N d)$ floats (often compressed via product quantization to a fraction of a byte per dimension), plus graph or inverted-list overhead for the ANN structure, typically $O(N)$ additional pointers.

Correctness argument. The probabilistic justification for RAG rests on the law of total probability. If $z$ is genuinely the (unknown) supporting document and the corpus $\mathcal{Z}$ is exhaustive, then

is exact by definition of marginalization — no approximation is needed if the sum runs over the whole corpus and the retriever's distribution $p_\eta(z\mid x)$ is the true posterior over which document is relevant. The approximation error introduced by RAG comes from two truncations: (1) restricting the sum to the top-$k$ retrieved documents (valid as an approximation because $p_\eta(z\mid x)$ is typically sharply peaked — a well-trained dense retriever assigns near-zero mass to irrelevant documents, so the truncated sum captures nearly all the probability mass); and (2) the retriever itself being an imperfect estimate of true relevance (an empirical error, bounded by retrieval recall@k, not a mathematical one). Thus RAG is exact under an idealized retriever and approximately correct in proportion to how close the real retriever's top-k recall is to 1. This gives a precise, falsifiable criterion for diagnosing failures: if the answer is wrong, either the retriever failed to place the needed document in the top-k (a recall failure, addressable by better retrieval or larger $k$), or the generator ignored/misused a correctly retrieved document (a grounding failure, addressable by better fine-tuning or prompting).

Honest tradeoffs. RAG adds a hard dependency on retrieval quality: a system with a superb generator and a weak retriever will still hallucinate, just now while sounding like it's citing evidence — arguably worse, since the citation lends false credibility. RAG also adds latency (an ANN lookup plus longer, evidence-laden context) and cannot handle queries whose answer requires reasoning not present verbatim in any single document (e.g., multi-hop arithmetic over scattered facts) without additional retrieval iterations or Tool Use & Function Calling.

Consider a customer-support RAG system over a small internal corpus, queried with: > "What is our refund policy for items damaged in shipping?" Suppose the indexed corpus contains four passages:

With $k = 2$, the retriever returns D1 and D2 (highest similarity). D1 is directly on-topic; D2 is tangentially related (returns policy generally, not damage-specific). The retrieval scores are converted to a normalized distribution via softmax over the top-k similarities: ``` softmax([0.81, 0.52]) ≈ [0.57, 0.43] (temperature-scaled example) ``` The generator, using Fusion-in-Decoder, encodes `(query, D1)` and `(query, D2)` independently, then the decoder cross-attends over both encoded representations while generating. Because D1's content directly answers the question and D2's does not, a well-trained generator learns (via its cross-attention weights, shaped by training signal) to lean heavily on D1's evidence despite D2 receiving nontrivial retrieval mass. The output: > "Damaged items are eligible for a refund if you submit photo evidence within 14 days of delivery. [Source: D1]" Trace of what would go wrong at each failure point, useful for debugging:

A second, contrasting case: query "What is our refund policy for lost luggage on international flights?" — no passage in the corpus addresses this. A well-behaved RAG system should recognize low maximum similarity (e.g., top score 0.31, well below the 0.81 seen above) and abstain or say "the available documents do not cover this," rather than force an answer from a low-relevance passage. This retrieval-confidence-aware abstention is precisely the mechanism connecting RAG to broader work on Hallucination & Factuality.

RAG synthesizes decades of information retrieval research with the sequence-to-sequence generation paradigm that matured through neural machine translation and pretrained transformers. Its immediate lineage runs through open-domain question answering (retrieve-then-read pipelines using classical sparse retrieval like TF-IDF/BM25), through the shift to learned dense retrieval, and finally into a single differentiable architecture combining both.

RAG underlies enterprise "chat with your documents" products, where a company's internal wikis, contracts, or support tickets are indexed and queried by an LLM interface without any model retraining. Search engines that generate synthesized answers with citations (rather than only a ranked link list) use retrieval-augmented pipelines to ground responses in live web content. Legal and medical assistant tools rely on RAG to point to source statutes or clinical guidelines, since an unsupported claim in these domains carries real liability. Coding assistants use retrieval over a codebase or documentation set so suggestions reflect the actual project rather than generic patterns memorized from public code. RAG is also a natural precursor to agentic systems that interleave retrieval with other actions, connecting directly to Tool Use & Function Calling, where retrieval becomes just one callable tool among many (calculators, APIs, code execution) that a model can invoke mid-generation. Because retrieval quality directly bounds output faithfulness, RAG is also the primary architectural lever discussed in work on Hallucination & Factuality, and its prompt-construction step (how retrieved text is formatted and inserted into context) draws directly on techniques from Prompting & Prompt Engineering.

Ask the tutor: How would you redesign a RAG pipeline to handle a question that requires combining facts from three separate documents through multi-step reasoning?