Reranking and Late Interaction: Cross-Encoders to ColBERT

Semantic and Hybrid Search ranked every section of this blog with BM25 and with embeddings, and fused the two. Both are cheap per document, because a search has to score every document. That cheapness has a price: the query and the document never meet. Each is reduced to its own vector, or its own bag of words, before they are compared. A reranker lets them meet, which costs far more per document, so it reads only the top few that a cheap first stage returns.

Table of contents

Bi-encoders and cross-encoders

A bi-encoder, the embedding model from Embeddings and Contrastive Learning, encodes the query and each document separately; the documents can be encoded once, ahead of time, and a query costs one encoding plus a dot product per document. A cross-encoder, first used for passage ranking by Nogueira and Cho (2019), reads the query and a document together, as one input, and outputs a relevance score. Every attention layer can then compare each query token with each document token, which is where its accuracy comes from, but nothing can be computed ahead of time: every query needs one full forward pass per document. Worked for this blog’s sections, the snapshot of 309 English sections that Semantic and Hybrid Search read (a rerun on the blog as committed finds 308, so the last digits can differ):

# Cost per query of the three designs, for this blog's 309 English
# sections, with a MiniLM-sized encoder: 22.7M parameters, about 11M of
# them the token-embedding table (a lookup), so roughly 2 x 11M FLOPs per
# token for the rest, ignoring attention. Queries are the section headings.
tok = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2", revision="1110a243fdf4706b3f48f1d95db1a4f5529b4d41")
secs = sections("en")
full = np.array([len(tok(s["body"]).input_ids) for s in secs])
read = np.minimum(full, 512)                                   # what a 512-token encoder reads
q = np.mean([len(tok(s["heading"]).input_ids) for s in secs])
N, flop_tok, K = len(secs), 2 * 11e6, 50
print(f"{N} sections, {full.mean():.0f} tokens each on average ({read.mean():.0f} read at most 512); "
      f"queries average {q:.1f} tokens")
bi = q * flop_tok + N * 384 * 2
cross = K * (q + read.mean()) * flop_tok
late = q * flop_tok + N * q * full.mean() * 128 * 2
print(f"bi-encoder: encode the query, then {N} dot products           {bi / 1e6:10.1f} MFLOPs per query")
print(f"cross-encoder, top {K}: {K} passes over query + section        {cross / 1e6:10.1f} MFLOPs per query")
print(f"late interaction (ColBERT, 128-d): every query x section token {late / 1e6:10.1f} MFLOPs per query")
print(f"  roughly K (q + L) / q = {K * (q + read.mean()) / q:,.0f}: the cross-encoder's multiple of one query encoding")

Full script: cost.py

The rest of the script scales the per-query cost to 10,000 and 1,000,000 documents and sizes each design’s index.

309 sections, 426 tokens each on average (361 read at most 512); queries average 9.0 tokens
bi-encoder: encode the query, then 309 dot products                198.3 MFLOPs per query
cross-encoder, top 50: 50 passes over query + section          407149.5 MFLOPs per query
late interaction (ColBERT, 128-d): every query x section token      501.3 MFLOPs per query
  roughly K (q + L) / q = 2,056: the cross-encoder's multiple of one query encoding
  at 10,000 documents: bi-encoder   0.21 GFLOPs, MaxSim over every document     10.0 GFLOPs
  at 1,000,000 documents: bi-encoder   0.97 GFLOPs, MaxSim over every document    981.6 GFLOPs

index size, every token indexed, 16-bit numbers:
  bi-encoder, one 384-d vector per section:                 0.23 MiB
  ColBERT, one 128-d vector per token:                     32.12 MiB
  ColBERTv2, 4-byte centroid id + 2-bit residuals:          4.52 MiB
  ColBERTv2, 4-byte centroid id + 1-bit residuals:          2.51 MiB

