Two-stage reranking

Open in Colab ▶️ Run this notebook in Colab

Two-stage reranking#

Track 02 - RAG · Notebook 05 · Runtime: ≈30 s on CPU

Prerequisites: 02_rag/02 (FAISS), 02_rag/03 (BM25).

Papers:

  • Nogueira & Cho 2019, Passage Re-ranking with BERT (1901.04085).

  • Xiao et al. 2024, BGE M3-Embedding (2402.03216).

What#

Production retrieval is almost always two-stage:

  1. First stage - a fast bi-encoder (or BM25) ranks the whole corpus and keeps top-N candidates. Cheap but coarse.

  2. Second stage - a slower cross-encoder scores each (query, doc) pair jointly and reranks the top-N. Expensive but accurate.

The ratio is asymmetric: cross-encoders are 30-100× slower per doc but only run on N docs (say 100), not the whole corpus (say 10⁶). Total latency is dominated by the first stage’s recall, not the second stage’s precision.

We simulate both stages: a noisy bi-encoder for first-stage, and an oracle-with-noise “cross-encoder” that has access to topic ground-truth plus a small error. Target: cross-encoder reranking raises NDCG@10 by ≥ 0.05 over the bi-encoder alone.

from llm_systems_cookbook.nb import bootstrap

import math

import numpy as np

s = bootstrap("02_rag_05_two_stage_reranking")

Corpus: 200 docs, 30 queries, graded qrels#

Each doc has a topic; a random fraction of docs in each topic are marked “true positives” for queries on that topic. First-stage scores are noisy projections of the true relevance; cross-encoder scores are much less noisy.

rng = np.random.default_rng(0)

N_DOCS = 200
N_QUERIES = 30
N_TOPICS = 10

# Each doc has a primary + secondary topic.
doc_primary = rng.integers(0, N_TOPICS, size=N_DOCS)
doc_secondary = rng.integers(0, N_TOPICS, size=N_DOCS)
query_topic = rng.integers(0, N_TOPICS, size=N_QUERIES)


def relevance(q_topic: int, i: int) -> int:
    if doc_primary[i] == q_topic:
        return 2
    if doc_secondary[i] == q_topic:
        return 1
    return 0


QRELS = np.stack([np.array([relevance(int(query_topic[q]), i) for i in range(N_DOCS)])
                  for q in range(N_QUERIES)])  # (Q, N)
print(f"corpus={N_DOCS}  queries={N_QUERIES}  topics={N_TOPICS}")
print(f"mean relevant docs per query: {np.sum(QRELS > 0, axis=1).mean():.1f}")

First stage - noisy bi-encoder#

We simulate the bi-encoder as noisy(relevance) + noise, where noise std is large enough to shuffle the top-20 quite aggressively. That’s the realistic pattern: bi-encoders miss doc-level nuance, so the top-10 often contains some irrelevant distractors.

def bi_encoder_scores(q: int, noise_std: float = 1.0) -> np.ndarray:
    return QRELS[q].astype(np.float32) + rng.normal(0, noise_std, size=N_DOCS).astype(np.float32)


def rank(scores: np.ndarray) -> np.ndarray:
    return np.argsort(-scores)


def ndcg_at_k(ranking: np.ndarray, qrels: np.ndarray, k: int = 10) -> float:
    rels = qrels[ranking[:k]]
    dcg = sum((2 ** r - 1) / math.log2(i + 2) for i, r in enumerate(rels))
    ideal = sorted(qrels, reverse=True)[:k]
    idcg = sum((2 ** r - 1) / math.log2(i + 2) for i, r in enumerate(ideal))
    return dcg / idcg if idcg > 0 else 0.0


bi_ndcgs = []
for q in range(N_QUERIES):
    sc = bi_encoder_scores(q, noise_std=1.2)
    r = rank(sc)
    bi_ndcgs.append(ndcg_at_k(r, QRELS[q]))
bi_mean = float(np.mean(bi_ndcgs))
print(f"bi-encoder NDCG@10 = {bi_mean:.3f}")

Second stage - cross-encoder#

Rerank only the top-N candidates. The cross-encoder has access to richer features (modelled here as lower-noise scores), so its ranking among those candidates is dramatically better than the bi-encoder’s ranking over the same set.

Two scoring configurations: top_n = 100 (generous) and top_n = 20 (aggressive). Aggressive reranking is cheaper but risks missing relevant docs that the bi-encoder dropped below rank 20.

