Cross-Encoder Reranker

A bi-encoder is a retriever that maps query and document into the same vector space and ranks by cosine. The two encodings never see each other, so the model has to compress everything useful about a document into a single vector, blind to the query. That is what makes it fast: one embedding per document at index time, one per query at query time. It is also the only way to rank at corpus scale.

The cost is precision. Two documents on the same overall topic can have nearly identical embeddings even when one answers the query and the other does not, and the bi-encoder cannot tell them apart.

A cross-encoder is a model that reads the query and the document together and outputs one relevance scalar. It receives [query] [SEP] [document] as a single sequence and runs full attention across the join, so every token of the document can attend to every token of the query. The score comes from full context.

Its cost is throughput. The bi-encoder embeds once and queries forever, while the cross-encoder runs once per (query, document) pair. On a 10-million-document corpus that is 10 million forward passes per query, which is unrunnable in a request budget.

The fix is staging. Retrieve the top-N with the bi-encoder, then rerank those N down to a top-K with the cross-encoder. N is small (50 to 200), and the cross-encoder's quality lift is concentrated where it matters. Total latency stays in the request budget. Total quality is the cross-encoder's quality, capped by the bi-encoder's recall at N.

The cross-encoder's input shape

The standard packing is [CLS] query_tokens [SEP] document_tokens [SEP]. The CLS-position output goes into a single linear head that emits the relevance scalar. Some implementations mean-pool instead of reading the CLS position, and the difference is small. The toy model in this lesson mean-pools over non-pad positions. What matters is that the model produces one number per pair.

A 22M-parameter cross-encoder (the published ms-marco-MiniLM-L-6-v2 weight class) is the typical production point. Smaller models lose quality faster than they save latency. Larger ones, like bge-reranker-v2-m3 at 568M parameters, are reserved for offline reranking or for first-page reranking where K is small.

Why this lesson trains a tiny one

A real cross-encoder is a finetuned encoder transformer, and in production you load a checkpoint and run it. Here the goal is to show you the shape of the model and the shape of the latency-quality curve, not to train a state-of-the-art ranker. So we build a small nn.Module with one transformer block, multi-head attention (4 heads by default), and one regression head. It is initialized deterministically from a seed, so the demo is reproducible without weights on disk.

The toy model is fit to the fixture's hand-labeled pairs, so relevant query-document pairs get higher predicted scores than irrelevant ones. The end-to-end pipeline reranks the bi-encoder's output with those scores. The demo prints the scores and both rankings for you to read. It does not compute a correlation or recall figure, so exercise 1 is where you measure one.

Latency vs quality

The two-stage pipeline has one tunable: N. Sweep N from 5 to 100 on a held-out query set and you get the curve.

These numbers illustrate the shape. They are not measurements from this fixture. The shape itself is real: there is a knee around 20 to 50 candidates where the rerank lift saturates, and past it you are paying for nothing.

Pick N from the eval curve and the latency budget together. The cross-encoder cannot raise recall above the bi-encoder's recall at N, so a low N caps quality, not just latency.

code/main.py implements:

• CrossEncoder - a small torch.nn.Module: token, type, and position embeddings, one transformer block with multi-head attention and feedforward, and a mean-pooled head producing one scalar.
• tokenize_pair(query, document) - packs the two strings into a single id sequence with type ids that mark the boundary, deterministic and stdlib.
• train_tiny(model, triples, epochs=60, lr=5e-3) - full-batch supervised training (Adam, MSE loss) on 14 hand-labeled (query, document, label) triples, 60 epochs by default, returning the per-epoch loss. Labels are graded: 1.0 relevant, 0.0 irrelevant, partial values in between.
• rerank(model, query, candidates, top_k=5) - scores all candidates in one batch and returns the top-k (candidate, score) pairs, the production interface.
• BiEncoder with mock_embed - a deterministic hashed bag-of-words stand-in for the bi-encoder, so the demo needs no embedding model.
• pipeline(query, retriever, reranker, top_n, top_k) - the two-stage flow, returning both lists and the latency of each stage.
• A demo main() that builds a hardcoded 8-document CORPUS, indexes it with BiEncoder (this is a mock, not lesson 65's hybrid retriever), trains the model, then for three queries retrieves top-N, reranks to top-K, prints both lists side by side, and reports the latency of each stage.

Run it:

python3 code/main.py

The output shows the bi-encoder's top-N, the cross-encoder's top-K, and a timing summary. With N = 8 over an 8-document corpus the rerank pool is the whole corpus, so the timings show per-stage cost, not the savings of staging. The three demo queries also appear in the training triples, but each with only three of the eight documents, and the model is tiny, so the demo ranking is not guaranteed to match the labels. For "centralized authorization check function" the trained reranker puts d8 (cancelling jobs) above d4 (the check_permission doc), and the lower ranks for the other queries are filled with documents it never saw labeled for them. Read the scores as a check that the plumbing runs, not an eval. In a real pipeline the cross-encoder only scores the retrieved N, not the full corpus, and the two-stage total stays within the request budget.

Cross-encoder scoring is not symmetric. rerank(q, d) and rerank(d, q) give different scores. Feed the query first. If you swap them by accident, recall collapses.

N too low hides the lift. With N = K the cross-encoder only reorders the same K documents the bi-encoder already returned. It cannot pull in one the bi-encoder ranked lower, so recall@K shows zero lift. Pick N at least three times K.

Training data leaks into the eval. If the hand-labeled training pairs include the eval queries, the rerank looks magical. Keep train and eval strictly separate, even on a fixture. The demo is itself in this position: its three queries all appear in TRAIN_TRIPLES.

Production weights are dense. A 22M-parameter cross-encoder is 88MB at float32. Plan the model server's memory before you promise sub-100ms p95.

Batching matters. A real cross-encoder runs the N candidates in one batch. This lesson does that in rerank: _batch_encode builds the batched id and type-id tensors with torch.tensor(...), and one forward pass scores every candidate. Skip batching and the latency multiplies by N.

Production patterns:

• Pin the bi-encoder, cross-encoder, and N together. Change any one and the eval is no longer valid.
• Cache the reranker's output by (query, document_id) hash. The same query against a stable corpus reranks to the same order, and a cache hit is a free latency cut.
• Log the rank-1 cross-encoder score. A query whose top-1 score falls below a corpus-specific threshold is an out-of-domain hit. Surface it to the LLM as "I am not confident".

Lesson 68 evaluates this two-stage pipeline end to end. Lesson 69 wires this reranker behind the hybrid retriever from lesson 65 and in front of the answer generator, as the second stage of the end-to-end system.

1. Sweep N from 5 to 50 and plot recall@1 of the reranked output. Find the knee on this fixture.
2. Vary epochs in train_tiny (try 5, 20, 60, 200) and measure the score margin between positive and negative pairs after each run.
3. Replace mean-pooling with a CLS-token head. Compare convergence on this fixture.
4. Add a second cross-encoder head that predicts a binary "is this answer in the document" label. Use both heads at inference; one to rank, one to threshold.
5. Replace the mock BiEncoder with the retriever from lesson 65 and chain the two stages. Measure the change in top-K versus bi-encoder alone.

• Nogueira, Cho, "Passage Re-ranking with BERT", 2019 - the canonical cross-encoder ranker paper
• Reimers, Gurevych, "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks", 2019 - on bi-encoders vs cross-encoders
• SentenceTransformers Cross-Encoders documentation
• BGE Reranker v2 model card
• Phase 19 lesson 65 - the hybrid retriever feeding this rerank stage
• Phase 19 lesson 68 - the eval that measures the lift this rerank delivers