ColBERT, annotated

How to read this page

Any dotted word explains itself when you hover it, tab to it or tap it; so does every symbol in every equation. Each idea climbs the same ladder: an everyday picture, a tiny example, a diagram or demo, the math, and why it still matters. The heart of the paper is one line of arithmetic, MaxSim, and the MaxSim demo lets you watch it score real (hand-made) token vectors. The retrieval lesson implements it in NumPy.

Abstract · original

“ColBERT introduces a late interaction architecture that independently encodes the query and the document using BERT and then employs a cheap yet powerful interaction step that models their fine-grained similarity.”Khattab and Zaharia (2020), Abstract

Everyday picture

A bi-encoder judges a job applicant from a one-line summary: fast, but crude. A cross-encoder interviews every applicant about every question: precise, but you can only interview a few. ColBERT keeps a note card for every sentence of every CV, written in advance. When a question comes in, each part of the question looks for its best-matching card in each CV, and the matches are added up. The CVs are prepared once; only the quick card comparison happens per question.

What the paper claims

• Effectiveness competitive with BERT-based rankers, and better than every non-BERT model tested.
• Over 170 times faster and about 14,000 times fewer FLOPs per query than BERT re-ranking.
• Because document vectors are precomputed and the scoring is pruning-friendly, it can also search a whole collection directly with a vector index, not only re-rank a shortlist.

Why it matters today

“Multi-vector” retrieval, keeping several vectors per document and matching them token by token, grew out of this paper and its successors, and sits between single-vector search and cross-encoder reranking in modern retrieval stacks.

1 Introduction: four ways to match a query and a document · original

Everyday picture

In 2019, BERT-based rankers raised search quality sharply on the MS MARCO benchmark, but at 100 to 1,000 times the cost of earlier models; the paper cites evidence that adding even 100 ms to response times measurably hurts users. The question the paper asks: can we keep BERT's understanding while paying far less per query?

Reading it: each panel is one family of ranking models, read from the bottom (query and document text) to the top (a relevance score). In (a), each side is squeezed into one vector and the vectors are compared: documents can be encoded in advance, but all detail is lost in the squeeze. In (b), every query word is compared with every document word in a grid that a neural network reads. In (c), BERT reads query and document glued together, so every word attends to every other: the most accurate, and nothing can be precomputed. In (d), ColBERT's choice, each side keeps one vector per token, encoded separately (so documents are precomputed), and the only interaction is a cheap step at the very end: for each query vector, find its best match among the document's vectors, then add up those best matches.

The trade-off in one picture

Reading it: the x-axis is the time to answer a query, on a log scale (each gridline is ten times slower); the y-axis is MRR@10, the ranking quality. Up and to the left is good. The grey models are fast but weak; the orange BERT rankers are strong but sit at 10 to 33 seconds. ColBERT (red) is about as strong as BERT-base at about 61 ms of re-ranking, and its end-to-end version, which searches the whole collection instead of re-ranking BM25's list, is stronger still at under half a second.

2 Related work · original

• Neural matching models (KNRM, Duet, ConvKNRM): interaction grids read by small networks. Cheap, but well below BERT.
• BERT rankers (Nogueira and Cho): feed the query and document together through BERT and score the [CLS] output. The state of the art, at great cost.
• Moving work offline: doc2query and docTTTTTquery add predicted questions to each document before building a BM25 index; DeepCT uses BERT to re-weight BM25's term frequencies. Fast, but well below BERT's precision.
• Representation models such as SNRM compress each document into one (sparse) vector; the paper's ablation tests a single dense BERT vector too and finds it clearly worse than late interaction.

3 ColBERT · original

3.1 Architecture · original

Everyday picture

Two BERT readers (in fact one model, told which role it plays) turn the query and the document into bags of contextual token vectors: “bank” in the query “river bank erosion” gets a river-ish vector. Then each query vector goes shopping in the document's bag for its single best match.

