Late Interaction (ColBERT and MaxSim)

Retrieval models sit on a spectrum of how early the query and the document are allowed to look at each other. Late interaction is the middle position. The query and the document are encoded separately, as in a bi-encoder, but each is kept as a set of per-token vectors instead of being squeezed into one. The comparison happens afterward, token against token. Khattab and Zaharia introduced the design in ColBERT.

Three positions on the spectrum

• Bi-encoder (single vector). Query and document each become one vector, and the score is a dot product or cosine. Document vectors are precomputed and indexed, so search is fast, but one vector must summarize everything in the text.
• Cross-encoder. The query and document are concatenated and read by one model, which outputs a relevance score. Every query-document pair needs a model run, so it can only rerank a short candidate list. It sees full interaction between query and document tokens and is usually the most accurate.
• Late interaction. Document token vectors are precomputed offline. At query time only the query is encoded, and a cheap operation compares its token vectors with each candidate’s.

MaxSim

Let the query have token vectors \(q_1,\dots,q_m\) and the document have \(d_1,\dots,d_n\), all normalized so that dot products are cosines. The score is

\[ S(q,d)=\sum_{i=1}^{m}\max_{1\le j\le n}\; q_i\cdot d_j. \]

For each query token, find the document token most similar to it and take that similarity; then add these maxima over the query tokens. A query token only needs one good match somewhere in the document, and the document tokens that match nothing do not hurt the score. In words: every query token asks “is there something in this document that matches me?” and the answers are summed.

A small example with hypothetical 3-dimensional token vectors, whose axes loosely stand for finance, river, and approval. The query tokens are “bank” and “loan”; the document tokens are “bank”, “approved”, “river”, and “the”.

import math
def n(v): m=math.sqrt(sum(x*x for x in v)); return [x/m for x in v]
dot=lambda a,b:sum(x*y for x,y in zip(a,b))
Q={"bank":n([1,0,0]),"loan":n([0.5,0,1])}
D={"bank":n([1,0.1,0]),"approved":n([0,0,1]),"river":n([0,1,0]),"the":n([1,1,1])}
for q,qv in Q.items():
    sims={t:round(dot(qv,tv),3) for t,tv in D.items()}
    print(q,sims,"max =",max(sims.values()))
print("MaxSim =",round(sum(max(dot(qv,tv) for tv in D.values()) for qv in Q.values()),3))

Output:

bank {'bank': 0.995, 'approved': 0.0, 'river': 0.0, 'the': 0.577} max = 0.995
loan {'bank': 0.445, 'approved': 0.894, 'river': 0.0, 'the': 0.775} max = 0.894
MaxSim = 1.889

“bank” is matched by the document’s “bank” token, and “loan”, which the document never mentions, is matched by “approved” at 0.894. The unrelated token “river” contributes nothing. A single averaged document vector would blend “river” and “approved” into one point, which dilutes both.

Cost and quality

Storage is the price. A document of 200 tokens stores 200 vectors. For 1 million such documents with 128-dimensional float16 vectors that is 51.2 GB, against 1.5 GB for one 768-dimensional float16 vector per document. In ColBERT the vectors are projected to a small dimension for this reason, and the search runs in two steps: approximate nearest-neighbor search over token vectors collects candidate documents, and MaxSim is then computed exactly on those candidates. The authors report query execution about two orders of magnitude faster, and four orders of magnitude fewer floating-point operations per query, than BERT cross-encoder ranking, with competitive effectiveness on passage ranking benchmarks.

In quality terms, late interaction is usually better than a single-vector bi-encoder at matching specific terms and sub-phrases, and usually below a cross-encoder, because query and document tokens never attend to each other inside the encoder. In practice it is used as a first-stage retriever or as a reranker over a candidate set, and its cost is engineering: bigger indexes, token-level search, and compression schemes.

References

• Khattab, O., Zaharia, M., ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT — https://arxiv.org/abs/2004.12832