Module 15: Retrieval-Augmented Generation
Introduction
Your model (m06–m08) knows exactly what was in its training data, frozen at a cutoff date, blended into billions of weights. Ask it about a document it never saw — your company’s wiki, a paper from last week, a fact it half-remembers — and it will confidently make something up. Retrieval-Augmented Generation (RAG) fixes this without touching the weights: at query time, search an external corpus for relevant passages and paste them into the prompt. The model then answers from fresh, specific, checkable text instead of parametric memory alone.
The engine underneath is nearest-neighbor search in an embedding space. Encode every document as a vector; encode the question the same way; retrieve the documents whose vectors point most nearly the same direction (highest cosine similarity); stuff them into the prompt; generate. This module builds that whole pipeline from scratch.
To keep it fully runnable with no trained encoder, we encode with classic TF-IDF vectors — term frequency times inverse document frequency, computable by hand. Modern RAG swaps TF-IDF for dense embeddings from a trained encoder (the m04 idea, learned end-to-end), but the pipeline is identical — only the encode step changes.
Why it matters for LLMs:
- Knowledge without retraining. New facts arrive by adding documents to the index, not by a training run. The model stays fixed; the corpus is live.
- Grounding and citations. The answer is conditioned on retrieved text you can show the user — the standard defense against hallucination.
- It is everywhere. Chat-with-your-docs, coding assistants over a repo, search copilots — almost every applied LLM system is a RAG system.
What You’ll Learn
After this module, you can:
- Explain the RAG pipeline — encode → search → retrieve → augment → generate — and why it beats parametric memory for fresh or private knowledge.
- Build TF-IDF document vectors from scratch (
tf, smoothedidf). - Rank documents by cosine similarity and retrieve the top-k.
- Chunk long documents, and de-duplicate results with Maximal Marginal Relevance (MMR).
- Assemble a grounded prompt and wrap it all in a
Retriever. - Build a dense bi-encoder from scratch, train it with the in-batch-negatives contrastive loss, and see it retrieve a paraphrase that TF-IDF provably cannot.
- Build ColBERT late interaction — one vector per token scored by MaxSim — the precomputable middle ground between a bi-encoder and a cross-encoder, and see it win by coverage where a single pooled vector is fooled by repetition.
- Train a Matryoshka embedding whose prefixes are all usable, so you can slice-and-renormalize one stored vector to any width — and see it retrieve from a 4-dim prefix where a plain embedding, truncated the same way, collapses.
- Build an HNSW approximate-nearest-neighbor index from scratch — a navigable small-world graph you walk downhill in
O(log N)hops instead of scanning allNvectors — and dial the beam widthefto trade recall for cost, the structure behind FAISS and every production vector database.
Prerequisites
This module requires familiarity with:
- Module 04: Embeddings — vectors as meaning; dense retrieval replaces TF-IDF with learned embeddings of exactly this kind.
- Module 08: Generation — RAG conditions generation on retrieved context; the decoding is unchanged.
- Module 20: Multimodality (optional) — the dense bi-encoder trains with the same contrastive loss CLIP uses to align images and text; here it aligns queries and documents.
Intuition: The RAG Pipeline
RAG is five steps, and only the middle three are new — the model at the end is the same one you already built. Step through what each stage produces:
NoteKey Insight
Only Encode → Search → Retrieve → Augment is retrieval; Generate is the model you already have. RAG is not a new kind of model — it is a way of choosing what goes in the prompt. Everything hard is in ranking documents by relevance.
The Math: TF-IDF and Cosine Similarity
To search by meaning we need documents as vectors. The classic recipe weights each term by how often it appears in a document (term frequency) against how rare it is across the corpus (inverse document frequency), so shared rare words count and boilerplate like “the” does not:
\text{tf}(t, d) = \text{count of } t \text{ in } d, \qquad \text{idf}(t) = \ln\!\frac{1 + N}{1 + \text{df}(t)} + 1, \qquad \text{tfidf}(t, d) = \text{tf}(t, d)\cdot\text{idf}(t),
where N is the number of documents and \text{df}(t) how many contain t. The +1 smoothing keeps every idf positive; a term in every document gets idf = 1.
Relevance is the angle between vectors, not their length — a long document is not more relevant just for being long — so we compare with cosine similarity:
\cos(\mathbf{q}, \mathbf{d}) = \frac{\mathbf{q}\cdot\mathbf{d}}{\lVert\mathbf{q}\rVert\,\lVert\mathbf{d}\rVert}.
Normalize both vectors to unit length and cosine is just a dot product. retrieval.py implements tfidf_matrix, cosine_similarity, and retrieve.
Code: Build a Retriever
The Retriever class fits a corpus once, then answers queries by cosine top-k:
from retrieval import Retriever
corpus = [
"The transformer architecture uses self-attention to relate every token to every other token.",
"Attention computes a weighted sum of value vectors using query-key similarity scores.",
"Byte-pair encoding builds a subword vocabulary by merging the most frequent adjacent pairs.",
"Photosynthesis lets plants convert sunlight, water, and carbon dioxide into glucose and oxygen.",
"Mount Everest is the tallest mountain above sea level, on the border of Nepal and Tibet.",
]
retriever = Retriever().fit(corpus)
for idx, score, doc in retriever.search("how does attention work?", k=2):
print(f"[{idx}] {score:.3f} {doc[:60]}...")[1] 0.233 Attention computes a weighted sum of value vectors using que...
[0] 0.185 The transformer architecture uses self-attention to relate e...The attention documents win because they share the rare, informative words of the query. Now assemble the retrieved passages into a grounded prompt — the “augment” step:
from retrieval import build_prompt
query = "how does attention work?"
hits = retriever.search(query, k=2)
prompt = build_prompt(query, [doc for _, _, doc in hits])
print(prompt)Answer the question using only the context below. If the answer is not in the context, say you don't know.
Context:
[1] Attention computes a weighted sum of value vectors using query-key similarity scores.
[2] The transformer architecture uses self-attention to relate every token to every other token.
Question: how does attention work?
Answer:That string is exactly what you would hand to generate() from m08. The model now answers from the retrieved context — and because you have the passages, you can show them as citations. Swap the TF-IDF encode for a trained dense encoder and nothing else in this pipeline changes.
Interactive: Watch Retrieval Rank the Corpus
Pick a question and see cosine similarity score every document. The top-k (here k = 3) are what gets retrieved and pasted into the prompt; everything else is ignored. Notice how the scores concentrate on the documents that share the query’s informative words.
TipTry This
- Switch between the attention query and the photosynthesis query. Watch the mass of similarity jump to a completely different pair of documents — the corpus is the same; only the query moved.
- Note the ignored documents still get a small nonzero score from incidental shared words (“the”, “of”). idf is what keeps those from dominating.
Chunking and Diversity
Two practical problems break naïve retrieval, and retrieval.py handles both.
Chunking. You do not index whole documents — a 50-page PDF has one vector that means nothing specific. You split it into small overlapping windows so each retrievable passage is focused, with a little overlap so a fact spanning a boundary survives in at least one chunk:
from retrieval import chunk_text
passage = "Retrieval augmented generation searches a corpus then conditions the model on the results"
print(chunk_text(passage, chunk_size=5, overlap=2))['Retrieval augmented generation searches a', 'searches a corpus then conditions', 'then conditions the model on', 'model on the results']Diversity. Plain top-k happily returns three near-identical passages, wasting the context window. Maximal Marginal Relevance (MMR) picks each next passage for relevance minus similarity to what is already chosen:
\text{next} = \arg\max_{c}\;\Big[\lambda\cdot\text{rel}(c) - (1-\lambda)\max_{s\in S}\text{sim}(c, s)\Big].
from retrieval import tfidf_matrix, encode_query, retrieve, mmr
docs = [
"transformer attention model",
"transformer attention model network", # near-duplicate of doc 0
"retrieval augmented generation search",
]
matrix, vocab, idf = tfidf_matrix(docs)
qv = encode_query("transformer attention", vocab, idf)
print("plain top-2:", [i for i, _ in retrieve(qv, matrix, k=2)]) # two near-duplicates
print("MMR top-2: ", mmr(qv, matrix, k=2, lambda_=0.3)) # swaps in the different docplain top-2: [0, 1]
MMR top-2: [0, 2]With \lambda = 1, MMR is just top-k; lower \lambda trades a little relevance for coverage.
Dense Retrieval: Matching Meaning, Not Words
Everything so far shares one blind spot. TF-IDF is a bag of words: two texts are similar only when they reuse the same tokens. Ask it for an automobile and it will never find a passage about a car — no shared word, cosine similarity exactly zero. Real questions and real documents paraphrase each other constantly, and lexical retrieval walks straight past the match.
Dense retrieval fixes this by learning the encoder. Instead of one dimension per vocabulary word, each text becomes a short, dense vector positioned so that things that mean the same thing sit nearby — even with no token in common. And here is the payoff the whole module has been promising: only the encode step changes. The same cosine_similarity → retrieve → build_prompt pipeline runs on top, unchanged. dense.py builds the learned encoder from scratch.
Intuition: Two Towers
The architecture is a bi-encoder (two towers): one encoder for queries, one for documents, both mapping into a single shared vector space. A text’s vector is just the mean of its token embeddings (the pooling Sentence-BERT uses), scaled to unit length so a dot product reads off as a cosine. Step through one tower:
NoteKey Insight
A dense encoder is nothing but embeddings (m04) plus pooling. The words never touch; the learned geometry does the matching. Everything hard has moved from the retrieval algorithm into how we train the two embedding tables.
The Math: Contrastive Training
How do the tables learn that “automobile” and “car” belong together? With in-batch-negatives contrastive learning — the exact objective m20 uses to align images with captions, here aligning queries with documents. Take a batch of matched (\text{query}_i, \text{doc}_i) pairs, encode both sides, and form the similarity matrix
S_{ij} = \frac{\mathbf{q}_i \cdot \mathbf{d}_j}{\tau}.
Each query’s own document (the diagonal S_{ii}) is the positive; every other document in the batch is a free negative. Cross-entropy pushes each query to pick its own document out of the batch, and — symmetrically — each document to pick its own query:
\mathcal{L} = \tfrac{1}{2}\Big[\text{CE}(S,\, \text{arange}) + \text{CE}(S^\top,\, \text{arange})\Big].
This is identical to CLIP’s loss (\tau is the temperature); the only difference is that both towers now read text. As training proceeds, the diagonal of S lights up while the off-diagonal goes dark — watch it happen:
Code: A Bi-Encoder From Scratch
BiEncoder holds the two embedding tables; train_bi_encoder runs the loss above. The corpus is deliberately adversarial for TF-IDF: every query and its correct document are synonyms with no shared word, so lexical similarity is exactly zero and only meaning can bridge them.
from dense import synonym_dataset, train_bi_encoder, DenseRetriever
pairs = synonym_dataset()
for q, d in pairs[:3]:
print(f"query “{q}” ↔ doc “{d}” (shared words: {set(q.split()) & set(d.split())})")
model, history = train_bi_encoder(pairs, epochs=400, seed=0)
print(f"\ncontrastive loss: {history['losses'][0]:.3f} → {history['losses'][-1]:.3f}")query “automobile motorcar” ↔ doc “car vehicle” (shared words: set())
query “physician doctor” ↔ doc “clinician medic” (shared words: set())
query “ocean sea” ↔ doc “marine maritime” (shared words: set())
contrastive loss: 5.754 → 0.000The pairs share no tokens, yet the encoder learns to align them. Now retrieve with a held-out phrasing — a single synonym the model never saw as a training query:
docs = [d for _, d in pairs]
retriever = DenseRetriever(model).fit(docs)
for query in ["automobile", "physician", "cash"]:
idx, score, doc = retriever.search(query, k=1)[0]
print(f"“{query}” → “{doc}” (cosine {score:.2f})")“automobile” → “car vehicle” (cosine 0.49)
“physician” → “clinician medic” (cosine 0.54)
“cash” → “currency funds” (cosine 0.64)DenseRetriever has the same fit/search interface as the TF-IDF Retriever and reuses the identical cosine_similarity and retrieve — the learned encode is the only thing that changed.
Interactive: Lexical vs. Dense on the Same Query
Here is the whole point in one picture. The query is “cash money” and its true match is a synonym passage (“currency funds”) with no word in common. TF-IDF scores every document at exactly 0 — a total tie, blind to the match — so its “top-1” is an arbitrary wrong guess. The trained dense encoder scores the correct passage highest.
TipTry This
Every TF-IDF bar is pinned at 0 because the query and its answer share no token — lexical retrieval literally cannot see the match. Scrub the training-snapshot slider above back to epoch 0: the dense scores start just as flat, and the signal is learned, not built in.
Interactive: The Embedding Space
Project every query and document vector down to 2D (PCA). Before training the towers are random noise; after training, each query lands next to its synonym document — the geometry that makes dense retrieval work. Lines connect the matched pairs.
NoteKey Insight
This is why production RAG is dense. A learned encoder retrieves by meaning, so a question and its answer match even when they share no vocabulary — exactly the paraphrase gap TF-IDF cannot close. The pipeline around it never changed.
Hybrid Retrieval: Best of Both
Dense retrieval fixed lexical retrieval’s blind spot — but it has one of its own. Ask for a passage about an exact literal token the encoder never learned — an error code E-4021, a commit hash, a function name, a product SKU — and the dense encoder shrugs: that token is out-of-vocabulary, so it carries no learned meaning and the query lands nowhere useful. Lexical retrieval, meanwhile, matches that literal token instantly.
So the two retrievers have complementary blind spots:
| matches | blind to | |
|---|---|---|
| Lexical (TF-IDF / BM25) | exact words, rare literals, names, codes | paraphrases (no shared word → score 0) |
| Dense (bi-encoder) | meaning, synonyms, paraphrases | rare literals it never learned |
Hybrid retrieval runs both and fuses their rankings, so a document either retriever likes rises to the top — recovering the union of what each finds alone. Two pieces make this work: a sharper lexical scorer (BM25) and a scale-free way to combine ranked lists (Reciprocal Rank Fusion).
The Math: BM25, TF-IDF’s Sharper Cousin
TF-IDF counts a term linearly: a word that appears 20 times contributes 20× a word that appears once. That over-rewards keyword stuffing. BM25 (Best Match 25, Robertson & Zaragoza) fixes two things — it saturates term frequency and normalizes for document length:
\text{score}(D, Q) = \sum_{t \in Q} \text{IDF}(t) \cdot \frac{f(t, D)\,(k_1 + 1)}{f(t, D) + k_1\left(1 - b + b\,\frac{|D|}{\text{avgdl}}\right)}
with the nonnegative IDF \text{IDF}(t) = \ln\!\left(\frac{N - \text{df}(t) + 0.5}{\text{df}(t) + 0.5} + 1\right). The term-frequency factor \frac{f(k_1+1)}{f + k_1(\cdots)} climbs fast, then flattens toward the asymptote k_1 + 1 — the 10th occurrence of a word barely beats the 2nd. Two knobs: k_1 (saturation, default 1.5) and b (length normalization, default 0.75; b=0 turns it off). What BM25 keeps from TF-IDF is the blind spot: a query term absent from a document contributes 0, so a paraphrase with no shared word still scores exactly 0.
Watch the saturation — the curve BM25 uses where TF-IDF draws a straight line:
TipTry This
Switch k_1 between 0.5 and 3.0. Small k_1 saturates sooner — the 2nd occurrence already nears the dashed asymptote, so extra repeats barely count. Large k_1 stays closer to the straight TF-IDF line for longer. TF-IDF is the limit k_1 \to \infty: no saturation at all.
The Math: Reciprocal Rank Fusion
Now fuse the two retrievers. The trouble: BM25 scores are unbounded positives, cosine similarities live in [-1, 1] — adding them directly is meaningless, and tuning a weight \alpha\,\text{bm25} + (1-\alpha)\,\text{cosine} needs calibration per corpus. Reciprocal Rank Fusion sidesteps all of it by throwing away the scores and keeping only the rank:
\text{RRF}(d) = \sum_{r \in \text{retrievers}} \frac{1}{k + \text{rank}_r(d)}
Each document scores the reciprocal of its 1-based rank in every list it appears in; the constant k (Cormack et al. use 60) dampens the top so no single retriever’s #1 can dominate. A document ranked well by both retrievers beats one ranked #1 by a single retriever and buried by the other — precisely the complementary-evidence behavior we want. Drive it:
NoteKey Insight
doc 2 is ranked #1 by the sparse retriever and #2 by the dense one — supported by both — so RRF floats it to the top, above doc 4 (dense’s #1 but sparse’s last) and doc 0 (sparse’s #2 but dense’s last). Fusion rewards agreement. Slide k toward 0 and the #1 ranks dominate (a single retriever can win outright); slide it up and the gaps flatten (deeper ranks matter more).
Code: Fuse a Lexical and a Semantic Retriever
BM25Retriever has the same .fit / .search API as the TF-IDF Retriever and the DenseRetriever, so HybridRetriever just fuses any list of them:
from hybrid import BM25Retriever, HybridRetriever, reciprocal_rank_fusion
docs = [
"runtime fault e4021 crashed the worker thread", # 0: rare literal 'e4021'
"car vehicle congestion downtown at rush hour", # 1: about automobiles
]
bm25 = BM25Retriever().fit(docs)
# BM25 nails the exact literal token...
print("BM25 'e4021': ", [i for i, s, _ in bm25.search("error e4021", k=2)], "scores>0:",
[round(s, 2) for _, s, _ in bm25.search("error e4021", k=2)])
# ...but is blind to a paraphrase with no shared word (all scores 0):
print("BM25 'automobile traffic':", [round(s, 3) for s in bm25.scores("automobile traffic").tolist()])BM25 'e4021': [0, 1] scores>0: [0.69, 0.0]
BM25 'automobile traffic': [0.0, 0.0]Reciprocal Rank Fusion is one pure function over ranked index lists:
# Two retrievers' rankings of 3 docs; doc 1 is liked by both.
sparse_ranking = [1, 0, 2]
dense_ranking = [1, 2, 0]
print(reciprocal_rank_fusion([sparse_ranking, dense_ranking], k=60))
# doc 1 wins with 1/61 + 1/61; the fusion never touched a raw score.[(1, 0.03278688524590164), (0, 0.03200204813108039), (2, 0.03200204813108039)]hybrid_search fuses only positive-evidence candidates: a retriever that scores a query at zero across the board (BM25 on a paraphrase; a dense encoder on an OOV literal) has no real ranking — just an arbitrary tie order — so it does not vote. That is what lets the fusion recover the union instead of averaging one retriever’s signal against the other’s noise.
Interactive: Hybrid Recovers the Union
The proof, on a corpus built to defeat each retriever alone. Three lexical queries share a rare literal token with their answer (BM25’s home turf); three semantic queries are pure synonyms with no shared word (dense’s home turf). We train the bi-encoder on the synonym half only — so the literal tokens stay out-of-vocabulary to it — then score each retriever’s recall@1:
Which retriever answered which query — a filled cell is a hit, the empty cells are each retriever’s blind spot:
NoteKey Insight
BM25 recovers the three lexical queries and scores the three semantic ones at exactly zero; dense does the mirror image. Each retriever alone tops out at recall 0.5. Fusing them recovers the union — recall 1.0, every query answered at rank 1. That is why serious RAG stacks run hybrid retrieval: BM25 for the exact terms, a dense encoder for the meaning, RRF to let each cover the other’s blind spot.
Re-ranking: Score the Shortlist Jointly
Every retriever so far — TF-IDF, the dense bi-encoder, the BM25 + RRF hybrid — shares one shape: it embeds the query and each document independently, then compares the vectors. That independence is the source of their speed. A bi-encoder can encode the whole corpus once, offline, so a query is one encode plus a dot product against a stored matrix. It is also their ceiling: the query and the document never actually meet. Each document is crushed into a single vector before the model has seen the query, so it can never ask “does this query word line up with that document word?”.
A cross-encoder removes the independence. It reads the (query, document) pair together — query tokens attending to document tokens and back — and emits one relevance score for that specific pair. Seeing both sides at once lets it read word order, binding, and negation that a bag-of-vectors bi-encoder averages away. The catch is that nothing can be precomputed: there is no query vector and no document vector to cache, only a score for a pair. Scoring a corpus of N documents means N full forward passes — you cannot run a cross-encoder over a million documents per query.
So real systems run two stages — retrieve, then re-rank: a cheap first stage (bi-encoder / BM25 / hybrid) pulls a shortlist of N candidates from the whole corpus for recall, and the slow, accurate cross-encoder re-scores only those N pairs for precision. Step through the pipeline:
NoteKey Insight
The two stages divide the labor: retrieve for recall (find the right documents cheaply, from the whole corpus), re-rank for precision (order a small shortlist accurately, expensively). The cross-encoder never touches the corpus — only the N candidates the first stage already narrowed to.
The Math: A Cross-Encoder, and the cross Flag
Our cross-encoder (rerank.py) concatenates the query and document into one sequence and runs a single self-attention block over it (m05), with a segment embedding marking which side each token is on:
input: [ q₁ q₂ … q_Lq | d₁ d₂ … d_Ld ] (query tokens, then document tokens)
└── one self-attention block (query↔document attention) ──┘
score: pool(query positions) · pool(document positions)The score is the (normalized, temperature-scaled) dot of a mean-pool over the query positions and a mean-pool over the document positions — the same cosine scoring the dense bi-encoder uses. What makes it a cross-encoder is the one thing the attention block adds: the query positions attend to the document positions, so the query’s pooled vector depends on the document.
Make that precise with a single flag. score(query, doc, cross=…) toggles the attention mask:
cross=True— full attention. The query and document tokens attend to each other. Joint. This is the cross-encoder.cross=False— a block-diagonal mask forbids cross-side attention: query tokens attend only to query tokens, document tokens only to document tokens. Each side is encoded in isolation, so the query’s pooled vector is provably independent of the document. The model has collapsed into a bi-encoder.
The “cross” in cross-encoder is exactly those query↔︎document attention edges — mask them and you are back to a bi-encoder. We can prove it: encode one query beside two different documents and check whether the query vector changed.
NoteKey Insight
On the left (cross=False) the two lines lie exactly on top of each other — the query vector is bit-for-bit identical whether it is paired with document A or B, so it cannot depend on the document: a bi-encoder. On the right (cross=True) they separate — the query’s representation now reflects the document it was scored against: a cross-encoder. demonstrate_independence() confirms it with torch.equal (no_cross → True, cross → False).
Code: Build and Train the Re-ranker
CrossEncoder is a small from-scratch transformer block; score returns one number for a pair. Because it re-scores a shortlist, we train it as a ranker with the pairwise (Bradley-Terry) loss from m12 — push each relevant document’s score above an irrelevant one’s, nothing more:
import torch
from rerank import CrossEncoder, train_cross_encoder, pair_score, side_is_independent
from retrieval import build_vocab
# (query, relevant_doc, irrelevant_doc) triples — the irrelevant doc is the
# word-reversal, same three words in the opposite order.
triples = [
("wolves hunt deer", "wolves hunt deer", "deer hunt wolves"),
("cats chase mice", "cats chase mice", "mice chase cats"),
]
model, history = train_cross_encoder(triples, dim=32, epochs=200, seed=0)
print(f"pairwise loss: {history['losses'][0]:.3f} → {history['losses'][-1]:.3f}")
# After training, the correctly-ordered document outscores its reversal:
q = "wolves hunt deer"
print(f"score(fact) = {pair_score(model, q, 'wolves hunt deer'):+.2f}")
print(f"score(reversal) = {pair_score(model, q, 'deer hunt wolves'):+.2f}")pairwise loss: 0.684 → 0.000
score(fact) = +9.84
score(reversal) = -4.36The cross flag is the whole bi-vs-cross distinction — the query representation is independent of the document only when the cross edges are masked off:
# The exact structural anchor: mask the cross edges → the query ignores the document.
print("cross=False, query independent of doc?",
side_is_independent(model, q, "cats chase mice", "deer hunt wolves", cross=False))
print("cross=True, query independent of doc?",
side_is_independent(model, q, "cats chase mice", "deer hunt wolves", cross=True))cross=False, query independent of doc? True
cross=True, query independent of doc? Falsererank re-scores a first stage’s candidates and reorders them; retrieve_then_rerank is the whole two-stage pipeline over any retriever:
from retrieval import Retriever
from rerank import retrieve_then_rerank
# The fact and its reversal share a bag of words, so TF-IDF ties them (identical
# scores) — it recalls both into the shortlist but cannot order them. The re-ranker breaks the tie.
docs = ["wolves hunt deer", "deer hunt wolves"]
first_stage = Retriever().fit(docs) # cheap, permutation-blind
print("first-stage scores:", [round(s, 2) for _, s, _ in first_stage.search("wolves hunt deer", k=2)])
top = retrieve_then_rerank(first_stage, model, "wolves hunt deer", pool=2, k=1)
print("re-ranked top-1: ", top[0][2]) # 'wolves hunt deer'first-stage scores: [1.0, 1.0]
re-ranked top-1: wolves hunt deerThe Headline: Re-ranking Recovers the Order the First Stage Is Blind To
Here is the division of labor made provable. The corpus pairs eight subject-verb-object facts with their reversals — “wolves hunt deer” vs “deer hunt wolves”, the same three words, the opposite meaning. A bag-of-words first stage (TF-IDF) is permutation-invariant: a fact and its reversal have the same bag of words, hence the same vector, hence — for any query — bit-for-bit identical scores. It recalls both but cannot order them, so recall@1 is a coin flip (0.5). The order-aware cross-encoder re-ranks that tied shortlist and lifts recall@1 to 1.0.
Why not just run the cross-encoder over everything? Because it costs one forward pass per document. Slide the corpus size and watch the gulf between scoring the whole corpus and re-ranking a fixed shortlist of 100:
NoteKey Insight
The first stage stalls at the 0.5 dashed line — it recalled the right document but, being order-blind, ties it with the distractor. The cross-encoder re-ranks the same shortlist to 1.0. And the cost chart is why it only ever sees a shortlist: re-ranking 100 candidates is a constant 100 forward passes no matter how large the corpus grows, while scoring the corpus itself scales with N.
TipTry This
- Watch it learn.
demonstrate_reranking()returns the traininglosses; the pairwise loss falls from ~0.69 (ln 2 — scores tied) toward 0 as the fact pulls ahead of its reversal. - Break the cross. Re-train with
train_cross_encoder(triples, cross=False)and re-rank withcross=False. The bi-encoder ablation can still use positional embeddings, so it may separate some pairs — but it must commit to a query vector before seeing the document. Compare its re-rank recall to the joint model’s. - Grow the shortlist. In
retrieve_then_rerank, raisepool. Bigger shortlists recover more first-stage misses but cost more cross-encoder passes — the recall / latency knob every production re-ranker tunes.
Late Interaction: the Precomputable Middle Ground
The two retrievers you just compared sit at opposite corners of one trade-off.
| encodes | interaction | precompute? | cost / query | |
|---|---|---|---|---|
Bi-encoder (dense.py) |
query, document → one vector each, independently | none — a single dot product | yes, offline | one encode + a vector search |
Cross-encoder (rerank.py) |
the pair → jointly, token-by-token | full, every token to every token | no | N forward passes |
Independence buys speed (index the corpus once); joint attention buys accuracy (read token-level matches). You seemingly have to pick one. ColBERT (Khattab & Zaharia, 2020) refuses the choice. It keeps the bi-encoder’s independence — query and document are still encoded separately, so document vectors precompute offline — but it stops pooling. Each text becomes one vector per token, and the interaction is deferred to scoring time. That deferral is the name: late interaction.
The scorer is MaxSim — for each query token, take its single best match over all document tokens, then sum:
S(q, d) \;=\; \sum_{i \,\in\, \text{query}} \; \max_{j \,\in\, \text{doc}} \; E_{q_i} \cdot E_{d_j}
Because the encoders never see each other, ColBERT precomputes like a bi-encoder; because MaxSim compares every query token to every document token, it recovers the fine-grained matching a single pooled vector throws away. It is literally the missing rung in the arc you have climbed: lexical → dense → hybrid → cross-encoder → late interaction. Step through what MaxSim computes:
NoteKey Insight
Late interaction is precisely independence in the encoders, interaction in the scorer. The document tower takes no query, so the corpus is indexed offline (fast, like a bi-encoder); MaxSim then lets each query token find its own best document token (fine-grained, like a cross-encoder). A bi-encoder is the special case with one token per side — MaxSim over 1×d matrices is a single dot product.
The Math: MaxSim, the Bi-encoder Limit, and Saturation
Three properties fall straight out of the formula, each an exact anchor.
Bounded. With L2-normalized token vectors every E_{q_i}\cdot E_{d_j} is a cosine in [-1, 1], so S \in [-L_q, L_q] for a query of L_q tokens — the score is a count of well-matched query tokens, at most one point per token.
Bi-encoder limit. Pool each side to a single token and MaxSim collapses to one dot product — the bi-encoder’s cosine. Late interaction strictly generalizes the bi-encoder: give it more tokens and it can only see more structure.
Saturation (the crux). The max caps each query token’s contribution at its one best document token. So a document cannot inflate its score by repeating a matching word — the second copy is never the max of anything new — and a long document is not diluted: extra tokens can only raise a max, never lower it. A single pooled vector has neither property. Pack a document with one query word and its pooled vector tilts hard toward it (repetition helps); pad a relevant document with filler and its pooled vector spreads thin, so its cosine drops (dilution hurts). This is exactly where a single vector loses and MaxSim wins.
Code: MaxSim and a ColBERT Retriever
colbert.py is the whole idea in three pieces: maxsim (the scorer), ColBERTModel (per-token encoders + score), and ColBERTRetriever (precompute an index, then MaxSim-search it). The document encoder takes no query, so a document’s vectors are the same for every query — that is what lets the index be built offline:
import torch
from colbert import maxsim, ColBERTModel, ColBERTRetriever, document_is_precomputable
from retrieval import build_vocab
# MaxSim by hand: query tokens e_x, e_y against doc tokens [e_x, e_x, e_y].
q = torch.tensor([[1., 0.], [0., 1.]])
d = torch.tensor([[1., 0.], [1., 0.], [0., 1.]])
print(f"MaxSim = max(1,1,0) + max(0,0,1) = {float(maxsim(q, d)):.1f}")
# A ColBERT model: per-token encoders + the MaxSim scorer.
torch.manual_seed(0)
docs = ["alpha beta gamma", "gamma delta epsilon", "alpha epsilon"]
model = ColBERTModel(build_vocab(docs), dim=16)
print(f"encode_doc('{docs[0]}') → {tuple(model.encode_doc(docs[0]).shape)} (one row per token)")
# Precomputable: the document's vectors do not depend on any query (torch.equal).
print(f"document_is_precomputable = {document_is_precomputable(model, docs[0])}")MaxSim = max(1,1,0) + max(0,0,1) = 2.0
encode_doc('alpha beta gamma') → (3, 16) (one row per token)
document_is_precomputable = TrueThe retriever’s fit encodes every document once into the index; search encodes the query and MaxSim-scores it against those stored matrices — no document is ever re-encoded at query time:
retriever = ColBERTRetriever(model).fit(docs)
for idx, score, doc in retriever.search("alpha gamma", k=2):
print(f" doc {idx} MaxSim={score:+.3f} {doc!r}") doc 0 MaxSim=-0.008 'alpha beta gamma'
doc 1 MaxSim=-0.626 'gamma delta epsilon'The Alignment MaxSim Induces
MaxSim is interpretable: it commits to one document token per query token, so you can read off which word matched which. With context-free one-hot tokens the grid is exactly 0/1 — a query token lights up wherever its word appears in the document, and MaxSim keeps the single brightest cell per row (a real contextual encoder softens these to cosines so synonyms light up too, but the mechanic is identical):
NoteKey Insight
Each query-token row keeps exactly one outlined cell — its MaxSim match. "alpha" matches the document’s alpha; "beta" matches the first beta and ignores the other two (they are never the max of anything). Summing the outlined cells gives the MaxSim score (here 2.0), and repeating beta bought the document nothing. That is saturation, made visible.
Coverage Beats Concentration
Here is the payoff, and it needs no training — just the one-hot vectors above, so every number is exact. Build a corpus where, for each query, the gold document covers both query aspects once (but is padded with filler) and a distractor piles on one aspect. Score every document two ways: MaxSim, and the single pooled vector a bi-encoder would use.
from colbert import demonstrate_saturation
sat = demonstrate_saturation()
print("query:", sat["queries"][0], " gold doc:", sat["relevant"][0])
for i, doc in enumerate(sat["docs"]):
print(f" MaxSim={sat['maxsim_scores'][0][i]:.2f} pool={sat['pool_scores'][0][i]:.2f} {doc!r}")
print(f"\nrecall@1 MaxSim={sat['recall']['maxsim']} sum-pool={sat['recall']['sum_pool']}")query: alpha beta gold doc: [0]
MaxSim=2.00 pool=0.58 'alpha beta filler one two three'
MaxSim=1.00 pool=0.71 'alpha alpha alpha alpha'
MaxSim=0.00 pool=0.00 'car ocean padding four five six'
MaxSim=0.00 pool=0.00 'ocean ocean ocean ocean'
MaxSim=0.00 pool=0.00 'gamma delta epsilon zeta'
recall@1 MaxSim=1.0 sum-pool=0.0Per-document scores for the first query (the gold document is d0) — watch the pooled vector rank a one-aspect distractor above the two-aspect gold:
NoteKey Insight
MaxSim retrieves the two-aspect gold for every query (recall@1 1.0); the pooled single vector picks a one-aspect distractor (recall@1 0.0 here). Repetition inflates the distractor’s pooled cosine while filler dilutes the gold’s — MaxSim, capping each query token at its best single match, is immune to both. This is the failure mode d0’s outlined bar makes concrete: the right document, ranked below a distractor by a single vector.
TipTry This
- Watch it learn meaning.
demonstrate_alignment()trains a contextualColBERTEncoderand retrieves paraphrases —"fast automobile"finds"quick car speeds down road", which shares no word (term-matching MaxSim scores it0). Itsrecallis1.0; the returnedlexical_scoresare all0. - Feel the saturation. In
maxsim_grid("alpha beta", "alpha alpha alpha alpha"), add or removealphacopies. The score stays1.0— the missingbetaaspect is never recovered by piling onalpha. - Reach the bi-encoder limit. Pool a query and a document to one vector each and confirm
maxsimequals their cosine. Late interaction with one token per side is a bi-encoder.
Matryoshka Retrieval: One Embedding, Many Sizes
Every dense embedding so far has been one fixed width — the bi-encoder emits a 32-dim vector and search is a dot product over all 32 numbers. But storing and comparing millions of full-width vectors is expensive, and different jobs want different budgets: a fast first pass would love a tiny vector, a careful re-rank wants the full one. What if a single trained embedding could be any size you ask for at query time?
Naively, you’d just chop off the tail — keep the first m numbers. For an ordinary encoder that is a disaster. Its training never asked the leading dimensions to mean anything on their own, so information is smeared across all 32 with no ordering; slice off the tail and you throw away meaning at random. Retrieval falls apart.
Matryoshka Representation Learning (MRL) — named after the Russian nesting dolls — trains the encoder so its prefixes are themselves good embeddings. The first 4 numbers are a coarse embedding, the first 8 a finer one, the full 32 the finest, each nested inside the next. One model, many sizes; to shrink a vector you just slice and renormalize. This is exactly what OpenAI’s text-embedding-3 (its dimensions parameter), Nomic Embed v1.5, and Google’s gemini-embedding ship — you pick the width, the API hands back a prefix.
NoteKey Insight
Every doll is a prefix of the next — the 4-dim embedding is the first 4 numbers of the 8-dim one, which is the first 8 of the 16-dim one, and so on. So one stored full-width vector is every smaller vector at once. The magic isn’t in the slicing (anyone can slice); it’s that MRL trained those leading dimensions to be meaningful on their own.
The Math: A Loss on Every Prefix
An ordinary bi-encoder minimizes one contrastive loss on the full-width vectors. MRL minimizes a sum of that same loss over a set of nested prefix widths M = \{4, 8, 16, 32\} — the granularities:
\mathcal{L}_{\text{MRL}} = \sum_{m \in M} c_m \cdot \text{InfoNCE}\big(\text{trunc}(q, m),\ \text{trunc}(d, m)\big)
where (as in the paper) every weight c_m = 1, and \text{trunc}(v, m) takes the first m dimensions and L2-renormalizes. Each term pushes a different prefix to put every query’s own document on the diagonal. The leading dimensions, supervised by every term, are forced to carry signal that survives truncation — the coarse-to-fine ordering the naive encoder never had.
Truncation itself is the whole trick, and it is tiny — slice, then renormalize (the slice changes a unit vector’s length, so you must rescale before comparing by cosine):
import torch
from matryoshka import truncate
v = torch.tensor([3.0, 4.0, 12.0, 0.0]) # a full 4-dim vector
print("first 2 dims, renormalized:", truncate(v, 2).tolist()) # [0.6, 0.8]
print("norm after truncation:", round(float(truncate(v, 2).norm()), 4))first 2 dims, renormalized: [0.6000000238418579, 0.800000011920929]
norm after truncation: 1.0matryoshka_loss is a one-line sum of the m08-era info_nce_loss over default_granularities(dim), and train_matryoshka is train_bi_encoder with that loss swapped in — same architecture, same corpus, only the objective differs. That is the fair comparison the payoff below rests on.
from matryoshka import default_granularities, train_matryoshka
from dense import synonym_dataset
dims = default_granularities(32) # [4, 8, 16, 32]
mrl_model, _ = train_matryoshka(synonym_dataset(), dim=32, dims=dims, epochs=200, seed=0)
print("granularities supervised:", dims)granularities supervised: [4, 8, 16, 32]The Payoff: Truncate and Barely Notice
Here is the headline. Two encoders trained on the same synonym corpus with the same seed — one with the Matryoshka loss, one plain — evaluated at each truncation width by top-1 retrieval accuracy (does each query retrieve its own document?). The full-width column ties; watch what happens as you slice.
NoteKey Insight
The plain model isn’t worse — at full width it’s identical. It simply never learned an ordering. MRL’s summed-prefix loss is the entire difference, and it buys you a dial: trade dimensions for accuracy at query time, from one stored vector, with no re-encoding and no re-indexing.
Adaptive Retrieval: A Cheap Shortlist, an Exact Answer
Graceful truncation unlocks a two-stage search the paper calls Adaptive Retrieval. Scan the whole corpus with a tiny prefix (cheap), keep a shortlist, then rescore just that shortlist with the full-width vector (accurate). Because both stages read prefixes of the same stored embedding, there’s nothing extra to store. Here the probe query is “”:
TipTry This
- Break the ordering. Train with
train_bi_encoder(plain) instead oftrain_matryoshka, then callretrieval_accuracy_by_dim— watch the 4- and 8-dim accuracy fall while full width stays perfect. - Dial the width.
MatryoshkaRetriever(mrl_model).fit(docs)then.search("cash money", dim=4)vsdim=32. On this corpus the top-1 is the same at 4 dims as at 32 — the stored vectors never change, only how many you compare. - Move the funnel. In
adaptive_retrieval, shrinkshortlist_kto 1 and confirm it can now miss the gold if stage 1’s coarse prefix ranked it second — the shortlist must be wide enough to catch what the rescore will promote.
Scaling Search: Approximate Nearest Neighbors
Every retriever in this module — TF-IDF, the dense bi-encoder, ColBERT, Matryoshka — ends the same way: score the query against every document, keep the top-k. That exhaustive scan is O(N) distance computations per query. At the toy corpora here it is instant. At a hundred million vectors — a real document store, a web index, a chatbot’s memory — comparing against all of them for every query is hopeless. The last line of this module hands you off to FAISS as “the library that makes cosine top-k fast.” This section is what FAISS is doing.
The move is to give up exactness for a huge speedup: return the approximate nearest neighbors — almost always the true ones — while touching a tiny, navigated fraction of the corpus. The structure the whole field converged on is HNSW (Hierarchical Navigable Small World graphs), the index inside FAISS, hnswlib, Qdrant, Weaviate, Milvus, and pgvector. It turns an O(N) scan into an O(\log N) walk.
Intuition: A Graph You Walk Downhill
Forget the flat list of vectors. Instead, build a graph: each vector is a node, linked to a few of its nearest neighbors. To find the neighbors of a query, you no longer check everyone — you walk downhill. Start at some node, look at its linked neighbors, step to whichever is closer to the query, and repeat. Each hop lands you nearer; a well-connected graph reaches the query’s neighborhood in a handful of hops instead of N comparisons.
One flat graph has a flaw: the greedy walk gets stuck in local minima and needs long “shortcut” edges to cross the space. HNSW’s fix is a hierarchy — like a skip list built over proximity graphs. Most nodes live only on the crowded bottom layer; exponentially fewer are promoted to each layer above. Search enters at the sparse top, where a few long hops cross the whole space fast, greedily descends to the nearest node it can find there, drops down a layer, and repeats. Only at the dense bottom does it open a real, wider beam to collect the final neighbors. The top layers are the express train; the bottom layer is the local stop.
Drive the bottom-layer walk below. The star is the query; grey dots are document vectors; faint lines are the graph’s edges. Step the search: the highlighted node is the one being examined, its neighbors get evaluated, the filled nodes are the best found so far (the beam of width ef), and the ringed nodes are the true nearest neighbors we hope to land on. Watch the beam hop across edges and close in — checking a fraction of the dots, never all of them.
NoteKey Insight
Nearest-neighbor search becomes graph navigation. You never compute the query’s distance to most vectors — you only ever evaluate the neighbors of nodes on your path. The hierarchy is what keeps that path short: a logarithmic number of hops, not a linear scan. The price is that it is approximate — a greedy walk can miss a true neighbor tucked behind a gap in the graph — and the search-beam width ef is the dial that buys the recall back.
The Math: Layers, Degree, and the log N Promise
Two parameters shape the graph. M is how many neighbors each node links to per layer (with a wider cap Mmax0 = 2M on the crowded bottom layer, where good connectivity matters most). ef is the width of the search beam — how many candidate-nearest nodes the search keeps alive at once. Bigger ef explores more of the graph: higher recall, more distance computations.
When a node is inserted, its top layer is drawn at random with an exponentially decaying distribution:
\ell = \big\lfloor -\ln(U) \cdot m_L \big\rfloor, \qquad U \sim \text{Uniform}(0, 1],
so layer 0 holds everyone and each higher layer is a factor e^{-1/m_L} rarer. The paper’s tuned choice is
m_L = \frac{1}{\ln M},
which makes that thinning factor exactly e^{-\ln M} = 1/M — one layer up, a 1/M slice of the nodes survives. That is what makes the top layers sparse enough to cross the space in long hops while the bottom stays dense enough to pin down the true neighbors. The number of layers grows like \log_M N, and the greedy descent does a bounded amount of work per layer, so search costs
O(\log N)\ \text{distance computations},
against the O(N) of the exhaustive scan — the entire reason the structure exists. Watch the exponential thinning: with M = 8, each layer up keeps about 1/8 of the nodes below it.
Code: Build and Search an HNSW Index
ann.py implements the whole index from scratch — layer assignment, the greedy search-layer beam, insertion with neighbor selection, and the layered search — faithful to Malkov & Yashunin’s Algorithms 1, 2, 3, and 5. Every distance it evaluates is counted, so we can measure the cost against a brute-force scan. Build an index and query it:
import torch
from ann import HNSW, brute_force_knn
torch.manual_seed(0)
data = torch.randn(2000, 32) # 2000 vectors in 32-dim space
query = torch.randn(32)
index = HNSW(M=8, ef_construction=48, seed=0).build(data)
approx = index.search(query, k=5, ef=32)
index.distance_computations = 0 # reset the counter, then search once more
index.search(query, k=5, ef=32)
exact = brute_force_knn(query, data, k=5)
print("HNSW top-5 :", [i for i, _ in approx])
print("exact top-5 :", [i for i, _ in exact])
print(f"HNSW touched {index.distance_computations} of {len(data)} vectors "
f"({100 * index.distance_computations / len(data):.1f}%)")HNSW top-5 : [1441, 1930, 124, 484, 115]
exact top-5 : [1441, 1930, 124, 484, 115]
HNSW touched 366 of 2000 vectors (18.3%)The graph found (nearly) the exact neighbors while comparing against a small fraction of the corpus. The approximation is real — at a narrow beam the greedy walk can miss a true neighbor — but widening ef closes the gap, and in the limit it is provably exact: give the beam enough width to reach every node and HNSW returns the same top-k as the exhaustive scan.
# With a wide enough beam, "approximate" becomes exact.
small = torch.randn(80, 16)
idx = HNSW(M=8, ef_construction=64, seed=1).build(small)
q = torch.randn(16)
wide = [i for i, _ in idx.search(q, k=5, ef=80)] # ef = N: full exploration
truth = [i for i, _ in brute_force_knn(q, small, k=5)]
print("wide-beam HNSW == brute force:", wide == truth)wide-beam HNSW == brute force: TrueThe Payoff: Recall for a Fraction of the Work
The dial is ef. Sweep it and two curves move in opposite ways: recall@k (how many true neighbors you recovered) climbs toward 1.0, while the distance computations per query rise — but stay far below the flat line the brute-force scan pays every single query. That gap is the win: near-exact answers for a small fraction of the work.
TipTry This
- Trade recall for speed. Call
demonstrate_hnsw(ef_values=(3, 5, 10, 50))and read the recall anddistance_computationscolumns — the narrow beam is fast but misses neighbors; the wide beam is near-exact but pricier. - Make it exact. Build an
HNSWand search withef = N(the corpus size); confirm its top-k equalsbrute_force_knn. Then dropeftokand watch a true neighbor slip out — that is the approximation, made visible. - Change the degree. Rebuild with
M=4vsM=16at the sameef. More links per node means a better-connected graph (higher recall) but a larger index and slower inserts — the build-time/query-time trade behind every vector database.
Common Pitfalls
- Chunk too big or too small. Whole-document chunks retrieve nothing specific; one-sentence chunks lose context. A few hundred words with light overlap is the usual sweet spot.
- Neither retriever alone is enough. Lexical (TF-IDF / BM25) misses paraphrases — “car” won’t retrieve “automobiles”, cosine exactly zero — and dense misses exact literals it never learned (error codes, hashes, names). The fix is not to pick one but to run both and fuse (the hybrid-retrieval section): BM25 for surface form, a dense encoder for meaning, RRF to let each cover the other’s blind spot.
- Fusing padded rankings, not candidates. Feed a retriever’s full ranking into RRF even when it has no real signal (all scores zero, an arbitrary tie order) and its noise outvotes the other retriever. Fuse only positive-evidence hits.
- Retrieving duplicates. Top-k over a corpus with repeats fills the context with the same fact. De-duplicate (MMR) or the model sees no new information.
- Stuffing too much context. More passages is not better — irrelevant context distracts the model and costs tokens. Retrieve few, retrieve well.
- Trusting retrieval blindly. If the answer isn’t in the corpus, a grounded prompt should let the model say so — the
build_promptinstruction asks for exactly that. RAG reduces hallucination; it does not abolish it. - Re-ranking can’t fix a bad shortlist. A cross-encoder only reorders what the first stage retrieved. If the gold document never made the top-N shortlist, no re-ranker can recover it — re-ranking raises precision, not recall. Size the first-stage
poolfor recall first; re-rank second. - Late interaction’s index is bigger. ColBERT stores one vector per token, not one per document, so its index is far larger than a bi-encoder’s and MaxSim is a per-token max, not a single dot. The precision is real, but so is the memory and compute — ColBERTv2’s residual compression exists precisely to shrink that index. And with a context-free encoder MaxSim is a max over a set, so it is still order-blind; the contextual token encoder is what lets it read phrases.
- Approximate means approximate. An HNSW index does not guarantee the exact top-k. At a narrow beam
efthe greedy walk can miss a true neighbor hidden behind a gap in the graph, and recall drops silently — the wrong document simply never gets scored. Tuneef(andM) for the recall you need and measure it against a brute-force scan on a held-out sample; never assume recall is 1.0 just because results came back.efmust be at leastk.
Exercises
Exercise 1: idf by hand
from retrieval import build_vocab, compute_idf
# For docs ["the cat sat", "the dog ran", "the cat ran"], compute idf by hand
# (N=3) for "the" (df=3), "cat" (df=2), "sat" (df=1). Confirm compute_idf agrees
# and that the rarest term has the highest idf.
# Your implementation here:Exercise 2: retrieve then augment
from retrieval import Retriever, build_prompt
# Fit a Retriever on 4–5 facts of your own, retrieve the top-2 for a question, and
# build_prompt them. Print the prompt and check the right facts were pulled in.
# Your implementation here:Exercise 3: when does MMR help?
from retrieval import tfidf_matrix, encode_query, mmr
# Build a corpus with two near-duplicate relevant docs and one different relevant
# doc. Show that lambda=1.0 returns both duplicates but lambda=0.2 returns a
# diverse pair.
# Your implementation here:Exercise 4: dense beats lexical
from retrieval import Retriever, cosine_similarity
from dense import synonym_dataset, train_bi_encoder, DenseRetriever
# Take a (query, doc) synonym pair from synonym_dataset() whose words are disjoint.
# 1. Confirm the TF-IDF cosine between the query and its correct doc is exactly 0.
# 2. Train a bi-encoder, then show DenseRetriever ranks that doc first.
# 3. Bonus: raise the InfoNCE `temperature` (e.g. 0.5) and watch the loss plateau
# higher — the diagonal never fully separates.
# Your implementation here:Exercise 5: hybrid beats either alone
from hybrid import BM25Retriever, HybridRetriever, reciprocal_rank_fusion
# 1. Confirm BM25's saturation: bm25_scores("x", ["x x x x x x x x"]) is higher than
# bm25_scores("x", ["x"]) but LESS than 8x it (sub-linear, unlike TF-IDF).
# 2. Fuse two rankings by hand with reciprocal_rank_fusion and verify a doc ranked
# #1 by BOTH beats a doc ranked #1 by only one (compute 1/(60+rank) yourself).
# 3. Build a 2-doc corpus: one about "car vehicle", one containing a rare literal
# like "e4021". Show BM25 finds the literal (score > 0) and scores the paraphrase
# query "automobile" at exactly 0 — the blind spot a dense retriever would cover.
# Your implementation here:Exercise 6: retrieve, then re-rank
from retrieval import Retriever
from rerank import train_cross_encoder, retrieve_then_rerank, reversal_corpus, permutation_tie
# The reversal corpus pairs each fact with its word-reversal (same bag of words).
# 1. Confirm the TF-IDF first stage is permutation-blind: permutation_tie(...) is True
# for a fact and its reversal (their encodings are bit-for-bit equal).
# 2. Train a CrossEncoder on (query, fact, reversal) triples, then retrieve_then_rerank
# a fact query and show the correctly-ordered fact is now rank 1.
# 3. Bonus: re-rank with cross=False (the bi-encoder ablation) and compare — how much
# of the win came from the query<->document attention edges?
# Your implementation here:Exercise 7: late interaction saturates
from colbert import maxsim, identity_token_matrix, sum_pool_score
from retrieval import build_vocab
# Query "alpha beta"; a gold doc covering both aspects vs a distractor piling on one.
# 1. With identity_token_matrix vectors, show MaxSim scores the gold 2.0 and the
# "alpha alpha alpha" distractor 1.0 (repetition earns no extra credit).
# 2. Show sum_pool_score ranks them the other way (the pooled cosine rewards the
# repeated aspect) — the exact failure MaxSim fixes.
# 3. Bonus: pool each side to ONE vector and confirm maxsim equals their cosine —
# the bi-encoder is late interaction with a single token per side.
# Your implementation here:Exercise 9: walk the graph, trade recall for cost
from ann import HNSW, brute_force_knn, ann_recall
import torch
# 1. Build an index over 500 random 16-dim vectors, then for one query compare
# HNSW's top-10 at ef=10 vs ef=100 against brute_force_knn. Print ann_recall and
# index.distance_computations for each (reset the counter before each search) —
# the wide beam should score higher recall at a higher cost, both well under 500.
# 2. Make it exact: search with ef equal to the corpus size and confirm the result
# equals brute_force_knn. Then set ef=k and find a query where a true neighbor
# drops out — that is the approximation.
# 3. Bonus: rebuild with M=4 and M=16 at a fixed ef and compare recall — a
# better-connected graph recovers more, at the price of a larger index.
# Your implementation here:Exercise 8: slice, renormalize, and break it
from matryoshka import truncate, train_matryoshka, retrieval_accuracy_by_dim, default_granularities
from dense import synonym_dataset, train_bi_encoder
# 1. Confirm truncation is a *slice*: truncate(v, 2) keeps v's first two numbers
# then renormalizes — check truncate(torch.tensor([3.,4.,12.,0.]), 2) == [0.6, 0.8].
# 2. Train a Matryoshka model and a plain train_bi_encoder on synonym_dataset() with
# the SAME seed. Print retrieval_accuracy_by_dim for both over default_granularities(32).
# The plain model's 4-dim accuracy should fall well below the Matryoshka model's 1.0
# while both stay 1.0 at full width.
# 3. Bonus: raise `smallest` in default_granularities so the coarsest supervised prefix
# is 16, not 4 — retrain, and watch the 4-dim accuracy sag (nothing trained it).
# Your implementation here:Summary
Key takeaways:
- RAG conditions generation on retrieved text — encode → search → retrieve → augment → generate — so the model answers from a live corpus, not just frozen weights.
- Retrieval is nearest-neighbor search — encode documents and the query into one vector space and rank by cosine similarity.
- TF-IDF is the from-scratch encoder — term frequency times smoothed idf makes shared rare words count; dense embeddings are the same pipeline with a learned encoder.
- Chunk long documents into small overlapping passages so each retrieved unit is focused.
- MMR buys diversity — pick passages for relevance minus redundancy so the context isn’t three copies of one fact.
- RAG’s win is grounding — fresh, private, citable knowledge with no retraining, and a real dent in hallucination.
- Dense retrieval matches meaning, not words — a learned bi-encoder trained with the in-batch-negatives contrastive loss (the same objective as CLIP) retrieves a paraphrase that shares no token with the query, which TF-IDF scores at exactly zero. Only the
encodestep changes; the rest of the pipeline is identical. - Hybrid retrieval fuses the two — BM25 (saturating, length-normalized lexical scoring) covers exact literals, a dense encoder covers meaning, and Reciprocal Rank Fusion (
Σ 1/(k+rank), scale-free,k=60) combines their rankings. On a mixed corpus each retriever alone hits recall 0.5 and the fusion recovers the union at 1.0 — the reason production RAG runs both. - Re-ranking scores the shortlist jointly — a cross-encoder reads the
(query, document)pair together (query tokens attending to document tokens), so it catches order and binding a bi-encoder’s independent pooling averages away — but it can’t precompute, so it costs one forward pass per pair and only ever re-ranks the first stage’s top-N. Retrieve for recall, re-rank for precision. A permutation-blind first stage ties a fact with its reversal (recall@1 = 0.5); the order-aware re-ranker lifts it to 1.0. Thecrossflag is the whole story: mask the query↔︎document edges and the cross-encoder degenerates into a bi-encoder. - Late interaction is the precomputable middle ground — ColBERT keeps one vector per token and scores by MaxSim (\sum_i \max_j E_{q_i}\cdot E_{d_j}): the encoders stay independent (so documents index offline, like a bi-encoder) but every query token finds its best document token (fine-grained, like a cross-encoder). A bi-encoder is the one-token-per-side special case. MaxSim saturates — capping each query token at its single best match — so it wins by covering every query aspect where a pooled vector is fooled by repetition (recall 1.0 vs 0.0 on the coverage corpus).
- Matryoshka embeddings size on demand — summing the contrastive loss over nested prefixes M = \{4, 8, …, d\} (c_m = 1) trains the leading dimensions to be meaningful alone, so you slice-and-renormalize one stored vector to any width. On the synonym corpus the Matryoshka encoder retrieves every concept from its 4-dim prefix (accuracy 1.0) where the plain encoder — identical at full width — slides toward chance; Adaptive Retrieval (shortlist at a short prefix, rescore at full width) then returns the same top result as a full-width scan for a fraction of the comparisons. This is the
dimensionsdial in OpenAI’stext-embedding-3, Nomic Embed v1.5, and Gemini embeddings. - HNSW makes search scale — every retriever above scores all
Ndocuments, anO(N)scan that dies at real scale. Hierarchical Navigable Small World graphs turn nearest-neighbor search into graph navigation: link each vector to a few near neighbors and walk downhill toward the query, in a hierarchy of graphs (layer \ell = \lfloor -\ln(U)\,m_L\rfloor, m_L = 1/\ln M) that keeps the walk to O(\log N) hops. The search beameftrades recall for cost — wide enough it recovers the exact top-k, narrow it is fast but approximate — while touching a small fraction of the corpus instead of all of it. This is the index inside FAISS, hnswlib, Qdrant, Weaviate, Milvus, and pgvector.
What’s Next
You can now ground a model in an external corpus. The natural next step is letting the model act: tool use and agents, where retrieval becomes one action among many (search, run code, call an API) inside a plan-act-observe loop — the same “put the right thing in the context” idea, now driven by the model itself.
Going Deeper
Core Papers:
- Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks — Lewis et al. (2020), the paper that named RAG and combined a retriever with a generator end-to-end.
- Dense Passage Retrieval for Open-Domain Question Answering — Karpukhin et al. (2020), the two-tower encoder with in-batch negatives built in the dense-retrieval section; beat BM25 by 9–19% on retrieval.
- Passage Re-ranking with BERT — Nogueira & Cho (2019), the BERT cross-encoder re-ranker the re-ranking section builds in miniature: joint
[query | passage]scoring over a first-stage shortlist. - ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT — Khattab & Zaharia (2020), the late interaction section built here: per-token vectors (precomputable like a bi-encoder) scored by token-level MaxSim (interaction like a cross-encoder).
- ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction — Santhanam et al. (2021), residual compression + denoised supervision that shrink the per-token index late interaction pays for.
- Matryoshka Representation Learning — Kusupati et al. (NeurIPS 2022), the nested-prefix loss the Matryoshka section builds: one encoder, many embedding sizes, sliced-and-renormalized, with up to 14× smaller vectors at the same accuracy and Adaptive Retrieval for large-scale speedups.
- Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks — Reimers & Gurevych (2019), the mean-pooling bi-encoder this section’s encoder mirrors.
- Representation Learning with Contrastive Predictive Coding — van den Oord et al. (2018), the InfoNCE loss behind the contrastive training.
- REALM: Retrieval-Augmented Language Model Pre-Training — Guu et al. (2020), retrieval baked into pre-training.
Practical Resources:
- The Probabilistic Relevance Framework: BM25 and Beyond — Robertson & Zaragoza (2009), the BM25 score and IDF built in the hybrid-retrieval section.
- Reciprocal Rank Fusion outperforms Condorcet and individual Rank Learning Methods — Cormack, Clarke & Büttcher (SIGIR 2009), the one-line rank-fusion the hybrid retriever uses, and the
k=60constant. - Efficient and Robust Approximate Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs — Malkov & Yashunin (2016; IEEE TPAMI 2020), the HNSW index the scaling-search section builds from scratch: a layered proximity graph greedily walked in
O(log N)hops, with theM/ef/mL = 1/ln Mparameters used here. Its predecessor is Malkov et al., Approximate nearest neighbor algorithm based on navigable small world graphs (Information Systems, 2014). - FAISS — the library that makes cosine top-k fast over millions of vectors, with HNSW (the section above) and IVF/PQ indexes.