Reading it: both sides climb the same stack: text with a marker token, BERT, a linear layer that shrinks each 768-number token vector to m = 128, and normalization to length 1. The document side (dashed box) runs once per document, offline, and also drops the vectors of punctuation. The query side runs once per query. They meet only in the MaxSim-and-sum box at the top, which has no learned parameters at all. That is why ColBERT is cheap: all the expensive BERT work is either precomputed or done once per query, never once per (query, document) pair.

3.2 Query and document encoders · original

Everyday picture

One translator, two hats. A “[Q]” sticker tells BERT it is reading a query; a “[D]” sticker, a document. Queries are short, so they are padded with blank [mask] slots, which BERT fills with guesses about what else the query might mean: a learned, soft form of query expansion the paper calls query augmentation.

Tiny example

The query “when was the eiffel tower built” becomes [CLS] [Q] when was the eiffel tower built [mask] [mask] … up to N_q = 32 tokens. BERT returns 32 vectors of 768 numbers; the linear layer shrinks each to 128; each is scaled to length 1. The result, E_q, is 32 unit vectors. A 70-token document yields about 70 unit vectors, minus any punctuation.

In words: “run BERT over the marked text, shrink every output vector with a linear layer, scale each to length 1; for documents, also throw away the punctuation vectors.”

With the numbers: query: 32 tokens → 32 × 768 from BERT → 32 × 128 → 32 unit vectors. A document of 70 tokens → at most 70 × 128. Because every vector has length 1, a dot product between two of them is exactly their cosine similarity, between −1 and 1. In two dimensions: (3, 4) normalizes to (0.6, 0.8) and (4, 3) to (0.8, 0.6), and their dot product 0.6 × 0.8 + 0.8 × 0.6 = 0.96 is their cosine similarity.

In Python:

import math
N_q, bert_width, m = 32, 768, 128
# BERT's output, then after the linear layer
(N_q, bert_width), (N_q, m)  # → ((32, 768), (32, 128))
# Normalize: divide by the length, so the length is 1
def normalize(v):
    length = math.sqrt(sum(x * x for x in v))
    return [x / length for x in v]
a, b = normalize([3, 4]), normalize([4, 3])
a, b  # → ([0.6, 0.8], [0.8, 0.6])
# dot product of unit vectors = cosine
round(sum(a_k * b_k for a_k, b_k in zip(a, b)), 2)  # → 0.96

3.3 Late interaction: MaxSim · original

“Intuitively, this interaction mechanism softly searches for each query term t_q ... against the document’s embeddings, quantifying the strength of the “match” via the largest similarity score between t_q and a document term t_d.”Khattab and Zaharia (2020), §3.1

Everyday picture

A checklist. For each thing the query asks about (“when?”, “Eiffel”, “tower”, “built”), scan the document for the one place that best answers it and note how good that match is. The document's score is the sum of the notes. A document that answers every part somewhere scores high, even if the parts are scattered.

In words: “for every query vector, find the highest dot product with any document vector, and add those maxima up.”

With the numbers: in the demo below, document A's five query rows have best matches 0.92 (“when” → “1889”), 0.97 (“eiffel” → “eiffel”), 0.97 (“tower” → “tower”), 0.98 (“built” → “completed”) and 0.77 ([mask] → “paris”), so S_q,d = 0.92 + 0.97 + 0.97 + 0.98 + 0.77 = 4.61. The “built” vector's best match is “completed” and the “when” vector's best match is “1889”, even though neither word appears in the query.

In Python:

