Advanced RAG: Hybrid Search + Cross-Encoder Reranking

Lab 4 showed the basic RAG pipeline: one embedding model, cosine similarity, done. In production, naive dense retrieval leaves 10–30% of queries on the table — especially ones involving rare keywords, acronyms, product names, or proper nouns. The industry's answer is a two-stage retrieval stack:

query ┬─→ [dense]   ─┐
       │              ├─→ merge top-K ─→ [cross-encoder reranker] ─→ top-N → LLM
       └─→ [BM25]    ─┘      (via RRF)

Three upgrades over vanilla RAG, each earned by empirical pain:

1. BM25 keyword retrieval — catches rare-term queries that dense embeddings sometimes miss
2. Reciprocal Rank Fusion (RRF) — the canonical way to blend heterogeneous score distributions
3. Cross-encoder reranker — a second, slower, more accurate model re-orders the top candidates

Why each piece exists

• BM25 (Robertson & Zaragoza, 2009) is a 40-year-old TF-IDF descendant. In 2024, it's still competitive with dense retrieval — and strictly better on rare-token queries where embeddings' semantic smoothing hurts. Sparse retrieval's strength: it scores based on exact term overlap, weighted by how distinctive each term is.
• RRF (Cormack et al., 2009) blends ranked lists by summing 1/(k + rank) across retrievers. Simple, parameter-light, beats most learned fusion methods. Default in Weaviate, Elasticsearch, Vespa.
• Cross-encoder reranker is a model that takes both query AND passage together and outputs a single relevance score. Much slower than dense retrieval (no precomputed vectors), but far more accurate per pair. Used to re-score the top 20–100 candidates from first-stage retrieval. BAAI's bge-reranker-base is the go-to open option; Cohere Rerank is the go-to paid option.

References to read alongside

• BGE paper (Xiao et al., 2023) — includes the reranker family.
• Reciprocal Rank Fusion (Cormack et al., 2009) — the paper. 4 pages. Read it.
• ColBERT (Khattab & Zaharia, 2020) — late-interaction, a cross-encoder/bi-encoder hybrid.
• MS MARCO passage ranking leaderboard — where all these retrieval models are benchmarked against each other.
• Vespa docs on hybrid search — the production playbook at one of the better open search platforms.
• OpenAI's cookbook on hybrid search — modern reference implementation.

---

The best way to see why hybrid search matters is to find a query where dense retrieval fails. Our corpus will include a few entries with distinctive proper nouns / acronyms — the kind of content dense embeddings smooth away.

import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

device = torch.device('cuda')

# Corpus includes rare proper nouns that test embeddings' generalization
corpus = [
    "Our internal CI/CD system is called NIMBUS-OPS. It auto-deploys every merged PR to staging within 90 seconds.",
    "The customer support portal is built on ZephyrDesk v3. It handles 50,000 tickets per day at 99.95% SLA.",
    "We use a custom observability stack combining Prometheus, Grafana, and Loki for logs aggregation.",
    "All employee badges are managed by the HALCYON access-control system, replaced in 2023 from the previous Mercury platform.",
    "Our data warehouse runs on Snowflake. Daily ETL batches process ~2 TB into fact tables for the BI team.",
    "Production Kubernetes clusters are managed through a GitOps layer built on ArgoCD + Kustomize.",
    "Engineering onboarding takes two weeks and includes a rotation through each of the six platform teams.",
    "Incident response follows the PagerDuty + Slack + Zoom bridge pattern. Ops rotations are weekly, 6 engineers per tier.",
    "Our internal ML platform, codenamed PRISM, runs on Ray and supports distributed training up to 128 GPUs.",
    "Employee expenses go through Concur Expense. Approval chains are based on manager org hierarchy with a 72-hour SLA.",
]
print(f'{len(corpus)} documents in corpus')

# Use BGE-small as the dense embedder (same as Lab 4)
embed_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5')
embed_model = AutoModel.from_pretrained('BAAI/bge-small-en-v1.5').to(device)
embed_model.eval()

def embed(texts, max_len=256):
    enc = embed_tokenizer(texts, padding=True, truncation=True, max_length=max_len, return_tensors='pt').to(device)
    with torch.no_grad():
        out = embed_model(**enc)
    return F.normalize(out.last_hidden_state[:, 0], dim=-1)

passage_embs = embed(corpus)
print(f'Embedded into {passage_embs.shape[1]}-dim vectors')

# Find a dense-retrieval failure case.
# Query contains a rare proper noun (HALCYON) that should clearly match doc 3.
# Dense embeddings tend to smooth away exact term identity in favor of general semantics — so they might
# retrieve a different access-control-ish doc first instead.

failure_query = 'Who provides our employee badge access system?'
q_vec = embed([failure_query])
scores = (q_vec @ passage_embs.T).squeeze(0)
dense_top1 = int(scores.argmax().item())

# Show top-3 for context
top_scores, top_idx = torch.topk(scores, k=3)
print(f'DENSE RETRIEVAL for query: {failure_query!r}\n')
for rank, (s, i) in enumerate(zip(top_scores.cpu().tolist(), top_idx.cpu().tolist()), 1):
    marker = ' <-- should be this one (HALCYON)' if i == 3 else ''
    print(f'  #{rank} [idx={i}, score={s:.3f}] {corpus[i][:100]}...{marker}')

TRUE_RELEVANT_IDX = 3  # The HALCYON doc
dense_failure = {
    'query': failure_query,
    'dense_top1_idx': dense_top1,
    'true_relevant_idx': TRUE_RELEVANT_IDX,
    'dense_top1_text': corpus[dense_top1],
}
if dense_top1 != TRUE_RELEVANT_IDX:
    print(f'\n DENSE MISSED: returned idx {dense_top1} instead of the true {TRUE_RELEVANT_IDX} (HALCYON doc).')
    print(' Rare proper nouns like HALCYON are exactly where keyword retrieval shines.')