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.

Conclusion

Integrating a Cross-Encoder reranker is one of the lowest-effort, highest-impact optimizations for RAG architectures. By leveraging the speed of Bi-encoders for initial filtering and the attention capability of Cross-encoders for final sorting, developers can construct scalable, highly accurate search systems.