# E_d: document A's vectors
doc = ["eiffel", "tower", "completed", "1889", "paris"]
# E_q_i · E_d_j, one row per query vector
sims = {
    "when":   [0.18, 0.08, 0.48, 0.92, 0.03],
    "eiffel": [0.97, 0.70, 0.05, 0.11, 0.58],
    "tower":  [0.70, 0.97, 0.10, 0.10, 0.21],
    "built":  [0.04, 0.09, 0.98, 0.45, 0.00],
    "[mask]": [0.20, 0.11, 0.57, 0.68, 0.77],
}
# which document vector each max_j picks
{q: doc[row.index(max(row))] for q, row in sims.items()}  # → {'when': '1889', 'eiffel': 'eiffel', 'tower': 'tower', 'built': 'completed', '[mask]': 'paris'}
# S_q,d = Σ_i max_j
round(sum(max(row) for row in sims.values()), 2)  # → 4.61

Reading it: rows are the query's token vectors (including one [mask] slot from query augmentation) and columns are the document's token vectors; each cell is a cosine similarity, darker for higher. The red-bordered cell in each row is that row's maximum: the best match for that query token anywhere in the document. The right-hand column shows the maxima and the total is the MaxSim score. Hover a row to see which document token it chose. Switch documents: doc B shares “eiffel” and “tower” but nothing answers “when” or “built”, so those rows find only weak matches. Doc C answers “when” and “built”, but for a different tower, so “eiffel” scores poorly. The bars below compare all three documents under three scoring rules: MaxSim, the average-similarity variant the paper's ablation tested, and a single mean-pooled vector per side, as a bi-encoder would use.

Training

ColBERT is trained on triples (query, relevant document, non-relevant document): each document is scored independently and a softmax cross-entropy over the two scores pushes the relevant one up, using the Adam optimizer. BERT is fine-tuned; the linear layer and the [Q] and [D] marker embeddings are trained from scratch. MaxSim itself has no parameters.

Why MaxSim, not something fancier?

Two reasons: it is very cheap, and it is pruning-friendly. Because each query vector's contribution is a maximum over document vectors, a vector index can find, for each query vector, the document vectors most similar to it across the whole collection, without scoring every document. Averaging instead of taking the maximum would destroy that property, and the ablation shows averaging is also less accurate.

3.4 Offline indexing · original

Everyday picture: write every document's note cards in advance, as efficiently as possible. The indexer uses several GPUs, groups documents of similar length together (in groups of B = 100,000, batches of b = 128) so little computation is wasted on padding, and runs the WordPiece tokenization in parallel on the CPU cores, which turned out to be a real bottleneck. Vectors are saved with 32 or 16 bits per number. The whole 8.8-million-passage MS MARCO collection takes about three hours on one server with four GPUs.

3.5 Top-k re-ranking · original

Everyday picture

BM25 hands over its top 1,000 documents. BERT re-ranking must run BERT 1,000 times on (query + document) sequences. ColBERT runs BERT once on the short query, loads the 1,000 documents' stored vectors onto the GPU, and does one batched matrix multiplication, a max and a sum.

Reading it: the x-axis is how many documents are re-ranked (k, log scale) and the y-axis is how many times more FLOPs BERT-base needs than ColBERT to re-rank them (log scale). The three points are the paper's own figures: about 180× at k = 10, 13,900× at k = 1,000 and 23,000× at k = 2,000, joined by straight lines for readability. The gap grows with k because ColBERT pays for BERT once per query, whatever k is, while BERT pays once per document. In ColBERT's case the actual bottleneck was not arithmetic but gathering and copying the stored vectors to the GPU; query encoding and scoring took only 13 of its 61 milliseconds.

3.6 End-to-end top-k retrieval · original

Everyday picture

Why rely on BM25's shortlist at all? Put every document token vector from the whole collection into one big nearest-neighbour index. For each of the query's 32 vectors, ask the index for the most similar document vectors anywhere; the documents those vectors came from are the candidates. Then score just those candidates exactly with MaxSim.

1. Filter: run N_q vector searches, each returning the top k′ document vectors. Map each back to its document: at most N_q × k′ candidate documents, K distinct.
2. Refine: score those K documents exactly with MaxSim, as in re-ranking.

