Contextual Retrieval in AI Systems

Sep 19, 2024

Core Problem

Traditional Retrieval-Augmented Generation (RAG) systems often fail because they remove context when encoding information into chunks, leading to poor retrieval accuracy.

Solution: Contextual Retrieval

A method that prepends chunk-specific explanatory context to each chunk before indexing. It uses two sub-techniques:

1. Contextual Embeddings
2. Contextual BM25

Key Performance Improvements

• Reduces failed retrievals by 49% (from 5.7% to 2.9% failure rate).
• Combined with reranking, reduces failed retrievals by 67% (from 5.7% to 1.9% failure rate).

When to Use This Method

• For small knowledge bases (<200k tokens / ~500 pages): Simply include the entire knowledge base in the prompt. Use Prompt Caching to reduce latency (>2x) and cost (up to 90%).
• For larger knowledge bases: Use Contextual Retrieval for a scalable solution.

How Traditional RAG Works (and its Limitation)

1. Split documents into small chunks.
2. Create semantic embeddings (vector) and TF-IDF/BM25 (lexical) indexes.
3. At runtime, retrieve top chunks using both methods and combine results.
4. Limitation: Splitting destroys context, making chunks ambiguous (e.g., "The company's revenue grew by 3%" lacks company/time context).

Implementing Contextual Retrieval

Step 1: Generate Context for Each Chunk

Use a model (like Claude) with a specific prompt to generate concise context for each chunk.

Prompt Template:

<document>
{{WHOLE_DOCUMENT}}
</document>
Here is the chunk we want to situate within the whole document
<chunk>
{{CHUNK_CONTENT}}
</chunk>
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.

Example Transformation:

• Original Chunk: "The company's revenue grew by 3% over the previous quarter."
• Contextualized Chunk: "This chunk is from an SEC filing on ACME corp's performance in Q2 2023; the previous quarter's revenue was $314 million. The company's revenue grew by 3% over the previous quarter."

Step 2: Index the Contextualized Chunks

Prepend the generated context (~50-100 tokens) to each chunk before creating both the semantic embeddings and the BM25 index.

Cost Optimization with Prompt Caching

Using Claude's prompt caching, the one-time cost to generate contextualized chunks is $1.02 per million document tokens (based on 800-token chunks, 8k-token documents).

Further Boost: Reranking

After initial retrieval (e.g., top 150 chunks), use a reranking model (e.g., Cohere) to score and select the most relevant chunks (e.g., top 20) to pass to the model. This adds a small latency/cost trade-off but significantly improves accuracy.

Key Implementation Considerations

1. Chunk Boundaries: Experiment with chunk size, boundaries, and overlap.
2. Embedding Model: Contextual Retrieval helps all models, but Gemini and Voyage embeddings were particularly effective.
3. Custom Prompts: Tailor the contextualizer prompt to your domain (e.g., include a glossary).
4. Number of Chunks: Passing top-20 chunks to the model was most effective in tests.
5. Always Run Evals: Test performance with and without the contextual prefix.

Summary of Findings

1. Embeddings + BM25 is better than embeddings alone.
2. Voyage and Gemini have the best-performing embeddings.
3. Top-20 chunks outperform top-10 or top-5.
4. Adding context to chunks improves retrieval accuracy significantly.
5. Reranking is better than no reranking.
6. Benefits stack: For maximum performance, combine Contextual Embeddings (Voyage/Gemini) + Contextual BM25 + Reranking + Top-20 chunks.

Getting Started

Developers can deploy their own Contextual Retrieval solution using the provided cookbook.