Advanced retrieval and RAG design
The pipeline of §7.6 handles a surprising share of real questions. It does one retrieval with the user's words, keeps the best few chunks, and asks the model to answer from them. But some questions break it in ways no amount of tuning fixes.
- "And what about the blue one?" The question only makes sense with the conversation before it.
- "Where did the founder of the company that makes the Aster kettle grow up?" The answer needs two retrievals, and the second depends on the first.
- "What are the main complaints across our 50,000 support tickets this year?" No single chunk contains the answer. It's a summary of everything.
- "What's the policy?" asked of a corpus with a dozen policies. The query is too vague to embed well.
This section is a toolbox for those cases, ordered from cheap to expensive. We'll rewrite and decompose queries, embed hypothetical answers instead of questions (HyDE), index documents at several granularities, build a knowledge graph for global questions (GraphRAG), follow a question across several retrievals, and preview what happens when the model itself drives the search. Then we'll settle the question everyone asks now that context windows reach a million tokens. Why retrieve at all? We'll answer with numbers, and finish with caching.
Many of these tools call a language model to rewrite a query, write a passage, summarize a section, or extract entities. Read "LLM" in this section as an instruction-tuned model, the kind Chapter 14 builds, which does such jobs from a one-line instruction. A base model can do several of them with a few solved examples in the prompt (§6.6), less reliably. The retrieval mechanics around each call are what this section teaches, and every one of them runs below with a hand-written stand-in for the model.
Rewriting and decomposing the query
The cheapest improvement is to fix the query before it reaches the retriever, usually with one call to a small LLM.
Conversational rewriting. Turn a follow-up into a standalone question using the chat history. "And what about the blue one?" becomes "What is the warranty on the blue Aster kettle?". Without this step, multi-turn RAG retrieves on fragments.
Rewrite, retrieve, read. Ma et al. (2023) trained a small model to rewrite queries for a frozen retriever and reader, using the final answer's quality as the reward. Even a prompted rewrite helps by expanding abbreviations, fixing typos, and adding the domain's vocabulary.
Multi-query expansion. Generate several paraphrases of the question, retrieve with each, and fuse the lists with RRF (§7.5). A document that several phrasings agree on rises to the top. Query2doc (Wang et al., 2023) goes further, having an LLM write a pseudo-document that answers the question and appending it to the query. That improved BM25 by 3 to 15 percent on MS MARCO and TREC DL without any fine-tuning, because the pseudo-document contains the vocabulary of real answers.
Step-back prompting (Zheng et al., 2023). Ask a more general question first ("What are the physics principles behind this?") and retrieve for both the general and the specific question.
Decomposition. Split a compound or multi-hop question into sub-questions, retrieve for each, and give the model all the results. "Compare the warranty and the return policy" becomes two retrievals. For multi-hop questions, where sub-question two depends on the answer to sub-question one, decomposition has to be sequential, and that's where we cross into agentic retrieval below.
When you retrieve for several sub-questions, you also have to decide how to merge their results into one context. Taking the top 5 from a single fused list can starve one sub-question entirely. Interleaving guarantees each one is represented.
Interleave results from several sub-queries
Implement interleave(lists, k). lists holds one ranked list of document ids per sub-query. Build the final context by round robin, taking the first item of each list in order, then the second item of each, and so on, skipping documents already taken, until you have k documents or every list is exhausted. Return the list of ids.
HyDE, search with a hypothetical answer
Here's a puzzle. Questions and answers look different. "Why does my kettle click off before the water boils?" is short, interrogative, and uses the asker's words. The passage that answers it is declarative and uses the manufacturer's words. An embedding model trained on question and answer pairs learns to bridge that gap. But a model trained without labels, or used in a domain far from its training data, compares a question with passages and often finds other questions.
Hypothetical Document Embeddings (HyDE, Gao, Ma, Lin & Callan, 2022) sidestep the gap. Ask an instruction-following LLM to write a passage that answers the question. It doesn't matter if the passage is wrong in its details. Embed that passage instead of the question, and search for real passages near it. With an encoder , a query , and sampled hypothetical documents , the paper's query vector is the average
which also counts the original query as one more hypothesis. In plain words, search with what an answer would look like, not with what the question looks like.
Why doesn't the made-up content poison the search? The authors' argument is that the encoder's dense bottleneck keeps the kind of passage, its topic, vocabulary, and structure, and loses fine details like the particular numbers the LLM invented. The nearest real passages share the shape of an answer, and the false details are not matched. With InstructGPT writing the passages and the unsupervised Contriever encoding them, HyDE substantially beat Contriever alone and was competitive with fine-tuned retrievers across web search, QA, and fact verification, in several languages.
Know its limits too. HyDE adds an LLM call to every query, so latency goes up. When the LLM knows nothing about the topic (your internal product names), its hypothetical answer can point in the wrong direction. And with a strong retriever trained on query and passage pairs, which already bridges the question-answer gap, the gain usually shrinks.
The HyDE query vector
Implement hyde_vector(q_vec, hyp_vecs) where q_vec is (d,) and hyp_vecs is (N, d). Normalize each input vector, average the N + 1 unit vectors with equal weight, and return the result renormalized to unit length (so it can be used for cosine search).
Many representations of one document
Chunking forces one compromise on every document, a single size for both matching and reading. You don't have to accept it. Index a document several ways and map every hit back to the text you actually want the model to read. This is multi-vector retrieval at the level of the index, several vectors per document, each pointing back to the same text, the coarse cousin of ColBERT's one vector per token (§7.3).
- Parent and child chunks. Index small child chunks (a few sentences) for sharp matching, and when a child matches, return its larger parent (the section or page) to the model. The match is precise and the context is complete. Practitioners report that this "small to big" pattern is one of the most effective single changes to a basic pipeline.
- Summaries. Index an LLM-written summary of each section or document alongside its chunks, so that broad questions can match the summary.
- Hypothetical questions. Generate, for each chunk, the questions it answers, and index those. It's the mirror image of HyDE, done once at indexing time instead of per query. The idea goes back to doc2query (Nogueira et al., 2019), which appended predicted queries to documents to help BM25.
- Recursive summary trees. RAPTOR (Sarthi et al., 2024) clusters chunk embeddings, summarizes each cluster, then clusters and summarizes the summaries, building a tree from details up to themes. Retrieval searches all levels at once, so a question can match a detail or a theme. Coupled with GPT-4, it improved the best reported accuracy on the QuALITY long-document benchmark by 20 points.
Parent-child retrieval has one subtle step. Several children of the same parent often match, so you have to aggregate their scores per parent and deduplicate before filling the budget.
Parent-child retrieval
Implement parents_for(child_scores, parent_of, parent_tokens, budget).
child_scoresis a dict from child chunk id to retrieval score.parent_ofmaps each child id to its parent id, andparent_tokensmaps each parent id to its length in tokens.- Score each parent by the maximum score of its children. Rank parents by that score (descending, ties broken by parent id), then keep them in order as long as they fit in the token budget, skipping any parent that doesn't fit.
Return the list of kept parent ids.
GraphRAG, for questions about the whole corpus
Vector RAG retrieves passages similar to the question. That fails completely for global questions such as "What are the main themes in this dataset?". No chunk is similar to that question, and the answer is spread across everything. Edge et al. (2024) framed this as query-focused summarization over a corpus and proposed GraphRAG, which builds a structure at indexing time that makes global summaries possible.
The pipeline has four stages.
- Extract a graph. An LLM reads every chunk and extracts entities (people, organizations, products, concepts), relationships between them with short descriptions, and optionally claims. Entities become nodes, and relationships become edges, weighted by how often they're observed.
- Find communities. Run hierarchical community detection, the Leiden algorithm in the paper, to partition the graph into groups of densely connected entities, then partition each group again, giving levels from a few broad communities down to many specific ones.
- Summarize communities. An LLM writes a report for every community at every level, describing its entities, relationships, and key claims. These summaries are the index.
- Answer by map-reduce. For a global question, pick a level. In the map step, each community summary is used, independently and in parallel, to write a partial answer along with a helpfulness score. In the reduce step, the most helpful partial answers are combined into the final response.
On "global sensemaking" questions over datasets of about a million tokens (podcast transcripts and news articles), with an LLM judging pairs of answers head to head, GraphRAG beat a conventional vector RAG baseline substantially on comprehensiveness and diversity of answers. For questions about specific facts, ordinary RAG remains the better tool, and GraphRAG implementations also offer a "local" mode that starts from entities matching the question and gathers their neighbors, relationships, and source chunks.
The costs are real. Extraction calls an LLM on every chunk, and community reports add more calls, so indexing a large corpus can cost orders of magnitude more than embedding it, and the graph must be updated as documents change. It's worth it when users ask about structure ("who works with whom", "what are the themes", "how did this evolve") rather than facts.
The graph-building core is classic algorithms, and it makes a good interview problem. Here we use co-occurrence within a chunk as the relationship, and connected components (with a minimum edge weight) as a simple stand-in for Leiden.
Build an entity graph and find its communities
Implement two functions.
entity_graph(chunk_entities)takes a list of entity lists, one per chunk, and returns a dict mapping each unordered pair of distinct entities that co-occur in a chunk to the number of chunks where they co-occur. Keys are tuples(a, b)witha < b. Count each pair at most once per chunk, even if an entity is repeated in it.communities(edges, min_weight)keeps the edges with weight at leastmin_weightand returns the connected components of the resulting graph, using union-find. Only entities touched by a kept edge appear. Each community is a sorted list, and the list of communities is sorted by size (largest first), then by its first entity.
Here's the map-reduce step in miniature. Each community gets a one-line summary (written by hand here, by an LLM in the real system), a stand-in "map" scores each summary's helpfulness for the question by word overlap, and "reduce" keeps the most helpful partial answers.
The real system replaces both stand-ins with LLM calls, and it's the community structure that makes the map step possible. There are a few hundred summaries to read instead of a million tokens, and each summary already distills the relationships the extraction found.
Multi-hop retrieval, and a preview of agentic search
Every method so far retrieves on a fixed schedule decided by the engineer, once per question. Multi-hop questions break that, because the second query depends on what the first one found. The fix is a loop. Retrieve, read the result, form the next query from it, and repeat. Here's the idea with a tiny corpus and a hand-written policy standing in for the model that reads.
One-shot retrieval finds only the kettle passage, the one that shares the question's words. The passage about the founder's childhood shares none of them, so no single query built from the question can reach it. The loop reads each result, extracts the entity the next hop needs, and reaches the answer in three steps. In a real system a language model does the reading and chooses the next query. The classic recipe is IRCoT (Trivedi et al., 2023), which interleaves retrieval with chain-of-thought (§18.1) and uses each new reasoning sentence as the next query.
Push the idea further and the model decides whether to retrieve at all, what to search for, and when it has enough, calling search as a tool in a loop. That is an agent, and the loop of reasoning and actions behind it is developed with reasoning and acting agents, which treat search as one tool among many. Two research lines are worth knowing by name now. Self-RAG (Asai et al., 2023) fine-tunes a model to emit special tokens that decide when to retrieve and whether a passage supports its output, and search agents such as Search-R1 (Jin et al., 2025) learn when and what to search by reinforcement learning on answer correctness, the verifiable-reward recipe of §18.3. The price of all this flexibility is latency (several model calls and searches per question), unpredictable cost, and new failure modes such as loops and drifting off the question. Bound the number of steps, log every step, and evaluate end to end on questions that actually need multiple hops.
Long context or retrieval?
Context windows have grown from 4K tokens to a million and beyond. If the whole corpus fits in the prompt, why not put it all in and skip retrieval? Sometimes you should. But it's a quantitative decision, and four things decide it.
Cost. Input tokens are billed per query. Putting a 500K-token corpus into every prompt costs about 100 times more per query than a RAG prompt of 5K tokens, at the same price per token. Prompt caching changes this (see below), but only for a static prefix.
Latency. Before the first output token, the model must process the entire prompt, the prefill phase of §6.5. Prefill time grows at least linearly with prompt length, and attention makes it worse than linear at long lengths, because every token attends to all the tokens before it. For a sense of scale, Yang et al. (2024) prefilled one million tokens through Llama 3 405B in 77 seconds by spreading the work over up to 128 H100 GPUs with context parallelism. A 5,000-token RAG prompt needs at least 200 times less prefill compute on the same model, since the linear layers shrink 200 times and attention far more.
Quality. Long-context models don't use all of their context equally well. Lost in the middle (§7.6) is one symptom. RULER (Hsieh et al., 2024) found that while 17 models all claimed context sizes of 32K or more, only about half kept satisfactory performance at 32K on tasks harder than finding a single needle, which is the gap between advertised and effective context in §12.7. On the other hand, Li et al. (2024) found that when long-context models are given enough resources, they outperform RAG on average across their benchmarks, while RAG remains far cheaper. Their Self-Route method lets the model decide per query whether its retrieved context is enough or whether to fall back to the full context, keeping most of the quality at a fraction of the cost. And Xu et al. (2023) found the two combine. A 4K-context model with retrieval matched a 16K-context model without it on long-context tasks, and retrieval improved models at every context size they tried.
Scale and freshness. A company's documents are usually far bigger than any context window, change daily, and carry per-user permissions. Retrieval handles all three naturally.
| situation | prefer |
|---|---|
| a single long document (a contract, a codebase file, a book), asked many questions | long context, with prompt caching |
| a corpus far larger than the context window | RAG |
| documents change often, or access depends on the user | RAG |
| tight latency or cost budgets | RAG |
| questions that need the whole document at once (summarize all of it) | long context, or GraphRAG-style summaries |
| a small, static knowledge base and many queries | long context with a cached prefix (cache-augmented generation) |
| anything high stakes | RAG with citations, so answers can be checked |
Here's the arithmetic the widget does, with deliberately round, made-up prices. Change them to your provider's.
At these prices, a 500K-token corpus costs about 50 cents a query in a plain long-context prompt, about 5 cents with a cached prefix, and well under a cent with RAG. Caching closes most of the gap for a static corpus, which is why "put the handbook in a cached system prompt" is a legitimate design for small, stable knowledge bases. For anything large, changing, or permissioned, retrieval wins on cost by orders of magnitude.
Caching
Caching appears at every layer of a RAG system, and each layer has its own pitfall.
- Embedding cache. Store each chunk's vector keyed by a hash of its text and the model version, so re-indexing only re-embeds chunks that changed. Changing the model invalidates everything (§1.8).
- Retrieval and answer cache. Exact-match caching of the normalized query string, with a time-to-live, and invalidation whenever a source document changes.
- Semantic cache. Reuse the answer of a previous query whose embedding is very close to the new one. It can save a lot on repetitive traffic, but it's dangerous for queries that differ in one crucial token. "Refund status of order 4512" and "of order 4521" embed almost identically. Use a high threshold, exclude queries with identifiers, and scope caches per user when answers depend on permissions.
- Prompt (prefix) caching, a preview. The KV cache of §6.5 holds the keys and values of every token processed so far, and a token's keys and values depend only on the tokens before it. So two requests that begin with the same tokens have identical caches for that shared beginning, and a server can keep it and skip that part of the prefill. Put the stable parts first (instructions, a fixed document set) and the variable parts last (retrieved passages, the question), and the stable prefix is computed once. Changing a single early token invalidates everything after it. Providers bill cached input tokens at a steep discount, which is what the cost cell above assumed. How servers store, share, and evict these caches belongs to serving systems with paged KV caches.
A semantic cache with a similarity threshold
Implement a class SemanticCache(threshold) with two methods. put(vec, answer) stores an answer under an embedding vector. get(vec) returns the stored answer whose vector has the highest cosine similarity with vec, provided that similarity is at least threshold, and None otherwise (including when the cache is empty). Normalize vectors so that their scale doesn't matter.
In PyTorch
Multi-query retrieval is a batched matrix multiply followed by rank fusion, and both run on the GPU.
Read more
- Gao, Ma, Lin & Callan (2022). Precise Zero-Shot Dense Retrieval without Relevance Labels (HyDE). Hypothetical document embeddings.
- Wang, Yang & Wei (2023). Query2doc: Query Expansion with Large Language Models. LLM-written pseudo-documents for sparse and dense retrieval.
- Edge et al. (2024). From Local to Global: A Graph RAG Approach to Query-Focused Summarization. Entity graphs, communities, and map-reduce answers.
- Sarthi et al. (2024). RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval. Summary trees for retrieval at several levels.
- Asai et al. (2023). Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection and Jiang et al. (2023), Active Retrieval Augmented Generation (FLARE). Deciding when to retrieve.
- Trivedi et al. (2023). Interleaving Retrieval with Chain-of-Thought Reasoning for Knowledge-Intensive Multi-Step Questions (IRCoT). Multi-hop retrieval guided by reasoning.
- Li et al. (2024). Retrieval Augmented Generation or Long-Context LLMs? A Comprehensive Study and Hybrid Approach. Self-Route, and the cost and quality trade-off.
- Hsieh et al. (2024), RULER: What's the Real Context Size of Your Long-Context Language Models?, and Yang et al. (2024), Context Parallelism for Scalable Million-Token Inference. Effective versus advertised context, and what it takes to prefill a million tokens.