Cross-encoder reranking

Deep dive

How it works

Parts:

• first stage — cheap retriever over the corpus (BM25, ANN such as HNSW, or their fusion); returns k candidate ids
• pair builder — makes (query, passage) inputs; splits long documents into overlapping passages (chunking)
• cross scorer — model taking both texts jointly (BERT-style encoder, T5, or LLM), one score per pair
• aggregator — document score = max (or top-2 mean) over its passages
• sorter and cache — keeps the best n (top-k); cache keyed on (query, doc id, model version)

1. Run the first stage for the top k.
2. Split each candidate into windows that fit the model's maximum length.
3. Batch the (query, window) pairs, sorted by length to limit padding; score them.
4. Take each document's max over its windows; sort; return the top n.
5. If another component needs a probability, calibrate on judged data; raw scores rank only within one query.

 corpus (N docs, encoded once)
   |  BM25 / ANN / fusion              ~ms, indexed
   v
 top-k ids --text--> [query ; passage] --> cross scorer --> score
   (top k)              one call per pair                     |
                                                              v
                                           max over passages, sort -> top-n

Why joint wins: a bi-encoder squeezes each text into one vector before seeing the other, so word order, negation and "from X to Y" versus "from Y to X" are largely lost. Why it cannot be indexed: the score depends on the pair, so there is nothing to precompute per document.

Model families: an encoder classifier (BERT reranking); a sequence-to-sequence model whose "true/false" token logits are the score (monoT5); an instructed LLM scoring one passage (pointwise) or ordering a window of passages (listwise, RankGPT). Distillation: train a bi-encoder (two-tower) on a cross-encoder's scores (knowledge distillation); RankGPT distilled an LLM's rankings into a small model that, per its abstract, beat a larger supervised one on BEIR. Late interaction (ColBERT) encodes separately but compares token vectors, so documents can be precomputed.

Code

The scorer is a toy standing in for a neural cross-encoder: bigram and proximity features exist only when query and text are read together, which a bag-of-words first stage cannot see. Cascade, windows, max-aggregation and metrics are real.

use std::time::Instant;

fn toks(s: &str) -> Vec<String> {
    s.split(|c: char| !c.is_alphanumeric()).filter(|w| !w.is_empty()).map(|w| w.to_lowercase()).collect()
}

/// Stage 1: BM25-lite on bags of words (k1 = 1.2, b = 0.75); a real system uses an inverted index.
fn bm25(q: &[String], docs: &[Vec<String>], i: usize) -> f64 {
    let n = docs.len() as f64;
    let avg = docs.iter().map(|d| d.len()).sum::<usize>() as f64 / n;
    q.iter().map(|t| {
        let tf = docs[i].iter().filter(|w| *w == t).count() as f64;
        let df = docs.iter().filter(|d| d.contains(t)).count() as f64;
        let idf = ((n - df + 0.5) / (df + 0.5) + 1.0).ln();
        idf * tf * 2.2 / (tf + 1.2 * (0.25 + 0.75 * docs[i].len() as f64 / avg))
    }).sum()
}

/// Toy "cross" scorer: sees query AND passage together, so it can use word order and distance.
/// Stand-in for a neural forward pass; same interface: (query, passage) -> score.
fn cross(q: &[String], p: &[String]) -> f64 {
    let bigrams = q.windows(2).filter(|b| p.windows(2).any(|c| c == *b)).count() as f64;
    let pos: Vec<usize> = q.iter().filter_map(|t| p.iter().position(|w| w == t)).collect();
    let span = if pos.len() > 1 { (pos.iter().max().unwrap() - pos.iter().min().unwrap() + 1) as f64 } else { 1.0 };
    2.0 * bigrams / (q.len().max(2) - 1) as f64 + pos.len() as f64 / q.len() as f64 / span.sqrt()
}