Reranking only 50 sections with a cross-encoder costs about 2,000 times the bi-encoder’s whole search; over all 309 it would be six times more again, and over a million documents, out of the question. The ratio is roughly $K (q + L) / q$ for $K$ candidates of length $L$ and a query of length $q$, since encoding the query dominates the bi-encoder’s cost: it grows with longer documents and shorter queries. Hence two-stage retrieval: a cheap first stage narrows the collection to $k$ candidates, and the expensive stage orders only those. The first stage’s recall at $k$ caps the whole pipeline: a document it misses can never be reranked into first place.

Late interaction: ColBERT

ColBERT (Khattab and Zaharia, 2020) sits between the two. Like a bi-encoder it encodes documents ahead of time, but it keeps one vector per token instead of one per document. At query time it encodes the query’s tokens and scores each document by MaxSim: every query token takes its best match among the document’s tokens, and the score is the sum of those best matches,

Worked on a three-word query, which becomes four word pieces, against two short documents. The token vectors here are from MiniLM, which wasn’t trained for this, so only the arithmetic is ColBERT’s:

import numpy as np
import torch
from transformers import AutoModel, AutoTokenizer

# Late interaction (ColBERT's MaxSim), worked on a 3-word query (4 word pieces). Every
# token of the query and of each document gets its own vector; each query
# token takes its best match among the document's tokens, and the score is
# the sum of those maxima. The token vectors here are all-MiniLM-L6-v2's,
# which was not trained for this: the arithmetic is ColBERT's, the numbers
# only illustrate it.
path, rev = "sentence-transformers/all-MiniLM-L6-v2", "1110a243fdf4706b3f48f1d95db1a4f5529b4d41"
tok, enc = AutoTokenizer.from_pretrained(path, revision=rev), AutoModel.from_pretrained(path, revision=rev).eval()
def vectors(text):
    ids = tok(text, return_tensors="pt")
    with torch.no_grad():
        h = enc(**ids).last_hidden_state[0, 1:-1]                     # drop [CLS] and [SEP]
    return tok.convert_ids_to_tokens(ids.input_ids[0, 1:-1]), torch.nn.functional.normalize(h, dim=-1).numpy()

q_tokens, Qv = vectors("paged kv cache")
for doc in ("vllm stores the kv cache in fixed-size pages, like virtual memory",
            "the cache line of a cpu holds sixty-four bytes"):
    d_tokens, Dv = vectors(doc)
    S = Qv @ Dv.T                                                     # every query token vs every document token
    best = S.argmax(1)
    print(f'document: "{doc}"')
    for i, t in enumerate(q_tokens):
        print(f"  {t:7s} best match {d_tokens[best[i]]!r:12s} {S[i, best[i]]:.2f}   (mean over the document {S[i].mean():.2f})")
    print(f"  MaxSim score {S.max(1).sum():.2f}; mean-of-similarities score {S.mean(1).sum():.2f}")

Full script: maxsim.py

document: "vllm stores the kv cache in fixed-size pages, like virtual memory"
  page    best match 'pages'      0.75   (mean over the document 0.29)
  ##d     best match 'like'       0.46   (mean over the document 0.30)
  kv      best match 'kv'         0.92   (mean over the document 0.30)
  cache   best match 'cache'      0.92   (mean over the document 0.35)
  MaxSim score 3.06; mean-of-similarities score 1.24
document: "the cache line of a cpu holds sixty-four bytes"
  page    best match 'cache'      0.35   (mean over the document 0.22)
  ##d     best match 'cache'      0.36   (mean over the document 0.24)
  kv      best match 'cpu'        0.24   (mean over the document 0.14)
  cache   best match 'cache'      0.87   (mean over the document 0.35)
  MaxSim score 1.81; mean-of-similarities score 0.95

The second document matches “cache” almost as well as the first (0.87 against 0.92), but nothing in it matches “kv” or “paged”, and MaxSim scores it far lower, 1.81 against 3.06. Why max and not mean: a query token should be judged by its best match, not diluted by every unrelated word in the document; averaging cuts both scores and the gap between them, to 1.24 against 0.95, and it would punish long documents for being long.

