Retrieval-Augmented Generation: Combining LLMs with Knowledge

In 2023, lawyers in multiple jurisdictions — including in India — were pulled up by courts for filing briefs that cited case law which simply did not exist. The pattern was always the same: someone asked an LLM to recall a precedent from memory, the model produced a fluent, correctly formatted citation complete with a judge's name and a year, and nobody checked it against an actual law database before it reached a courtroom. The model was not being deceptive. It was doing exactly what a language model trained to predict plausible next tokens is built to do — and "a citation that looks like the citations in my training data" is a very different target from "a citation that refers to a real judgment." If you were building the retrieval layer behind an Indian legal-research assistant — the kind of tool that has to search millions of AIR and SCC-reported judgments before it lets a model answer — this is exactly the failure you are hired to prevent. Retrieval-Augmented Generation (RAG) is the architecture that prevents it, and this chapter goes past the "retrieve some text and paste it into the prompt" mental model into the actual mechanism: the probability model the retriever and generator jointly define, how that retrieval is made fast enough to run at production scale over tens of millions of documents, and where the architecture's own math creates a failure mode students consistently miss.

Why "paste retrieved text into the prompt" is not the whole story

Every RAG system has two trainable (or at least two distinct) components. A retriever takes a query x and returns a small set of candidate documents z from a corpus, each with a relevance score. A generator is a sequence-to-sequence or decoder-only language model that produces the answer y conditioned on both the query and whichever document it was handed. The naive way to combine them — the way most first tutorials present it — is to retrieve the single best-scoring document and concatenate it into the prompt as extra context. That works, but it throws away information: the retriever's confidence that the document is actually relevant never enters the generator's probability at all. A document that scored barely above the retrieval threshold is treated with exactly the same weight as one that scored far above every competitor.

Patrick Lewis and colleagues at Facebook AI Research formalized a better alternative in their 2020 NeurIPS paper, Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. Instead of picking one document, they treat the retrieved document z as a latent variable and marginalize over it — the retriever's relevance score becomes a literal probability weight inside the generator's output distribution, not just a filter on what text gets shown. Their retriever was Dense Passage Retrieval (DPR), from Karpukhin et al., EMNLP 2020: two independent BERT-base encoders, a query encoder q(·) and a passage encoder d(·), each mapping text to a 768-dimensional vector, trained so that a query's vector has high dot product with the vectors of passages that actually answer it (trained with in-batch negatives plus mined hard negatives — passages that look topically similar but don't answer the question, which is exactly what forces the encoder to learn semantic relevance rather than keyword overlap). The retrieval score is:

p_η(z | x) ∝ exp( d(z)ᵀ · q(x) )

a softmax over dot products, computed only over the top-k documents returned by an approximate nearest-neighbor search (more on why "approximate" is unavoidable in a moment). The generator in the original paper was BART-large, a 400-million-parameter encoder-decoder. Given that setup, the paper defines two mathematically distinct ways to combine retriever and generator, and the difference between them is the part most treatments skip.

RAG-Sequence and RAG-Token: two different marginalizations

RAG-Sequence treats one retrieved document as the sole source of evidence for the entire output sequence. It generates a full candidate answer conditioned on each retrieved document separately, then mixes the resulting sequence-level probabilities:

p_Sequence(y | x) = Σ_{z ∈ top-k} p_η(z | x) · p_θ(y | x, z)

where p_θ(y | x, z) = Π_i p_θ(y_i | x, z, y_1:i-1) is the generator's probability of the whole sequence given that one document. Every token of a given hypothesis is generated as if z were the only evidence in the room; documents are only ever compared at the level of the finished sequence.

RAG-Token marginalizes over documents separately at every generated token:

p_Token(y | x) = Π_{i=1}^{N} Σ_{z ∈ top-k} p_η(z | x) · p_θ(y_i | x, z, y_1:i-1)

Retrieval still happens once per query — the top-k set and the weights p_η(z|x) don't change mid-generation — but which document's *reading* of the evidence dominates can shift from token to token, because the generator's conditional p_θ(y_i | x, z, y_1:i-1) is recomputed fresh at each step for each document. In effect, RAG-Token behaves like a per-token mixture-of-experts over the retrieved documents, letting the model pull the subject from one passage and a supporting detail from another within the same answer. RAG-Sequence, precisely because it commits one document to the whole sequence before mixing, cannot do that — it is closer to "pick the best single witness and let it tell the whole story, then average across witnesses." These are not implementation details; they are different probability distributions over answers, and — as the worked example below shows — for anything longer than a single output token they generally assign different probabilities to the same candidate answer.

Worked example: three passages, one contested fact

