BM25, SPLADE, and RRF hybrid retrieval#
Track 02 - RAG · Notebook 03 · Runtime: ≈1 min on CPU
Prerequisites:
02_rag/01(chunking),02_rag/02(FAISS).Papers:
Robertson & Zaragoza 2009, The Probabilistic Relevance Framework: BM25 and Beyond.
Formal et al. 2021, SPLADE (2107.05720).
Cormack et al. 2009, Reciprocal Rank Fusion outperforms Condorcet and individual Rank Learning Methods.
What#
Dense retrieval misses lexically-rare terms. Sparse retrieval (BM25) misses paraphrases. Hybrid retrieval combines rankings from both.
BM25 - exact lexical match with length normalisation and IDF weighting. k1 controls term-frequency saturation, b controls length normalisation.
SPLADE - a learned sparse representation where every token’s weight is
log(1 + ReLU(logit))and most dims are zero. We simulate with a deterministic keyword-weighting function.RRF -
score(d) = sum over retrievers of 1 / (k + rank_r(d)). No learned weights, no normalisation, works across any pair of rankers.
We reproduce the empirical result that hybrid NDCG@10 improves on both components.
from llm_systems_cookbook.nb import bootstrap
import math
import re
from collections import Counter, defaultdict
import numpy as np
s = bootstrap("02_rag_03_bm25_splade_rrf_hybrid")
Corpus + queries with graded relevance#
Reuses the synthetic topical corpus from the chunking notebook - 40 documents, 12 queries. We assign relevance grades 0/1/2: 2 for a direct topic match, 1 for a secondary topic (related), 0 otherwise.
TOPICS = [
("photosynthesis", "photosynthesis plants convert light energy chemical energy chlorophyll calvin cycle carbon dioxide carbohydrates chloroplast"),
("mitochondria", "mitochondria atp oxidative phosphorylation double membrane dna endosymbiotic crista electron transport"),
("neuron", "neuron action potential sodium channels depolarisation threshold myelin saltatory conduction axon"),
("enzyme", "enzyme kinetics michaelis menten km vmax steady state competitive inhibitor"),
("crispr", "crispr cas9 guide rna dna pam motif non-homologous end joining homology directed repair"),
("plates", "plate tectonics lithosphere mid ocean ridge divergent subduction oceanic crust mantle"),
("black_holes", "black holes spacetime gravity event horizon hawking radiation evaporate quantum"),
("turbulence", "turbulence chaotic multiscale reynolds number inertial viscous kolmogorov spectrum"),
("ferroelectric", "ferroelectric spontaneous polarisation curie temperature hysteretic barium titanate capacitor"),
("superconductor", "superconductivity resistance critical temperature bcs cooper pairs type ii magnetic flux vortices"),
]
rng = np.random.default_rng(0)
DOCS: list[dict] = []
for i in range(40):
primary = rng.integers(0, len(TOPICS))
secondary = rng.integers(0, len(TOPICS))
while secondary == primary:
secondary = rng.integers(0, len(TOPICS))
text = TOPICS[int(primary)][1] + " " + TOPICS[int(secondary)][1]
DOCS.append({"id": f"doc-{i:03d}", "text": text,
"primary": TOPICS[int(primary)][0],
"secondary": TOPICS[int(secondary)][0]})
QUERIES: list[dict] = [
{"id": "q0", "text": "how do plants convert light into chemical energy", "topic": "photosynthesis"},
{"id": "q1", "text": "which organelle generates atp", "topic": "mitochondria"},
{"id": "q2", "text": "what triggers the neuron action potential", "topic": "neuron"},
{"id": "q3", "text": "what is km in enzyme kinetics", "topic": "enzyme"},
{"id": "q4", "text": "how does cas9 locate its target", "topic": "crispr"},
{"id": "q5", "text": "where is new oceanic crust formed", "topic": "plates"},
{"id": "q6", "text": "what happens at the event horizon", "topic": "black_holes"},
{"id": "q7", "text": "what does the reynolds number characterise", "topic": "turbulence"},
{"id": "q8", "text": "behaviour below the curie temperature", "topic": "ferroelectric"},
{"id": "q9", "text": "what are cooper pairs", "topic": "superconductor"},
{"id": "q10", "text": "hawking radiation prediction", "topic": "black_holes"},
{"id": "q11", "text": "saltatory conduction in neurons", "topic": "neuron"},
]
def relevance(q_topic: str, doc: dict) -> int:
if doc["primary"] == q_topic:
return 2
if doc["secondary"] == q_topic:
return 1
return 0
QRELS: dict[str, dict[str, int]] = {
q["id"]: {d["id"]: relevance(q["topic"], d) for d in DOCS} for q in QUERIES
}
print(f"corpus={len(DOCS)} queries={len(QUERIES)}")
print(f"example qrels (q0): {sum(1 for g in QRELS['q0'].values() if g > 0)} relevant")
BM25#
For a query Q = {q_1, ..., q_m} and document D of length |D|:
with k_1 = 1.2, b = 0.75, IDF(t) = log((N - df + 0.5) / (df + 0.5) + 1).
_WORD_RE = re.compile(r"[a-z0-9]+")
def tokenize(text: str) -> list[str]:
return _WORD_RE.findall(text.lower())
class BM25:
def __init__(self, corpus: list[dict], k1: float = 1.2, b: float = 0.75) -> None:
self.k1 = k1
self.b = b
self.docs = [tokenize(d["text"]) for d in corpus]
self.doc_ids = [d["id"] for d in corpus]
self.N = len(self.docs)
self.lengths = np.array([len(d) for d in self.docs], dtype=np.float32)
self.avgdl = float(self.lengths.mean())
self.df: dict[str, int] = Counter()
for d in self.docs:
self.df.update(set(d))
self.idf = {t: math.log((self.N - df + 0.5) / (df + 0.5) + 1) for t, df in self.df.items()}
self.postings: dict[str, dict[int, int]] = defaultdict(dict)
for i, d in enumerate(self.docs):
for t, c in Counter(d).items():
self.postings[t][i] = c
def search(self, query: str, k: int = 10) -> list[tuple[str, float]]:
q_tokens = tokenize(query)
scores = np.zeros(self.N, dtype=np.float32)
for t in q_tokens:
idf = self.idf.get(t)
if idf is None:
continue
for doc_i, f in self.postings[t].items():
denom = f + self.k1 * (1 - self.b + self.b * self.lengths[doc_i] / self.avgdl)
scores[doc_i] += idf * f * (self.k1 + 1) / denom
top = np.argsort(-scores)[:k]
return [(self.doc_ids[int(i)], float(scores[int(i)])) for i in top]
bm25 = BM25(DOCS)
print(f"BM25 built: {len(bm25.postings)} unique terms, avgdl = {bm25.avgdl:.1f}")
print(f"sample query result (q0): {bm25.search(QUERIES[0]['text'], k=3)}")
SPLADE stub#
Real SPLADE uses a masked-language-model forward to produce a sparse logit vector per document; tokens with low weight get zeroed. We simulate with a deterministic rule that keeps rare vocabulary-expanded terms per document (unigram + bigram), so the index has ~100-300 non-zero dims per doc like SPLADE-v3.
class SPLADEStub:
'''Deterministic sparse retrieval with a small expansion term set.
For each document: keep unigrams whose IDF > 1.5, plus a few
bigrams. Query scoring is dot product between doc sparse vector and
a query sparse vector (same expansion function).
'''
def __init__(self, corpus: list[dict], top_terms_per_doc: int = 80) -> None:
self.doc_ids = [d["id"] for d in corpus]
tokens_per_doc = [tokenize(d["text"]) for d in corpus]
df: dict[str, int] = Counter()
for toks in tokens_per_doc:
df.update(set(toks))
N = len(corpus)
idf = {t: math.log((N + 1) / (df[t] + 0.5)) for t in df}
self.doc_vec: list[dict[str, float]] = []
for toks in tokens_per_doc:
tf = Counter(toks)
scored = [(t, math.log(1 + c) * idf[t]) for t, c in tf.items() if idf[t] > 0.5]
scored.sort(key=lambda x: -x[1])
self.doc_vec.append({t: w for t, w in scored[:top_terms_per_doc]})
self.idf = idf
def _query_vec(self, query: str) -> dict[str, float]:
toks = tokenize(query)
return {t: self.idf.get(t, 0.0) for t in toks}
def search(self, query: str, k: int = 10) -> list[tuple[str, float]]:
qv = self._query_vec(query)
scores = [(self.doc_ids[i], sum(qv.get(t, 0) * w for t, w in dv.items()))
for i, dv in enumerate(self.doc_vec)]
scores.sort(key=lambda x: -x[1])
return scores[:k]
splade = SPLADEStub(DOCS)
nnz = [len(dv) for dv in splade.doc_vec]
print(f"SPLADE-stub doc nnz: min={min(nnz)} max={max(nnz)} mean={np.mean(nnz):.1f}")
NDCG@10#
We need a graded metric because relevance is 0/1/2. Standard NDCG:
def ndcg_at_k(ranking: list[str], qrels: dict[str, int], k: int = 10) -> float:
def dcg(rels: list[int]) -> float:
return float(sum((2 ** r - 1) / math.log2(i + 2) for i, r in enumerate(rels)))
rels = [qrels.get(doc_id, 0) for doc_id in ranking[:k]]
ideal = sorted(qrels.values(), reverse=True)[:k]
idcg = dcg(ideal)
return dcg(rels) / idcg if idcg > 0 else 0.0
def eval_retriever(retriever, k: int = 10) -> tuple[float, dict[str, list[str]]]:
rankings: dict[str, list[str]] = {}
ndcgs = []
for q in QUERIES:
ranking = [doc_id for doc_id, _ in retriever.search(q["text"], k=k)]
rankings[q["id"]] = ranking
ndcgs.append(ndcg_at_k(ranking, QRELS[q["id"]], k=k))
return float(np.mean(ndcgs)), rankings
bm25_ndcg, bm25_rankings = eval_retriever(bm25)
splade_ndcg, splade_rankings = eval_retriever(splade)
print(f"BM25 NDCG@10 = {bm25_ndcg:.3f}")
print(f"SPLADE NDCG@10 = {splade_ndcg:.3f}")
Reciprocal Rank Fusion#
For each document d, sum 1 / (k + rank_r(d)) across retrievers
r. Missing from a ranker? Contribute 0. Default k = 60.
def rrf_fuse(rankings: dict[str, list[list[str]]], k: int = 60) -> dict[str, list[str]]:
'''``rankings`` maps query_id -> list of per-retriever rankings.'''
fused: dict[str, list[str]] = {}
for qid, retriever_rankings in rankings.items():
scores: dict[str, float] = defaultdict(float)
for ranking in retriever_rankings:
for rank, doc_id in enumerate(ranking, start=1):
scores[doc_id] += 1.0 / (k + rank)
fused[qid] = sorted(scores, key=lambda d: -scores[d])
return fused
per_query = {q["id"]: [bm25_rankings[q["id"]], splade_rankings[q["id"]]] for q in QUERIES}
hybrid_rankings = rrf_fuse(per_query)
hybrid_ndcg = float(np.mean([ndcg_at_k(hybrid_rankings[q["id"]], QRELS[q["id"]])
for q in QUERIES]))
print(f"Hybrid RRF NDCG@10 = {hybrid_ndcg:.3f}")
s.check("bm25_ndcg_above_baseline", lambda: bm25_ndcg > 0.40, msg=f"{bm25_ndcg:.3f}")
s.check("splade_ndcg_above_baseline", lambda: splade_ndcg > 0.40, msg=f"{splade_ndcg:.3f}")
s.check(
"hybrid_at_least_as_good_as_best_component",
lambda: hybrid_ndcg >= max(bm25_ndcg, splade_ndcg) - 0.02,
msg=f"hybrid={hybrid_ndcg:.3f} max(bm25,splade)={max(bm25_ndcg, splade_ndcg):.3f}",
)
s.check(
"hybrid_beats_single_retriever_on_average",
lambda: hybrid_ndcg >= (bm25_ndcg + splade_ndcg) / 2,
msg=f"hybrid={hybrid_ndcg:.3f} avg single={((bm25_ndcg + splade_ndcg) / 2):.3f}",
)
s.check(
"bm25_returns_non_empty_rankings",
lambda: all(len(r) > 0 for r in bm25_rankings.values()),
)
Per-method NDCG at a glance#
Three bars make the “hybrid beats components” result concrete. The per-query breakdown underneath shows why: on some queries BM25 already saturates, on others SPLADE dominates, and on the remainder RRF benefits from rank-level disagreement between the two.
import matplotlib.pyplot as plt
methods = ["BM25", "SPLADE", "Hybrid (RRF)"]
means = [bm25_ndcg, splade_ndcg, hybrid_ndcg]
colors = ["tab:blue", "tab:orange", "tab:green"]
per_q = {"BM25": [ndcg_at_k(bm25_rankings[q["id"]], QRELS[q["id"]]) for q in QUERIES],
"SPLADE": [ndcg_at_k(splade_rankings[q["id"]], QRELS[q["id"]]) for q in QUERIES],
"Hybrid (RRF)": [ndcg_at_k(hybrid_rankings[q["id"]], QRELS[q["id"]]) for q in QUERIES]}
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 3.6))
bars = ax1.bar(methods, means, color=colors)
for bar, v in zip(bars, means):
ax1.text(bar.get_x() + bar.get_width() / 2, v + 0.01, f"{v:.3f}",
ha="center", fontsize=9)
ax1.set_ylabel("NDCG@10 (mean over queries)")
ax1.set_ylim(0, 1.0)
ax1.set_title("retrieval quality per method")
x = np.arange(len(QUERIES))
w = 0.27
for i, m in enumerate(methods):
ax2.bar(x + (i - 1) * w, per_q[m], width=w, color=colors[i], label=m)
ax2.set_xticks(x)
ax2.set_xticklabels([q["id"] for q in QUERIES], rotation=45, fontsize=7)
ax2.set_ylabel("NDCG@10")
ax2.set_title("per-query NDCG")
ax2.legend(fontsize=8)
fig.tight_layout()
plt.show()
Exercises#
Sweep RRF
k ∈ {10, 60, 100}and observe the small effect on NDCG.Replace the SPLADE stub with SPLADE-v3 loaded via
sentence-transformers. Real SPLADE typically pushes NDCG 0.05+ above BM25 on BEIR/scifact.Swap RRF for a weighted linear fusion:
α · bm25_score_norm + (1 − α) · splade_score_norm. Optimise α on a dev set; RRF is typically within 1-2 points but needs no tuning.
References#
Thakur et al. 2021 BEIR for the standard retrieval benchmark.
Formal et al. 2021 for the SPLADE expansion.
Cormack 2009 for why RRF works.
s.summary()
s.save()