ColBERT gets most of a cross-encoder’s token-level matching; the original paper reports reranking more than 170 times faster than BERT, with 13,900 times fewer FLOPs per query. On this blog’s 309 sections its cost is close to a bi-encoder’s, because encoding the query dominates. On a large collection, scoring every document by MaxSim would not be: at a million documents it is about a thousand times a bi-encoder’s search, by the table’s last lines. So ColBERT doesn’t score everything. It first gathers candidates with a nearest-neighbor search per query token, through centroids in ColBERTv2’s engine, PLAID (Santhanam et al., 2022b), and runs MaxSim only on those. (ColBERT also pads every query to 32 tokens with a special mask token, which its ablations show helps.)

The other price is the index: one vector per token, 32 MiB for this blog’s sections against a quarter of a MiB for one vector per section, both in 16 bits. ColBERTv2 (Santhanam et al., 2022a) compresses each 128-number token vector to the 4-byte ID of its nearest centroid plus a 1- or 2-bit residual per dimension: 20 or 36 bytes instead of 256, about 13 or 7 times smaller here, with little loss. On why max and not mean, ColBERT’s own ablation is the evidence: swapping the maximum for an average lost noticeably on MS MARCO.

Training rerankers and retrievers: hard negatives

A reranker is trained on queries paired with relevant and irrelevant documents: scored one at a time against a label (Nogueira and Cho’s pointwise cross-entropy), or relevant against irrelevant (pairwise or listwise losses). What it learns depends on the irrelevant ones. Random documents are too easy: almost anything beats a random document. Hard negatives, documents a first stage ranked high but that aren’t relevant, teach the distinctions a reranker exists to make; Nogueira and Cho took theirs from BM25’s top thousand. The same technique trains retrievers.

Hard negatives carry a risk: some are relevant passages nobody labeled, and training against them teaches the wrong thing. ColBERTv2’s answer is denoised supervision. A first ColBERT model retrieves the top passages for each training query; a cross-encoder scores them; and ColBERT is trained to match those scores across 64 passages per query rather than trusting the binary labels, distillation from a teacher (see Knowledge Distillation). The teacher was ms-marco-MiniLM-L-6-v2, the very cross-encoder measured below.

LLM rerankers

A language model can rerank too: ask for a relevance score per document (pointwise; monoT5 reads it from the probability of the word “true”, Nogueira et al., 2020), or show it the query and the candidates and ask for an ordering (listwise; RankGPT, Sun et al., 2023). Listwise reranking has to fit the candidates in the context window, so it slides a window over them, and it is sensitive to the order they are shown in. It brings everything the model knows, and the cost is a language model’s: hundreds of tokens per candidate, through a model hundreds of times larger than a MiniLM. In practice it is a third stage, over the cross-encoder’s top few, or a way to label training data for a smaller reranker, which is what RankGPT’s authors did.

Measured: reranking the benchmark from Semantic and Hybrid Search

rerank.py reuses retrieval.py and sections.py, the two files of Semantic and Hybrid Search, up to that post’s encoder measurements: importing retrieval.py loads the sections, BM25, the dense encoders and fusion. It runs on the CPU; on an Apple M4 laptop the whole run takes about 35 minutes. The first stage is the best ranker from Semantic and Hybrid Search for each language: BM25 fused with dense passages by RRF for English, BM25 alone for Vietnamese. The second stage is a cross-encoder over that stage’s top 50, either the English ms-marco-MiniLM-L6-v2 or the multilingual mmarco-mMiniLMv2-L12-H384-v1. The heart of it is the loop over queries in two_stage:

K = 50
REVISION = {"cross-encoder/ms-marco-MiniLM-L6-v2": "233902d25c440f23af6f7d6e94d2946bac0bee0a",
            "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1": "1427fd652930e4ba29e8149678df786c240d8825"}