With the numbers: N_q = 32 and k′ = 1,000 give at most 32,000 candidate hits, usually far fewer distinct documents, out of 8.8 million.

The index is FAISS's IVF with product quantization (IVFPQ): the vectors are clustered into P = 2,000 partitions in the experiments, each search probes the p = 10 nearest partitions, and each 128-number vector is compressed to s = 16 one-byte codes. See the HNSW companion and the vector index lesson for how these indexes work.

4 Evaluation · original

The benchmark and the metric

MS MARCO: 8.8 million web passages gathered from Bing results for about a million real queries, each query typically labelled with one relevant passage. The official metric is MRR@10: for each query, 1 divided by the rank of the first relevant passage in the top 10 (0 if none), averaged.

Tiny example: three queries whose first relevant passage is at rank 1, rank 3, and not in the top 10 score 1, 1/3 and 0, so MRR@10 = (1 + 0.333 + 0) / 3 = 0.444, reported as 44.4.

In words: “average, over all queries, one over the position of the first right answer, counting nothing if it is not in the top 10.”

With the numbers: (1/1 + 1/3 + 0) / 3 = 0.444.

In Python:

# rank_q of the first right answer; None: not in the top 10
ranks = [1, 3, None]
Q = len(ranks)
# (1/|Q|) Σ_q 1/rank_q
round(sum(1 / r if r is not None and r <= 10 else 0 for r in ranks) / Q, 3)  # → 0.444

Implementation: learning rate 3 × 10⁻⁶, batch size 32, N_q = 32 query vectors, m = 128 dimensions, 200,000 training iterations for MS MARCO.

4.2 Re-ranking and 4.3 end-to-end results

Reading it: the top block re-ranks BM25's top 1,000; the bottom rows search the whole collection. ColBERT re-ranking matches BERT-base's quality (34.9 against 34.7) at 61 ms instead of 10.7 seconds, with 7 billion FLOPs instead of 97 trillion. End-to-end, ColBERT is more accurate than when re-ranking (36.0), because it finds relevant passages BM25 never retrieved: its recall of relevant passages in its top 50 (82.9%) beats official BM25's recall in its top 1,000 (81.4%). On TREC CAR the story repeats: 31.3 MAP against BERT-base's 31.0 and BERT-large's 33.5.

4.4 Ablations

On a smaller 5-layer version, the paper tests: a single [CLS] vector per side (like a bi-encoder) instead of late interaction, which is considerably worse; average similarity instead of the maximum, which is worse; and no query augmentation, which is noticeably worse. Every piece earns its place. (The paper reports these in a figure, so we describe them without numbers.)

4.5 Index size: the price of many vectors · original

Everyday picture

Note cards for every word take much more shelf space than one summary card per document. How much, and how much can you shrink them?

Reading it: our back-of-envelope calculator for 8.8 million passages. The default of about 68 vectors per passage is inferred from the paper's reported 286 GiB for 128 dimensions at 4 bytes (286 GiB ÷ (8.8M × 128 × 4 bytes) ≈ 68). Set m = 24 and 2 bytes: about 27 GiB, the paper's smallest setting, which it reports costs only about 1 point of MRR@10 (33.9 against 34.9). The right-hand card is the same collection with one vector per passage, for comparison: multi-vector storage is roughly “vectors per passage” times larger. Reducing that cost is exactly what the follow-up work tackled.

5 Conclusions · original

Encoding queries and documents separately into fine-grained token vectors, and letting them interact only through a cheap, pruning-friendly MaxSim, keeps most of BERT's precision while cutting query cost by orders of magnitude, and makes end-to-end neural retrieval from a large collection practical.

What changed since 2020

Where it sits in a modern stack: single-vector bi-encoders (DPR, Sentence-BERT) are cheapest and most common; late interaction costs more storage for finer matching; cross-encoder rerankers are the most precise and are applied only to a shortlist. The retrieval lesson compares all three.

Glossary

Every term with hover guidance on this page, in one place.