def cross_encoder_score(q: int, doc_ids: np.ndarray, noise_std: float = 0.15) -> np.ndarray:
    return QRELS[q, doc_ids].astype(np.float32) + \
           rng.normal(0, noise_std, size=len(doc_ids)).astype(np.float32)


def two_stage_ndcg(top_n: int) -> float:
    ndcgs = []
    for q in range(N_QUERIES):
        first = rank(bi_encoder_scores(q, noise_std=1.2))
        cand = first[:top_n]
        rerank_scores = cross_encoder_score(q, cand)
        reorder = cand[np.argsort(-rerank_scores)]
        tail = first[top_n:]
        full = np.concatenate([reorder, tail])
        ndcgs.append(ndcg_at_k(full, QRELS[q]))
    return float(np.mean(ndcgs))


tn_100 = two_stage_ndcg(top_n=100)
tn_20 = two_stage_ndcg(top_n=20)
tn_200 = two_stage_ndcg(top_n=N_DOCS)
print(f"two-stage (top_n=100):  NDCG@10 = {tn_100:.3f}")
print(f"two-stage (top_n=20):   NDCG@10 = {tn_20:.3f}")
print(f"two-stage (top_n=200=all): NDCG@10 = {tn_200:.3f}  (ceiling)")

s.check(
    "reranking_raises_ndcg_at_top_n_100",
    lambda: tn_100 >= bi_mean + 0.05,
    msg=f"bi={bi_mean:.3f}  two-stage(100)={tn_100:.3f}",
)
s.check(
    "reranking_at_top_n_20_close_to_ceiling",
    lambda: tn_20 >= tn_200 - 0.10,
    msg=f"top_20={tn_20:.3f}  top_all={tn_200:.3f}",
)
s.check(
    "ceiling_top_200_reaches_high_ndcg",
    lambda: tn_200 >= 0.85,
    msg=f"ceiling = {tn_200:.3f}",
)
s.check(
    "first_stage_is_nontrivial",
    lambda: bi_mean >= 0.50,
    msg=f"bi-encoder NDCG = {bi_mean:.3f}",
)

Quality vs modelled latency#

Pair NDCG with a latency model (bi-encoder 0.1 ms/doc, cross-encoder 5 ms/doc). Swept over top_n, the bar chart shows recall rising quickly and saturating, while latency climbs linearly. The knee around top_n = 20-50 is where most production systems sit.

import matplotlib.pyplot as plt

BI_MS, CE_MS = 0.1, 5.0
sweeps = [0, 5, 10, 20, 50, 100, N_DOCS]
ndcgs = [bi_mean] + [two_stage_ndcg(top_n=n) for n in sweeps[1:]]
lats  = [BI_MS * N_DOCS + CE_MS * n for n in sweeps]
labels = ["bi only"] + [f"+rerank {n}" for n in sweeps[1:]]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 3.6))
x = np.arange(len(sweeps)); w = 0.38
ax1.bar(x - w/2, ndcgs, width=w, color="tab:blue", label="NDCG@10")
ax1b = ax1.twinx()
ax1b.bar(x + w/2, lats, width=w, color="tab:orange", alpha=0.8, label="latency")
ax1.set_xticks(x); ax1.set_xticklabels(labels, rotation=30, fontsize=8)
ax1.set_ylabel("NDCG@10", color="tab:blue")
ax1b.set_ylabel("latency (ms)", color="tab:orange")
ax1.set_ylim(0, 1.0); ax1.set_title("quality vs latency vs top_n")

ax2.plot(lats, ndcgs, "o-", color="tab:purple")
for i, lab in enumerate(labels):
    ax2.annotate(lab, (lats[i], ndcgs[i]), fontsize=7,
                 xytext=(4, 4), textcoords="offset points")
ax2.set_xlabel("latency (ms / query)"); ax2.set_ylabel("NDCG@10")
ax2.set_title("Pareto curve"); ax2.grid(True, alpha=0.3)
fig.tight_layout(); plt.show()

Exercises#

  1. Sweep top_n ∈ {5, 10, 20, 50, 100, 200} and plot NDCG vs latency (assume cross-encoder per-pair latency is 5 ms, bi-encoder is 0.1 ms per doc). Find the knee.

  2. Two-stage with three stages: bi-encoder → ColBERT → cross-encoder. ColBERT catches bi-encoder misses; cross-encoder does the final polish. Very common in production search stacks.

  3. Real reranker: pip install sentence-transformers, use cross-encoder/ms-marco-MiniLM-L-6-v2. Measure p50 latency on a 100-doc batch on your hardware.

References#

  • BGE reranker model cards.

  • MS-MARCO leaderboard reranker section.

s.summary()
s.save()