Reranking: Cross-Encoders for Precise Information Retrieval › Code Implementation
Language AI HandbookAll chapters Reranking: Cross-Encoders for Precise Information Retrieval › Code ImplementationOne section · Sections in this chapter8 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Let us implement a complete retrieve-then-rerank pipeline using the sentence-transformers library, which provides convenient interfaces for both bi-encoders and cross-encoders. This implementation will make the concepts concrete and give you a working starting point for your own retrieval systems.
# Install required packages (run once locally, then comment out):
# import subprocess
# import sys
# subprocess.check_call([sys.executable, "-m", "pip", "install", "-q",
# "sentence-transformers", "scikit-learn", "numpy"])# Sample corpus of technical documentation
documents = [
"To reset your password, click 'Forgot Password' on the login page. Enter your registered email address, and we will send you a reset link.",
"Account recovery options include password reset via email, security questions, or contacting support. If you no longer have access to your registered email, use the alternative verification methods available in your account settings.",
"Email notifications can be configured in your account settings. You can choose to receive updates about password changes, login attempts, and security alerts to any verified email address.",
"Two-factor authentication adds an extra layer of security to your account. You can enable 2FA using an authenticator app or SMS verification.",
"If you suspect unauthorized access to your account, immediately change your password and review recent login activity in your security dashboard.",
"Premium users have access to priority support and advanced security features including hardware key authentication.",
"Mobile app users can enable biometric authentication for faster, secure login using fingerprint or facial recognition.",
"Account deletion is permanent and cannot be undone. All data associated with your account will be removed within 30 days.",
]
query = "How do I reset my password if I forgot the email associated with my account?"We begin by loading a bi-encoder for initial retrieval. We will use a small model suitable for demonstration, though production systems might use larger embeddings. The bi-encoder encodes both the query and documents into dense vectors that capture semantic meaning, letting us to quickly identify the subset of documents that are most likely to be relevant based on vector similarity.
## Load bi-encoder for initial retrieval
print("Loading bi-encoder...")
import numpy as np
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
bi_encoder = SentenceTransformer("all-MiniLM-L6-v2")
## Encode corpus (this would be precomputed and indexed in production)
print("Encoding documents...")
doc_embeddings = bi_encoder.encode(documents, convert_to_tensor=True)
query_embedding = bi_encoder.encode([query], convert_to_tensor=True)
## Compute similarities
similarities = cosine_similarity(
query_embedding.cpu().numpy(), doc_embeddings.cpu().numpy()
)[0]
## Get top-k from bi-encoder
k = 4 # Retrieve top 4 for reranking
top_k_indices = np.argsort(similarities)[::-1][:k]
top_k_docs = [documents[i] for i in top_k_indices]
top_k_scores = similarities[top_k_indices]Top-4 documents from bi-encoder: 1. [Score: 0.750] To reset your password, click 'Forgot Password' on the login page. Enter your re... 2. [Score: 0.698] Account recovery options include password reset via email, security questions, o... 3. [Score: 0.473] If you suspect unauthorized access to your account, immediately change your pass... 4. [Score: 0.447] Email notifications can be configured in your account settings. You can choose t...
The bi-encoder returns candidates based on semantic similarity. Notice that documents mentioning "password" and "email" rank highly, but the specific condition about forgetting the email is not well captured by vector similarity alone. The bi-encoder identifies documents in the correct semantic neighborhood but lacks the discriminative power to verify whether the specific constraint mentioned in the query is satisfied by each document.
Now we load a cross-encoder to rerank these candidates. Cross-encoders are typically larger models that process query-document pairs jointly. Unlike the bi-encoder, which produced independent vectors, the cross-encoder will process the concatenated text of the query with each candidate document, letting fine-grained interaction between the query's specific constraints and the document's content.
## Load cross-encoder for reranking
print("Loading cross-encoder...")
import time
from sentence_transformers import CrossEncoder
cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
## Prepare pairs for cross-encoder
pairs = [[query, doc] for doc in top_k_docs]
## Time the reranking
start_time = time.time()
cross_scores = cross_encoder.predict(pairs)
rerank_time = time.time() - start_time
## Sort by cross-encoder scores (argsort gives indices of sorted elements, so reranked_indices[i] is the original position)
reranked_order = np.argsort(cross_scores)[::-1]
reranked_docs = [top_k_docs[i] for i in reranked_order]
reranked_scores = cross_scores[reranked_order]
## For display: reranked_indices[i] should give the original rank (1-indexed) of the i-th reranked doc
reranked_indices = reranked_order + 1 # Convert to 1-indexed original ranksReranking completed in 61.4ms Reranked results: 1. [Score: 7.156, originally #1] To reset your password, click 'Forgot Password' on the login page. Enter your re... 2. [Score: 4.691, originally #2] Account recovery options include password reset via email, security questions, o... 3. [Score: -2.313, originally #4] Email notifications can be configured in your account settings. You can choose t... 4. [Score: -2.403, originally #3] If you suspect unauthorized access to your account, immediately change your pass...
The cross-encoder reranks the documents, potentially changing the order based on fine-grained relevance signals that the bi-encoder missed. The document about alternative verification methods for lost email access should now rank higher than the generic password reset instructions. This reordering demonstrates how the cross-encoder identifies the specific relevance signals that distinguish a satisfactory answer from a superficially similar but ultimately unhelpful document.
Let us compare the latency characteristics of both approaches to understand the computational tradeoff. This benchmark illustrates why we reserve the cross-encoder for only a small set of candidates rather than applying it to the entire corpus.
# Benchmark latency comparison
print("Latency Benchmark")
print("=" * 50)
# Bi-encoder timing (query encoding only, assuming pre-indexed docs)
n_trials = 10
bi_times = []
for _ in range(n_trials):
start = time.time()
q_emb = bi_encoder.encode([query])
# Simulate dot product with precomputed index (very fast)
_ = cosine_similarity(q_emb, doc_embeddings.cpu().numpy()[:k])
bi_times.append(time.time() - start)
# Cross-encoder timing
cross_times = []
for _ in range(n_trials):
start = time.time()
_ = cross_encoder.predict(pairs)
cross_times.append(time.time() - start)Bi-encoder (query + top-4 retrieval): 10.34ms +/- 1.02ms Cross-encoder (score 4 pairs): 13.15ms +/- 1.70ms Slowdown factor: 1.3x
The cross-encoder is significantly slower, often by one to two orders of magnitude. This demonstrates why we use it only on a small candidate set rather than the entire corpus. In production systems, this latency difference necessitates careful architectural decisions about how many documents to rerank and whether to employ optimization techniques such as batching, quantization, or cascade filtering.
Finally, let us examine the attention patterns to understand how the cross-encoder makes its decisions. We will use the cross-encoder's ability to output attention weights. This visualization helps illustrate the token-level interactions that enable the cross-encoder to distinguish between documents that share vocabulary but differ in their relevance to the specific query.
# Use the worked example's relevant document so the simulated alignments remain
# stable even if an upstream model version ranks the small demo corpus differently.
attention_document = documents[1]
# Tokenize to show alignment
query_tokens = query.split()
doc_tokens = attention_document.split()
use_book_style()
plt.rcParams["figure.figsize"] = (6.0, 4.0)
# Create a simplified visualization of token-level interaction
# In practice, you would extract actual attention weights from the model
fig = plt.figure()
ax = plt.gca()
# Simulate attention pattern showing key alignments
# Real attention would be extracted from model layers
attention_matrix = np.random.rand(len(query_tokens), len(doc_tokens)) * 0.1
# Highlight key alignments
query_keywords = ["forgot", "email", "account"]
doc_keywords = ["access", "registered", "email", "alternative", "verification"]
for i, qt in enumerate(query_tokens):
for j, dt in enumerate(doc_tokens):
qt_lower = qt.lower().strip("?,.\"'")
dt_lower = dt.lower().strip(".,;\"'")
# Strong attention for keyword matches
if qt_lower in ["email", "password", "account"] and dt_lower in [
"email",
"password",
"account",
]:
attention_matrix[i, j] = 0.8
elif qt_lower == "forgot" and dt_lower in ["access", "longer"]:
attention_matrix[i, j] = 0.7
elif qt_lower == "reset" and dt_lower in ["recovery", "reset"]:
attention_matrix[i, j] = 0.6
image_artist = ax.imshow(
attention_matrix, cmap="Blues", aspect="auto", interpolation="nearest"
)
# Set ticks
ax.set_xticks(range(len(doc_tokens)), doc_tokens, rotation=45, ha="right")
ax.set_yticks(range(len(query_tokens)), query_tokens)
ax.set_xlabel("Document Tokens")
ax.set_ylabel("Query Tokens")
ax.set_title("Cross-Encoder Attention: Query-Document Alignment")
ax.grid(False)
fig.colorbar(image_artist, ax=ax, label="Attention Weight")
plt.show()
Query: How do I reset my password if I forgot the email associated with my account? Document visualized: Account recovery options include password reset via email, security questions, or contacting support...
The attention visualization reveals how the cross-encoder aligns specific query tokens with relevant document tokens. Unlike the bi-encoder, which produces a single similarity score, the cross-encoder can match "forgot" with "no longer have access" even though these phrases share no words. This shows the power of contextualized cross-attention. This ability to recognize semantic equivalence across different phrasings is what enables the cross-encoder to correctly identify Doc B as the most relevant, despite its lower lexical overlap with the query compared to Doc A.