/// Windows of `w` tokens (stride w/2), one score per pair; document score = max over windows.
fn rerank(q: &[String], docs: &[Vec<String>], cand: &[usize], w: usize, pairs: &mut usize) -> Vec<usize> {
    let mut s: Vec<(usize, f64)> = cand.iter().map(|&i| {
        let d = &docs[i];
        let starts = (0..d.len().max(1)).step_by(w / 2).take_while(|&a| a == 0 || a + w / 2 < d.len());
        let best = starts.map(|a| { *pairs += 1; cross(q, &d[a..(a + w).min(d.len())]) }).fold(f64::MIN, f64::max);
        (i, best)
    }).collect();
    s.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
    s.into_iter().map(|x| x.0).collect()
}

/// nDCG@k with gain 2^g - 1; `rel[i]` is the grade of doc i.
fn ndcg(rank: &[usize], rel: &[u32], k: usize) -> f64 {
    let dcg = |g: Vec<u32>| g.iter().take(k).enumerate().map(|(i, &g)| (2f64.powi(g as i32) - 1.0) / (i as f64 + 2.0).log2()).sum::<f64>();
    let mut ideal = rel.to_vec();
    ideal.sort_by(|a, b| b.cmp(a));
    dcg(rank.iter().map(|&i| rel[i]).collect()) / dcg(ideal)
}

fn main() {
    let corpus = [
        "flights from paris to boston: fares and schedules",           // 0 wrong direction
        "cheap flights from boston to paris, daily nonstop service",   // 1 relevant
        "boston paris flights",                                        // 2 keyword soup
        "paris travel guide with museum tickets and flights to rome",  // 3
        "apple stock price today and analyst targets",                 // 4 relevant
        "apple pie: the price of stock for broth, and pie crust tips", // 5
        "stock market news: apple, price moves on the nasdaq",         // 6 partial
        "how to prune an apple tree; fruit stock price at the market", // 7
    ];
    // (query, graded relevance per doc)
    let qrels = [("flights from boston to paris", [0, 3, 1, 0, 0, 0, 0, 0]), ("apple stock price", [0, 0, 0, 0, 3, 0, 1, 0])];
    let docs: Vec<Vec<String>> = corpus.iter().map(|d| toks(d)).collect();
    let k = 4;
    let (mut n, mut m, mut pairs) = ([0.0; 2], [0.0; 2], 0);
    for (q, rel) in &qrels {
        let qt = toks(q);
        let mut first: Vec<usize> = (0..docs.len()).collect();
        first.sort_by(|&a, &b| bm25(&qt, &docs, b).total_cmp(&bm25(&qt, &docs, a)).then(a.cmp(&b)));
        let mut second = rerank(&qt, &docs, &first[..k], 8, &mut pairs); // only the top k are re-scored
        second.extend_from_slice(&first[k..]);
        for (j, r) in [&first, &second].iter().enumerate() {
            n[j] += ndcg(r, rel, k);
            m[j] += 1.0 / (1 + r.iter().position(|&i| rel[i] > 0).unwrap()) as f64;
        }
        println!("{q:30} bm25 {:?} rerank {:?}", &first[..k], &second[..k]);
        if rel[1] == 3 { assert!(first[0] == 0 && second[0] == 1); } // known case: reversed route fixed
    }
    let nq = qrels.len() as f64;
    println!("nDCG@{k} {:.3} -> {:.3}   MRR {:.3} -> {:.3}", n[0] / nq, n[1] / nq, m[0] / nq, m[1] / nq);
    assert!(n[1] > n[0] && m[1] > m[0]);

    let qt = toks("flights from boston to paris");
    let (mut p, t) = (0usize, Instant::now());
    for _ in 0..200_000 { std::hint::black_box(rerank(&qt, &docs, &[0, 1, 2, 3], 8, &mut p)); }
    println!("toy cross scorer: {:.2e} pairs/s", p as f64 / t.elapsed().as_secs_f64());
}

