Two-Stage Retrieval: Optimizing RAG Pipelines with Cross-Encoder Rerankers

In search and Retrieval-Augmented Generation (RAG) systems, we must continuously balance speed and accuracy.

When searching over millions of documents, comparing a query embedding to every document using a deep transformer model is computationally prohibitive. To solve this, typical vector databases use Bi-Encoder architectures. However, while Bi-encoders are incredibly fast, they lose fine-grained attention interaction between queries and documents.

Two-stage retrieval systems resolve this by retrieving candidates using a fast Bi-Encoder (Stage 1) and then re-scoring them using a powerful Cross-Encoder (Stage 2).

Two-Stage Retrieval with Cross-Encoders

Bi-Encoders vs. Cross-Encoders

The structural differences between Bi-encoders and Cross-encoders are fundamental:

1. Bi-Encoder: The query and document are processed independently by a transformer model to generate two fixed-size vectors u and v. The similarity score is computed via a cheap vector operation:

s = \cos(u, v) = \frac{u \cdot v}{\|u\| \|v\|}

2. Cross-Encoder: The query and document are concatenated and fed into the transformer together:

\text{Input} = \text{[CLS]} \,\, \text{Query} \,\, \text{[SEP]} \,\, \text{Document} \,\, \text{[SEP]}

The model applies self-attention across all tokens simultaneously. This allows the query tokens to directly attend to specific document words, leading to a much more accurate relevance score, albeit at a significantly higher computational cost.

Setting up a Reranking Pipeline with SentenceTransformers

You can implement a two-stage retrieval pipeline using the popular sentence-transformers library:

from sentence_transformers import SentenceTransformer, CrossEncoder
import numpy as np

# 1. Load Bi-Encoder (Stage 1) and Cross-Encoder (Stage 2)
bi_encoder = SentenceTransformer("all-MiniLM-L6-v2")
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

# Documents and query
documents = [
    "Python is an interpreted, high-level, general-purpose programming language.",
    "Rust is a multi-paradigm, general-purpose programming language designed for performance and safety.",
    "FastAPI is a modern, fast (high-performance), web framework for building APIs with Python.",
    "Chroma is an open-source AI application database designed for developer productivity.",
]
query = "What is a fast web framework for Python?"

# Stage 1: Vector Search (Bi-Encoder)
query_emb = bi_encoder.encode(query)
doc_embs = bi_encoder.encode(documents)
scores = np.dot(doc_embs, query_emb) / (np.linalg.norm(doc_embs, axis=1) * np.linalg.norm(query_emb))

# Get top 3 candidates
top_k_indices = np.argsort(scores)[::-1][:3]
candidates = [documents[i] for i in top_k_indices]
print("Stage 1 Candidates:", candidates)

# Stage 2: Reranking (Cross-Encoder)
pairs = [[query, doc] for doc in candidates]
rerank_scores = cross_encoder.predict(pairs)

# Sort candidates by rerank scores
reranked_indices = np.argsort(rerank_scores)[::-1]
final_results = [candidates[i] for i in reranked_indices]
print("Stage 2 Reranked Results:", final_results)

Production Tradeoffs and Performance

By utilizing a two-stage pipeline:

• You retain sub-millisecond retrieval speeds across large databases by searching only top K candidates (e.g., K = 100).
• You filter out irrelevant context that passed vector similarity but lacked semantic alignment.
• The average performance improvement on standard benchmarks (like BEIR) is often 10% to 20% in NDCG@10 when incorporating a reranker.

Sizing the Candidate Set

The one number that decides whether this pipeline helps or hurts is K, the candidate count handed to the reranker. Cross-encoder cost is linear in K — a MiniLM-sized reranker runs roughly 5-15ms per pair on GPU and an order of magnitude slower on CPU, so K = 100 is a 100ms-plus tax on every query, sitting on the critical path where users feel it. Meanwhile the recall ceiling is set entirely by stage one: if the correct document isn't in the top K from the bi-encoder, no amount of reranking will conjure it. Reranking can only reorder what retrieval already found.

That gives you a concrete tuning procedure rather than a guess. Measure recall@K for your bi-encoder on a labelled eval set at K = 20, 50, 100, 200. The curve almost always flattens somewhere — commonly around 50 for well-chunked corpora — and that knee is your K. Going past it buys you nothing but latency; stopping short of it silently caps your ceiling no matter how good the reranker is.

Two caveats worth internalising before you ship this. Cross-encoder scores are not comparable across queries — they're logits from a binary relevance head, so a 4.2 on one query and a 4.2 on another mean different things, and any absolute cutoff you hardcode ("drop everything below 0") will behave inconsistently in production. Filter by rank, or calibrate per-query against the top score. And the off-the-shelf ms-marco models used here were trained on short web-search queries; on long conversational questions or domain-specific corpora such as legal or clinical text, an untuned reranker can score worse than the bi-encoder it's correcting. Benchmark it against your own data before assuming the reported 10-20% NDCG lift transfers.