Suppose the legal-research assistant's general-knowledge cousin is asked: "Who won the 2023 Cricket World Cup final?" (Australia beat India by six wickets at the Narendra Modi Stadium, Ahmedabad, on 19 November 2023.) The retriever's approximate search over the corpus returns three candidate passages, each already reduced by DPR's encoders to a toy 3-dimensional embedding for this illustration (a real DPR vector is 768-dimensional; the mechanism is identical):

• d_a — the correct passage, describing the actual 2023 final and its result: embedding (2.4, 1.0, 1.0)
• d_b — a passage about the 2019 final (England beat New Zealand at Lord's after a Super Over), topically close but temporally wrong: embedding (1.0, 1.4, −1.0)
• d_c — a passage about the IPL, cricket-related but not about the World Cup at all: embedding (0.2, 0.6, 0.5)

The query embedding is q(x) = (1.0, 0.5, −0.2). The retrieval scores are dot products:

score(d_a) = 1.0×2.4 + 0.5×1.0 + (−0.2)×1.0 = 2.4 + 0.5 − 0.2 = 2.7
score(d_b) = 1.0×1.0 + 0.5×1.4 + (−0.2)×(−1.0) = 1.0 + 0.7 + 0.2 = 1.9
score(d_c) = 1.0×0.2 + 0.5×0.6 + (−0.2)×0.5 = 0.2 + 0.3 − 0.1 = 0.4

Softmax normalizes these into the retrieval distribution p_η(z|x). Using e^2.7 ≈ 14.8797, e^1.9 ≈ 6.6859, e^0.4 ≈ 1.4918, the partition sum is Z ≈ 23.0575, giving:

p_η(d_a) ≈ 14.8797 / 23.0575 ≈ 0.6453
p_η(d_b) ≈  6.6859 / 23.0575 ≈ 0.2900
p_η(d_c) ≈  1.4918 / 23.0575 ≈ 0.0647

Now suppose the generator, reading each passage independently, assigns the following probability to producing the first output token as the word "Australia": p_θ("Australia" | x, d_a) = 0.90 (the passage says so directly), p_θ("Australia" | x, d_b) = 0.35 (the passage is about a different final and doesn't mention Australia; the generator falls back partly on its own pretrained prior, which correctly remembers Australia as a frequent World Cup winner, so it hedges rather than committing), and p_θ("Australia" | x, d_c) = 0.10 (the IPL passage is India-centric and gives the generator no signal favoring Australia at all). Because the answer here is a single token, RAG-Sequence and RAG-Token coincide exactly, and the marginal probability is:

p(y = "Australia" | x) = Σ_z p_η(z|x) · p_θ("Australia" | x, z)
 = 0.6453×0.90 + 0.2900×0.35 + 0.0647×0.10
 = 0.5808 + 0.1015 + 0.0065
 = 0.6888

The following code reproduces this calculation exactly, using only the standard library:

import math

# Query and document embeddings (toy 3-D stand-ins for 768-D DPR vectors)
q = (1.0, 0.5, -0.2)
docs = {
    "d_a": (2.4, 1.0, 1.0),   # correct: 2023 final result
    "d_b": (1.0, 1.4, -1.0),  # wrong year: 2019 final
    "d_c": (0.2, 0.6, 0.5),   # off-topic: IPL passage
}

def dot(u, v):
    return sum(a * b for a, b in zip(u, v))

scores = {name: dot(q, vec) for name, vec in docs.items()}
# scores == {"d_a": 2.7, "d_b": 1.9, "d_c": 0.4}

exp_scores = {name: math.exp(s) for name, s in scores.items()}
Z = sum(exp_scores.values())
p_eta = {name: v / Z for name, v in exp_scores.items()}
# p_eta ≈ {"d_a": 0.6453, "d_b": 0.2900, "d_c": 0.0647}

p_theta_australia = {"d_a": 0.90, "d_b": 0.35, "d_c": 0.10}

p_answer = sum(p_eta[z] * p_theta_australia[z] for z in docs)
print(round(p_answer, 4))   # 0.6888

Tracing it line by line: dot computes the three scores exactly as above; exp_scores and Z reproduce the softmax; p_eta is the dictionary {"d_a": 0.6453..., "d_b": 0.2899..., "d_c": 0.0647...}; the final sum multiplies each weight by the matching generator probability and adds them, so print outputs 0.6888, matching the hand derivation.

Making retrieval fast enough to matter: approximate nearest-neighbor search

The three-document example only works because someone already narrowed the corpus down to three candidates. The original RAG paper's retrieval index was built over Wikipedia split into roughly 21 million 100-word passages (21,015,324, per the DPR paper's exact count) — and a production Indian legal-search system would face a similarly large table of judgments, statutes, and commentary. Finding the true top-k nearest neighbors of a 768-dimensional query vector by brute force means computing a dot product against every single passage vector: roughly N × 2d floating-point operations per query, where N ≈ 21×10⁶ and d = 768 (768 multiplications plus 768 additions per comparison). That's on the order of 21×10⁶ × 1536 ≈ 3.2×10¹⁰ operations for one query — order-of-magnitude enough to make interactive latency (a budget usually well under a second) impossible if paid in full for every question.

This is why real systems use Maximum Inner Product Search (MIPS) approximations rather than brute force, and the dominant algorithm — used inside the FAISS library that the RAG paper itself relies on — is Hierarchical Navigable Small World graphs (HNSW), from Malkov and Yashunin's 2018 IEEE TPAMI paper. HNSW builds a multi-layer graph over the corpus: the top layer contains a sparse sample of documents connected by long-range links, and each layer below adds more documents and shorter, denser links, until the bottom layer contains every document in the corpus. A query is answered by starting at an entry point in the sparse top layer and greedily walking to whichever neighboring node is closer to the query, exactly like a "six degrees of separation" traversal — cheap because each layer has few nodes to compare against. When no neighbor in the current layer improves on the current best node, the search drops down one layer at that same node and repeats, now with access to more, closer-together nodes for finer-grained comparisons. By the time it reaches the bottom layer, it is already near the true nearest neighbors and only has to explore a small local neighborhood to finalize the top-k. The expected cost is close to O(d · log N) rather than O(d · N) — for N = 21×10⁶, log₂N ≈ 24.3, so the search visits on the order of a few dozen small candidate sets instead of comparing against all 21 million vectors, which is the difference between an interactive system and one that times out.

Common misconception: "more retrieved documents always makes RAG more accurate"

Students who have only seen the "retrieve top-k and stuff it in the prompt" version of RAG tend to assume that raising k is a free win — more evidence, more accuracy. The worked example above disproves this directly. If the retriever had returned only d_a (k = 1), the marginal probability of the correct answer would simply be p_θ("Australia" | x, d_a) = 0.90, since there is nothing else to average against. Once the wrong-year and off-topic passages are added at k = 3, the correct answer's probability is pulled down to 0.6888 — a full 21 percentage points lower — purely because two low-relevance documents now carry nonzero weight in the mixture and drag the average toward their weaker, less-grounded generator outputs. This is not a bug in the arithmetic; it is what marginalizing over a latent variable necessarily does when some of that variable's outcomes are noisy. The correct mental model is a precision/recall tradeoff, not a monotonic improvement curve: too small a k risks the retriever missing the one document that actually contains the answer (a recall failure with no chance of recovery downstream), while too large a k dilutes the generator's effective attention with irrelevant material and, as shown here, can measurably lower the probability mass on the correct answer even when the right document is present in the set. The original paper found accuracy on open-domain QA tasks generally rising as k grows from very small values and then flattening or degrading well before k reaches the tens — the practical fix is not "always retrieve more" but improving retriever precision itself (harder negative mining during DPR training, or adding a cross-encoder re-ranking pass after the initial ANN search) so that the top-k set is small and clean rather than large and padded with plausible-looking noise.

A different fusion strategy: Fusion-in-Decoder

RAG-Sequence and RAG-Token both fuse evidence through an explicit probability mixture — a weighted average of separately-computed distributions. Izacard and Grave's Fusion-in-Decoder (EACL 2021) takes a structurally different approach to the same problem: every retrieved passage is encoded independently by the encoder (so encoding cost scales linearly and stays parallelizable even with a hundred retrieved passages), but all of those encoded representations are then concatenated and handed to a single decoder, which cross-attends over the entire concatenation at once while generating the answer. There is no softmax mixing of separate output distributions — the "fusion" happens once, inside the decoder's attention, rather than as an explicit weighted sum of per-document answers after the fact. This tends to scale better to larger k (FiD systems commonly retrieve dozens to a hundred passages, where RAG-style mixtures become expensive since each document requires its own full decoder pass) at the cost of losing the clean, interpretable per-document probability weight that makes RAG-Sequence and RAG-Token easy to reason about mathematically, as in the worked example above.

Active recall

Attempt each question before reading its answer.

1. Write the RAG-Sequence and RAG-Token marginalization formulas and state, in one sentence, the structural difference between them.
2. The query embedding changes to q' = (0.5, 0.5, 0.5) (assume the generator probabilities p_θ for each document are unchanged). Recompute the retrieval scores, the softmax weights p_η(z|x), and the new marginal probability of "Australia."
3. Using the original query q = (1.0, 0.5, −0.2), the retriever is reconfigured to use k = 2, dropping d_c. Recompute p_η(d_a), p_η(d_b), and the marginal probability of "Australia." Does removing the least relevant document raise or lower grounding confidence, and why?
4. Why is brute-force nearest-neighbor search over a 21-million-document Wikipedia index infeasible at interactive latency, and what does HNSW do differently? Give the order-of-magnitude argument.
5. True or false: "Since RAG-Sequence and RAG-Token both marginalize over the same top-k documents with the same retrieval weights p_η, they must always assign the same probability to any given candidate answer." Justify your answer.
6. For the legal-research assistant from the opening scenario, propose one retriever-side change (not a prompt-engineering change) that would reduce the dilution effect demonstrated in question 3.

Answers

1. RAG-Sequence: p(y|x) = Σ_z p_η(z|x)·p_θ(y|x,z), where the whole sequence y is generated conditioned on one fixed document before mixing across documents. RAG-Token: p(y|x) = Π_i Σ_z p_η(z|x)·p_θ(y_i|x,z,y_<i), which re-mixes across documents independently at every generated token. Structurally: RAG-Sequence commits to one document per full hypothesis and averages sequence-level probabilities; RAG-Token averages token-level probabilities and can effectively draw on different documents for different tokens within the same answer.

2. Scores: score(d_a) = 0.5(2.4)+0.5(1.0)+0.5(1.0) = 1.2+0.5+0.5 = 2.2; score(d_b) = 0.5(1.0)+0.5(1.4)+0.5(−1.0) = 0.5+0.7−0.5 = 0.7; score(d_c) = 0.5(0.2)+0.5(0.6)+0.5(0.5) = 0.1+0.3+0.25 = 0.65. Note d_b and d_c are now nearly tied. Using e^2.2 ≈ 9.0250, e^0.7 ≈ 2.0138, e^0.65 ≈ 1.9155, Z ≈ 12.9543: p_η(d_a) ≈ 0.6967, p_η(d_b) ≈ 0.1555, p_η(d_c) ≈ 0.1479. New marginal: 0.6967(0.90) + 0.1555(0.35) + 0.1479(0.10) = 0.6270 + 0.0544 + 0.0148 = 0.6962. The ripple: even though d_a's absolute score dropped (2.7 → 2.2), its relative dominance over the other two rose slightly (0.6453 → 0.6967) because the new query is even less aligned with d_b and d_c's distinguishing dimensions, so the correct answer's marginal probability rises slightly, from 0.6888 to about 0.6962 — a reminder that what matters is relative score gaps, not absolute magnitudes.

3. With only d_a and d_b, Z = 14.8797 + 6.6859 = 21.5656. p_η(d_a) = 14.8797/21.5656 ≈ 0.6900; p_η(d_b) ≈ 0.3100. Marginal: 0.6900(0.90) + 0.3100(0.35) = 0.6210 + 0.1085 = 0.7295. This is higher than the k=3 value of 0.6888, because d_c was the most irrelevant document and contributed almost nothing but dilution — removing it lets the remaining weight concentrate more heavily on d_a. This confirms the precision point from the misconception section: removing a low-relevance document can improve grounding, it is not automatically a loss of information.

4. Brute force costs roughly N × 2d operations per query; with N ≈ 21×10⁶ and d = 768, that's about 3.2×10¹⁰ operations — far too slow for a sub-second interactive budget. HNSW instead organizes the corpus into a multi-layer graph, sparse at the top and dense at the bottom, and answers a query by greedily descending from a sparse entry point, comparing against only a small local neighborhood at each layer rather than the whole corpus. Expected cost is close to O(d·log N); with log₂(21×10⁶) ≈ 24.3, the number of comparisons needed drops by several orders of magnitude relative to brute force.

5. False in general. For a single-token answer the two formulas are algebraically identical (the product in RAG-Token has only one factor, so it collapses to the same expression as RAG-Sequence) — which is exactly why question 1's example worked for both. For sequences of two or more tokens, RAG-Sequence computes Σ_z p_η(z)·Π_i p_θ(y_i|x,z,y_<i) (a sum of per-document sequence products) while RAG-Token computes Π_i Σ_z p_η(z)·p_θ(y_i|x,z,y_<i) (a product of per-token sums). A sum-of-products is not algebraically equal to a product-of-sums except in degenerate cases (only one document has nonzero weight, or every document gives identical per-token conditionals), so the two formulas generally diverge once the answer has more than one token.

6. Improve retriever precision rather than retrieval breadth: fine-tune the DPR-style query/passage encoders with hard negatives specifically drawn from same-topic-wrong-year or same-topic-wrong-forum confusions (the exact d_b/d_c pattern above), or insert a cross-encoder re-ranking stage that re-scores the ANN-retrieved top-k before it reaches the generator, so that low-relevance passages are demoted out of the mixture before they get to dilute the marginal probability — rather than trying to fix the problem after the fact with a better prompt.

Think About It

Think about this: How would you explain retrieval-augmented generation: combining llms with knowledge to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where retrieval-augmented generation: combining llms with knowledge is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting retrieval-augmented generation: combining llms with knowledge to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind retrieval-augmented generation: combining llms with knowledge, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.