Compiled and run (release): BM25 ranks the reversed route "paris to boston" first (same words, shorter text); the reranker puts the right document first. Over the two toy queries nDCG@4 went 0.822 → 0.983, MRR 0.750 → 1.000; on query 2 it still lifts an irrelevant document over a partial match (toy features).

Complexity and constants

• Per query: first stage plus k model calls; each is O(L^2) attention in the joint length L (query + passage) plus O(L) feed-forward.
• Measured here (toy scorer, one core of an Apple M4 Pro, includes allocation and sorting): 7.6-11.5 x 10^6 pairs/s over three runs, order of 100 ns per pair. A neural model is orders of magnitude slower (order of ms per pair on CPU); never extrapolate this figure.

Efficiency

• Cost: k forward passes over L joint tokens per query. Small encoder: order of ms per pair on CPU, order of 0.1-1 ms batched on GPU (not measured here). Bound: compute for long inputs, memory bandwidth for small batches. Toy scorer: about 0.1 µs per pair (measured).
• Levers:
  • smaller k — latency falls linearly in k (cost model) — when judged recall@k is flat beyond it
  • shorter inputs — attention work falls roughly with L^2 — truncate, or score titles and best passages
  • length-sorted batches — fewer padded tokens, better GPU use (practice) — always when k pairs are scored together
  • caching — skips repeated pairs — repeated queries and stable documents, key includes model version
  • distillation or late interaction — moves cost to an index lookup (ColBERT precomputes document representations) — when k calls do not fit the budget
• Naive → standard → state of the art:
  • cross-encoder over the whole corpus · O(N) model calls, infeasible
  • BM25 or ANN top-k, cross-encoder on k · O(k) calls
  • fused first stage, small k, timeout fallback, distilled or late-interaction first cut · fewest calls at equal quality

Implementation notes

• Choosing k: on 50+ judged queries plot first-stage recall@k and post-rerank nDCG@10; take the knee. Vespa bounds its second ranking phase by a rerank count (phased ranking).
• Truncation: query and passage share the model's maximum length (512 tokens for BERT-style encoders); reserve room for the query, cut at sentence boundaries, or window and take the max.
• Calibration: logits shift with query length and model; for "show nothing if low" fit Platt or isotonic scaling per model on judged pairs.
• Latency budget: first stage + k · t_pair / parallelism; one batch for all pairs; on timeout fall back to first-stage order.
• Format: match training exactly (separators, prompt template); a wrong template silently loses quality.

Pitfalls

• Reranking a list that lacks the answer: nDCG cannot exceed first-stage recall@k. Measure it first.
• Comparing scores across queries or models, then thresholding.
• Dropping the first-stage signal on exact identifiers; fuse by rank (fusion) or learn weights (learning to rank).
• Domain shift: BEIR found BM25 a robust baseline, re-rankers best on average but costly (BEIR); evaluate on your own queries (metrics).
• Listwise LLM rerankers: order sensitivity, missing or invented ids in the output, cost growing with window size.
• Near-duplicates all score high; diversify afterwards (MMR).

Variants and alternatives

In production

• Search engines run phased ranking (Vespa); RAG pipelines rerank retrieved chunks before filling the prompt.
• Open models: MS MARCO-trained cross-encoders (sentence-transformers) and BGE rerankers; MS MARCO and BEIR are the usual benchmarks.
• Same cascade in e-commerce, patent and legal prior-art, code and biomedical search, and ticket deduplication.

Implementations

Sources

• Nogueira, Cho, Passage Re-ranking with BERT, 2019. paper
• Khattab, Zaharia, ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT, 2020. paper · code
• Nogueira, Jiang, Lin, Document Ranking with a Pretrained Sequence-to-Sequence Model, 2020. paper
• Sun et al., Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents, 2023. paper · code
• Bajaj et al., MS MARCO: A Human Generated MAchine Reading COmprehension Dataset, 2016. paper
• Thakur et al., BEIR: A Heterogenous Benchmark for Zero-shot Evaluation of Information Retrieval Models, 2021. paper · code