Learning to Rank & Retrieval Models
1. The Ranking Problem
Given a query and a set of documents , the ranking problem aims to find a permutation that orders documents by relevance:
where is the relevance function and is the discount factor emphasizing top positions.
2. BM25 (Best Matching 25)
2.1 BM25 Score Formula
BM25 is a probabilistic retrieval function based on the BM25 scoring formula:
where:
- : term frequency of term in document
- : document length
- : average document length in the corpus
- : term frequency saturation parameter (typically 1.2-2.0)
- : length normalization parameter (typically 0.75)
2.2 IDF (Inverse Document Frequency)
where is the total number of documents and is the number of documents containing term .
2.3 BM25 Properties
- Term saturation: As increases, the score approaches a limit:
-
Length normalization: Longer documents are penalized proportionally to their length relative to average.
-
IDF weighting: Rare terms contribute more to relevance.
2.4 BM25 Variants
BM25+ (Lv & Zhai, 2011): Adds lower bound to prevent zero scores for non-matching documents:
SPL2 (Lv & Zhai, 2009): Uses Poisson approximation:
3. Learning to Rank Framework
3.1 Ranking Pipeline
4. Pointwise Ranking
4.1 Regression
Treat relevance as a continuous value:
4.2 Classification
Treat relevance as discrete classes:
4.3 Limitations
Pointwise methods don't capture the relative ordering between documents. They treat each query-document pair independently.
5. Pairwise Ranking
5.1 RankNet
RankNet (Burges et al., 2005) learns to rank pairs:
Given a pair where :
The probability that is ranked higher than :
where is the model score.
5.2 LambdaRank
LambdaRank extends RankNet by weighting pairs by their impact on NDCG:
The gradient becomes:
5.3 Listwise Pairwise Loss
The pairwise loss over all document pairs in a query:
6. Listwise Ranking
6.1 ListNet
ListNet (Cao et al., 2007) optimizes the listwise ranking loss using a softmax-based probability distribution:
Given scores , the probability of permutation :
The loss minimizes the KL divergence between the predicted and ground-truth distributions:
6.2 ListMLE
ListMLE (Xia et al., 2008) uses maximum likelihood estimation:
6.3 Softmax Cross-Entropy Loss
The listwise softmax cross-entropy loss:
where is the relevance label (can be multi-grade).
7. Learned Sparse Retrieval
7.1 SPLADE (Sparse Lexical and Expansion Model)
SPLADE learns sparse representations by expanding queries with predicted terms:
where is the contextualized representation of token .
The resulting representation is sparse (many zeros) and can be indexed using inverted indices.
7.2 Comparison
| Aspect | Dense | Sparse (SPLADE) | BM25 |
|---|---|---|---|
| Representation | Dense vector | Sparse vector | Term counts |
| Index | ANN (HNSW) | Inverted index | Inverted index |
| Storage | High (768d) | Medium (30K vocab) | Low |
| Semantics | Strong | Moderate | Weak |
| Exact match | Weak | Strong | Strong |
8. ColBERT (Contextualized Late Interaction)
8.1 Architecture
ColBERT encodes queries and documents independently but uses token-level interaction for scoring:
where and are L2-normalized token embeddings.
8.2 MaxSim Operation
For each query token, find the maximum similarity with any document token:
The total score is the sum over all query tokens:
8.3 ColBERT-2
Improvements in ColBERT-2:
- Residual compression: Quantize document embeddings
- Token pruning: Reduce document length dynamically
- Weighted query expansion: Learn importance weights
8.4 ColBERT Training
where is the positive document and are hard negatives.
9. Cross-Encoder Fine-Tuning
9.1 Training Objective
Cross-encoders are trained with binary cross-entropy:
where for relevant documents and for non-relevant.
9.2 Hard Negative Mining
The effectiveness of cross-encoders depends heavily on hard negatives:
9.3 Knowledge Distillation
Distill cross-encoder knowledge into bi-encoder:
where and are scores from the cross-encoder and bi-encoder respectively.
10. Implementation
from sentence_transformers import SentenceTransformer, CrossEncoder
import faiss
import numpy as np
class RetrievalSystem:
def __init__(self):
# Bi-encoder for retrieval
self.bi_encoder = SentenceTransformer("BAAI/bge-large-en-v1.5")
# Cross-encoder for re-ranking
self.cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
# FAISS index
self.index = faiss.IndexFlatIP(1024) # Inner product
self.documents = []
def index_documents(self, documents):
self.documents = documents
embeddings = self.bi_encoder.encode(
[d["text"] for d in documents],
normalize_embeddings=True,
show_progress_bar=True
)
self.index.add(embeddings.astype(np.float32))
def retrieve(self, query, top_k=100):
# Stage 1: Bi-encoder retrieval
q_emb = self.bi_encoder.encode(
[query], normalize_embeddings=True
).astype(np.float32)
scores, indices = self.index.search(q_emb, top_k)
candidates = [self.documents[i] for i in indices[0]]
return candidates, scores[0]
def rerank(self, query, documents, top_k=10):
# Stage 2: Cross-encoder re-ranking
pairs = [(query, d["text"]) for d in documents]
ce_scores = self.cross_encoder.predict(pairs)
# Combine with bi-encoder scores
ranked = sorted(
zip(documents, ce_scores),
key=lambda x: x[1],
reverse=True
)[:top_k]
return ranked
def search(self, query, top_k=10):
candidates, _ = self.retrieve(query, top_k=100)
results = self.rerank(query, candidates, top_k)
return results
11. Evaluation Metrics
11.1 Precision@K
11.2 Recall@K
11.3 MAP (Mean Average Precision)
11.4 NDCG@K
11.5 MRR (Mean Reciprocal Rank)
where is the rank of the first relevant document.
References
- Robertson & Zaragoza (2009). "The Probabilistic Relevance Framework: BM25 and Beyond." Foundations and Trends in IR.
- Burges et al. (2005). "Learning to Rank using Gradient Descent." ICML.
- Cao et al. (2007). "Learning to Rank: From Pairwise Approach to Listwise Approach." ICML.
- Khattab & Zaharia (2020). "ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT." SIGIR.
- Formal et al. (2021). "SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking." SIGIR.
- Nogueira & Cho (2019). "Passage Re-ranking with BERT." arXiv.