def two_stage(title, corpus, queries, stage1, ce_name, passages=False):
    ce = CrossEncoder(ce_name, device="cpu", revision=REVISION[ce_name])
    docs = [s["body"] for s in corpus]
    before, after, seconds = [], [], []
    for q, target in queries:
        s1 = stage1(q)
        top = np.argsort(-s1)[:K]
        before.append(int((s1 >= s1[target]).sum()))
        t0 = time.perf_counter()
        if passages:                                                     # score each 80-word passage, keep the best
            pairs, owner = [], []
            for j, i in enumerate(top):
                w = docs[i].split()
                for p in range(0, len(w), 80):
                    pairs.append((q, " ".join(w[p:p + 80]))); owner.append(j)
            sc = ce.predict(pairs, batch_size=64)
            s2 = np.full(K, -np.inf); np.maximum.at(s2, np.array(owner), sc)
        else:
            s2 = ce.predict([(q, docs[i]) for i in top], batch_size=K)
        seconds.append(time.perf_counter() - t0)
        if target in top:
            after.append(int((s2 >= s2[list(top).index(target)]).sum()))
        else:
            after.append(before[-1])                                     # stage 1 missed it; rerank can't help

Full script: rerank.py

The rest of two_stage prints recall@1 and MRR before and after reranking, the paired MRR difference with a 95% bootstrap interval, and the median reranking time. The script then runs it for each language and cross-encoder.

English (297 queries), cross-encoder ms-marco-MiniLM-L6-v2:
  stage 1 alone:    recall@1  40%  MRR 0.530  (right section in its top 50: 90%)
  + re-rank top 50: recall@1  45%  MRR 0.575  difference +0.045 [+0.012, +0.079]
  re-ranking time: 1468 ms per query on the CPU (50 pairs)
English (297 queries), cross-encoder ms-marco-MiniLM-L6-v2, best 80-word passage:
  stage 1 alone:    recall@1  40%  MRR 0.530  (right section in its top 50: 90%)
  + re-rank top 50: recall@1  40%  MRR 0.535  difference +0.004 [-0.032, +0.040]
  re-ranking time: 1737 ms per query on the CPU (50 pairs)
English (297 queries), cross-encoder mmarco-mMiniLMv2-L12-H384-v1:
  stage 1 alone:    recall@1  40%  MRR 0.530  (right section in its top 50: 90%)
  + re-rank top 50: recall@1  39%  MRR 0.516  difference -0.014 [-0.054, +0.025]
  re-ranking time: 2377 ms per query on the CPU (50 pairs)
Vietnamese (257 queries), cross-encoder mmarco-mMiniLMv2-L12-H384-v1:
  stage 1 alone:    recall@1  42%  MRR 0.539  (right section in its top 50: 91%)
  + re-rank top 50: recall@1  41%  MRR 0.529  difference -0.010 [-0.055, +0.037]
  re-ranking time: 2317 ms per query on the CPU (50 pairs)
Vietnamese (257 queries), cross-encoder ms-marco-MiniLM-L6-v2:
  stage 1 alone:    recall@1  42%  MRR 0.539  (right section in its top 50: 91%)
  + re-rank top 50: recall@1  35%  MRR 0.487  difference -0.053 [-0.098, -0.007]
  re-ranking time: 976 ms per query on the CPU (50 pairs)

Five runs, one per pairing of language and reranker, plus one that reads passages:

