Why Two-Stage Retrieval Works
Understand the recall-precision trade-off in single-stage retrieval and how a fast coarse retriever followed by a slow but accurate re-ranker gets the best of both worlds.
Why Two-Stage Retrieval Works is a free AI Engineering Academy lesson on CoddyKit — lesson 1 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Engineering Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.
The Recall-Precision Trade-off in Retrieval
Every retrieval system faces a fundamental trade-off: recall measures how many relevant documents you find (did you miss any?), while precision measures how accurate the top results are (how many retrieved docs are actually relevant?). Maximizing both simultaneously is computationally expensive. Fast retrievers sacrifice precision for recall; precise rankers sacrifice speed for accuracy.
Bi-Encoder vs Cross-Encoder: The Core Distinction
The two types of models at the heart of two-stage retrieval differ in how they see the query and document. A bi-encoder encodes the query and each document independently and measures similarity between their vectors — fast but limited by independent encoding. A cross-encoder sees the query and document concatenated as a single input, enabling deep interaction between them — highly accurate but O(n) complexity over the candidate set.
# Bi-encoder: compute query embedding ONCE, compare to all doc embeddings
# O(1) query encoding + O(n) dot products via ANN index = fast
query_vec = embed(query) # done once
results = vector_index.search(query_vec, top_k=100) # fast ANN search
# Cross-encoder: re-scores (query, doc) pairs jointly
# O(k) forward passes for k candidate documents = slow but accurate
for doc in results[:100]:
score = cross_encoder.score(query, doc.text) # joint scoringStage 1: Fast Coarse Retrieval
The first stage is a fast retriever — typically a bi-encoder with an approximate nearest neighbor index or a BM25 index — that retrieves a large candidate set (50-200 documents) with high recall but modest precision. The goal is not to be accurate; it is to not miss relevant documents. We cast a wide net and accept some false positives, knowing the second stage will clean them up.
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
# Stage 1: retrieve 100 candidates (high recall, modest precision)
vectorstore = FAISS.from_documents(documents, OpenAIEmbeddings())
coarse_retriever = vectorstore.as_retriever(
search_kwargs={'k': 100} # large candidate set
)
candidates = coarse_retriever.invoke(query)
print(f'Stage 1: retrieved {len(candidates)} candidate documents')Stage 2: Accurate Cross-Encoder Re-ranking
The second stage takes the candidate set from stage 1 and re-scores each (query, document) pair using a cross-encoder that reads both together. Because it processes only 50-200 candidates (not the full corpus), it can afford the expensive joint encoding. The cross-encoder's deep attention over the concatenated input makes it far more accurate at estimating true relevance than a bi-encoder.
from sentence_transformers import CrossEncoder
reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
def rerank(query: str, candidates: list[str], top_k: int = 5) -> list[str]:
# Score each (query, document) pair jointly
pairs = [[query, doc] for doc in candidates]
scores = reranker.predict(pairs)
# Sort by score descending
ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
return [doc for doc, _ in ranked[:top_k]]
candidate_texts = [doc.page_content for doc in candidates]
final_docs = rerank(query, candidate_texts, top_k=5)
print(f'Stage 2: selected top {len(final_docs)} documents after re-ranking')Why This Combination Works
The two-stage design exploits a key asymmetry: the first stage's fast ANN search scales to millions of documents in milliseconds, while the second stage's accurate cross-encoder operates only on the small candidate pool. You get the scalability of approximate search with the accuracy of exact joint scoring. The overall pipeline is both fast and highly accurate — something neither stage achieves alone.
Latency Profile of Two-Stage Retrieval
In a typical two-stage pipeline: stage 1 (vector ANN search over 1M documents) takes 5-20ms; stage 2 (cross-encoder over 100 candidates) takes 100-500ms depending on document length and hardware. The total latency budget is 150-600ms — acceptable for most applications. GPU acceleration in stage 2 can reduce re-ranking to under 30ms for short documents, making the pipeline competitive with single-stage retrieval on latency-sensitive applications.
import time
def two_stage_search(query, coarse_retriever, reranker, top_k=5):
t0 = time.perf_counter()
candidates = coarse_retriever.invoke(query) # stage 1
t1 = time.perf_counter()
candidate_texts = [c.page_content for c in candidates]
final_docs = rerank(query, candidate_texts, top_k) # stage 2
t2 = time.perf_counter()
print(f'Stage 1 (retrieval): {(t1-t0)*1000:.1f}ms')
print(f'Stage 2 (re-ranking): {(t2-t1)*1000:.1f}ms')
print(f'Total: {(t2-t0)*1000:.1f}ms')
return final_docsChoosing the Right Candidate Set Size
The first stage candidate set size is a critical hyperparameter. Too small (say, 10) and relevant documents may be missed before re-ranking even starts. Too large (say, 500) and stage 2 latency blows up. The recall at N curve — how many relevant documents are captured at different values of N — guides this choice. Typical sweet spots are between 50 and 150 candidates, where recall is near-saturation but latency remains manageable.
def recall_at_n(coarse_retriever, test_queries, golden_relevant, n_values):
for n in n_values:
recalls = []
for query, relevant in zip(test_queries, golden_relevant):
# Temporarily set k to n
coarse_retriever.search_kwargs['k'] = n
results = coarse_retriever.invoke(query)
retrieved_ids = {r.metadata.get('id') for r in results}
relevant_found = len(set(relevant) & retrieved_ids)
recalls.append(relevant_found / len(relevant))
avg = sum(recalls) / len(recalls)
print(f'N={n}: recall={avg:.3f}')Hybrid First Stage + Cross-Encoder Second Stage
The most powerful two-stage configuration pairs a hybrid retriever (dense + BM25) as the first stage with a cross-encoder as the second stage. Hybrid retrieval maximizes first-stage recall by combining semantic and keyword matching, and the cross-encoder then accurately selects the most relevant documents from the combined candidate pool. This configuration consistently achieves state-of-the-art retrieval quality on benchmarks.
from langchain.retrievers import EnsembleRetriever
# Stage 1: hybrid retrieval for maximum recall
hybrid_retriever = EnsembleRetriever(
retrievers=[bm25_retriever, vector_retriever],
weights=[0.4, 0.6],
)
# Stage 2: cross-encoder re-ranking for high precision
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
cross_encoder_model = HuggingFaceCrossEncoder(model_name='cross-encoder/ms-marco-MiniLM-L-6-v2')
compressor = CrossEncoderReranker(model=cross_encoder_model, top_n=5)
from langchain.retrievers import ContextualCompressionRetriever
two_stage = ContextualCompressionRetriever(
base_compressor=compressor,
base_retriever=hybrid_retriever,
)Commercial Re-ranking APIs
If you want cross-encoder accuracy without managing your own model, both Cohere Rerank and Jina AI Reranker offer cloud-hosted re-ranking APIs. You send a query and a list of document texts, and receive relevance scores back. These APIs use large cross-encoder models (often 500M+ parameters) that outperform self-hosted small cross-encoders, at the cost of additional API latency (50-300ms) and pricing per reranked document.
import cohere
co = cohere.Client('YOUR_API_KEY')
def cohere_rerank(query: str, documents: list[str], top_k: int = 5):
response = co.rerank(
model='rerank-english-v3.0',
query=query,
documents=documents,
top_n=top_k,
)
return [
{'text': documents[r.index], 'score': r.relevance_score}
for r in response.results
]
final = cohere_rerank(query, candidate_texts, top_k=5)
for doc in final:
print(f'Score {doc["score"]:.3f}: {doc["text"][:80]}')When Two-Stage Is Overkill
Two-stage retrieval adds complexity and latency compared to single-stage. It is not always necessary. For small corpora under 10,000 documents, a single cross-encoder over the full set may be fast enough. For applications where latency under 100ms is critical and accuracy gains are modest, single-stage dense retrieval may be preferable. Use two-stage when you have a large corpus, high accuracy requirements, and can afford 200-500ms retrieval latency.
Three-Stage Retrieval for Extreme Scale
For corpora of tens of millions of documents, a three-stage pipeline is sometimes used: first stage retrieves 10,000 candidates with ANN, second stage re-ranks to 100 using a fast small cross-encoder, and third stage re-ranks to 5 using a large powerful cross-encoder. Each stage applies a more expensive and accurate model to a smaller set. This architecture is used by large-scale search engines and document Q&A systems.
Quick Check
Test your understanding of why two-stage retrieval works from this lesson.
Lesson Recap
In this lesson you learned: bi-encoders are fast but limited to independent query-document encoding, cross-encoders are accurate through joint encoding but too slow for full-corpus search, and two-stage retrieval combines both: a fast first stage for high recall followed by an accurate second stage for high precision. The first stage retrieves many more candidates than needed to avoid missing relevant documents. Next up we implement cross-encoder re-ranking with Cohere and BGE.
Frequently asked questions
Is the “Why Two-Stage Retrieval Works” lesson free?
Yes — the full text of “Why Two-Stage Retrieval Works” is free to read here on the web, and the AI Engineering Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Engineering Academy course, upgrade to CoddyKit PRO.
What will I learn in “Why Two-Stage Retrieval Works”?
Understand the recall-precision trade-off in single-stage retrieval and how a fast coarse retriever followed by a slow but accurate re-ranker gets the best of both worlds. You practise AI Engineering Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.
Do I need any experience to start AI Engineering Academy?
No prior experience is required. AI Engineering Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 1 of 4, so you can start here or from the beginning and move at your own pace.
How long does the “Why Two-Stage Retrieval Works” lesson take?
Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.
Can I write and run code in this AI Engineering Academy lesson?
Yes. Every AI Engineering Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.