ML System Design · ML Systems Lab

RAG: Retrieval-Augmented Generation from Architecture to Production

LLMs hallucinate when they do not know the answer. RAG fixes this by retrieving relevant documents at inference time and conditioning the LLM's generation on them. The architecture is: encode the query, retrieve from a vector store, prepend retrieved chunks to the prompt, generate. Simple in principle; full of engineering decisions in production. This is the full system: chunking, embedding, retrieval, re-ranking, and generation.

Large language models have a fundamental limitation: their knowledge is frozen at training time. They cannot answer questions about events after their training cutoff, cannot access proprietary internal documents, and hallucinate when asked about facts they are uncertain about. Retrieval-Augmented Generation (RAG, Lewis et al., 2020) addresses this by augmenting the LLM with a retrieval system that fetches relevant documents at inference time.

The basic RAG pipeline

Query → encode query → retrieve top-k chunks → prepend chunks to prompt → LLM generates answer conditioned on chunks. The retrieval component is a semantic search system (see Posts 70 and 79). The generation component is any capable LLM. The simplest RAG implementation is a few hundred lines of code using LangChain or LlamaIndex. Production RAG has a dozen more engineering decisions.

Chunking: the overlooked critical step

Documents must be split into chunks before embedding. Chunk size determines what the retrieval system can find and what fits in the LLM context. Too small (< 100 words): chunks lack enough context for the LLM to generate coherent answers. Too large (> 1000 words): embedding a long chunk averages its meaning, retrieving less precisely. Standard: 200-500 word chunks with 50-100 word overlap between adjacent chunks (sliding window) so that answers spanning chunk boundaries are not lost.

Semantic chunking: split on natural semantic boundaries (paragraph breaks, section headers, sentence boundaries) rather than fixed token counts. Hierarchical chunking: embed both small chunks (for precise retrieval) and large parent chunks (for full-context generation); retrieve small, generate from large.

Embedding model choice

The quality of retrieval depends entirely on the embedding model mapping both queries and document chunks into a shared semantic space. Bi-encoder models (sentence-transformers, OpenAI embeddings, Cohere embeddings) encode queries and documents independently. Quality varies significantly: MTEB benchmark measures retrieval quality across 56 datasets. For domain-specific retrieval (medical, legal, code), fine-tune a general embedding model on in-domain (query, relevant document) pairs.

Vector stores: storing and querying embeddings at scale

Chunk embeddings are stored in a vector store with ANN (approximate nearest neighbour) indexing. Open-source options: FAISS (Facebook, in-memory, fastest), Chroma (simple, good for development), Weaviate, Qdrant (production-grade with persistence and filtering). Managed options: Pinecone, Weaviate Cloud. Key capabilities: metadata filtering (retrieve only documents from the last 30 days, or only from a specific department), hybrid search (combine dense vector retrieval with keyword BM25 filtering), real-time upsert (update embeddings when source documents change).

Re-ranking: the quality multiplier

Initial vector retrieval (top-k = 50-100) is fast but imprecise. A re-ranking step scores each retrieved chunk against the query more carefully. Cross-encoder re-rankers (e.g., Cohere Rerank, BGE Reranker) take the concatenated (query, chunk) pair and produce a relevance score — much more accurate than the cosine similarity of independently encoded vectors. Re-rank the top-50 retrieved chunks with a cross-encoder, then pass the top-5 to the LLM.

The generation step: prompting the LLM

Prompt structure: system message (defines the assistant's role and rules, e.g., "You are a helpful assistant. Answer only based on the provided context. If the answer is not in the context, say 'I don't know.'"), context (retrieved chunks, clearly demarcated), user query. The context injection location matters — LLMs attend more strongly to context at the beginning and end of the prompt than in the middle (the "lost in the middle" problem). Place the most relevant chunks at the beginning or end.

Failure modes in production RAG

Retrieval failure: the relevant document is not retrieved. Cause: wrong embedding model, wrong chunk size, query not matching document vocabulary. Fix: hybrid retrieval (BM25 + dense), query rewriting. Context overflow: retrieved chunks exceed the LLM context window. Fix: hierarchical summarisation, reduce top-k. Hallucination despite retrieval: the LLM ignores the retrieved context. Fix: stronger system prompt, use an instruction-tuned model, implement citation checking. Stale knowledge: documents updated but embeddings not refreshed. Fix: document change detection pipeline with incremental re-embedding.

Try on Colab: build a RAG system over a set of Wikipedia articles (50 articles on a consistent topic). Use sentence-transformers for embeddings, FAISS for the vector store, and GPT-3.5 or Claude Haiku for generation. Ask factoid questions. Compare: (1) LLM without retrieval (baseline, observe hallucination), (2) LLM with top-3 chunk retrieval, (3) LLM with top-50 retrieval + cross-encoder re-ranking. Measure answer accuracy. Add citation (require the LLM to cite the specific chunk) and measure citation accuracy.

Continue interactively
Read this post inside ML Systems Lab — with Simplify toggle, interview Q&As, inline glossary, and the MLE Path forward pointer.
Open in MSL →