• A reranker trained on the right data helps. The English MS MARCO cross-encoder lifts English first-place hits from 40% to 45%, MRR $+0.045$ with a 95% interval of $[+0.012, +0.079]$ over queries. The first stage had the right section in its top 50 for 90% of queries, the ceiling on first-place hits for a reranker over 50; it gets part of the way.
• It needs to read the section. Scoring 80-word passages and keeping each section’s best, instead of reading the whole section cut at 512 tokens, erases the gain ($+0.004$, $[-0.032, +0.040]$): the relevant evidence for a heading is spread through its section.
• A bigger, multilingual reranker isn’t automatically better. The multilingual cross-encoder, twelve layers against six and twice the compute per pair, doesn’t improve English ($-0.014$, $[-0.054, +0.025]$), and on Vietnamese it only matches BM25 ($-0.010$, $[-0.055, +0.037]$). Both rerankers were trained on the same web questions, the multilingual one on mMARCO, MS MARCO machine-translated into 13 languages including Vietnamese, from a model distilled from one covering about a hundred languages. So the gap is in the model, not the domain; this experiment can’t say whether translation noise or capacity spread over many languages is to blame.
• A reranker in the wrong language hurts. The English cross-encoder on Vietnamese drops first-place hits from 42% to 35% ($-0.053$, $[-0.098, -0.007]$): it overrides a good first stage with scores it can’t compute well.
• Reranking is the expensive step. The run above timed it on the CPU of an Apple M4 laptop while other jobs ran on its GPU: about 1 to 1.5 seconds per query for the six-layer model over 50 sections, 2.3 seconds for the twelve-layer one. Run again on the same laptop with nothing else running (its output), the six-layer model took 0.8 to 1.1 seconds and the twelve-layer one 1.6. The first stage is a dot product and a word-count lookup per section; the cost table above puts the ratio between them in the thousands.

A reranker is a model, with a training distribution. Check it on the system’s own queries against the first stage alone before putting it in the pipeline, as for any other model.

Where this leads

• Building RAG assembles retrieval-augmented generation: chunking, query rewriting and caching, around the pipeline built in Vector Search, Semantic and Hybrid Search, and this post.

Takeaways

• A cross-encoder reads query and document together and scores more accurately than a bi-encoder, at one forward pass per document: about 2,000 times the bi-encoder’s whole search, for 50 documents here.
• Two-stage retrieval lets the cheap stage pick the candidates and the expensive one order them; the first stage’s recall at $k$ is the ceiling.
• ColBERT keeps one vector per token and scores by MaxSim, near a bi-encoder’s cost with token-level matching; its index is per token, and ColBERTv2’s compression cut it 7 to 13 times here.
• On this blog, the English MS MARCO reranker added $+0.045$ MRR in English, but only reading whole sections; the same model hurt Vietnamese ($-0.053$), and a multilingual one helped neither. Rerankers have training distributions too.

References

Start with Nogueira and Cho for the cross-encoder reranker, then Khattab and Zaharia for ColBERT and MaxSim; ColBERTv2 adds the compression and the distillation from a cross-encoder.

• Nogueira and Cho, Passage Re-ranking with BERT (2019): the first cross-encoder reranker for passages, and its pointwise training.
• Khattab and Zaharia, ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT (SIGIR, 2020): late interaction, MaxSim and one vector per token.
• Santhanam, Khattab, Saad-Falcon, Potts and Zaharia, ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction (NAACL, 2022a): ColBERTv2’s residual compression and distillation from a cross-encoder.
• Santhanam, Khattab, Potts and Zaharia, PLAID: An Efficient Engine for Late Interaction Retrieval (CIKM, 2022b): PLAID, ColBERTv2’s engine that prunes candidates by centroids.
• Nogueira, Jiang, Pradeep and Lin, Document Ranking with a Pretrained Sequence-to-Sequence Model (Findings of EMNLP, 2020): monoT5, a pointwise reranker that reads relevance from the word “true”.
• Sun et al., Is ChatGPT Good at Search? Investigating Large Language Models as Re-Ranking Agents (EMNLP, 2023): RankGPT, listwise reranking with an LLM.
• Bonifacio et al., mMARCO: A Multilingual Version of the MS MARCO Passage Ranking Dataset (2021): MS MARCO machine-translated into 13 languages, Vietnamese among them; the data behind the multilingual cross-encoder.