ColBERTv2 - late interaction

Open in Colab ▶️ Run this notebook in Colab

ColBERTv2 - late interaction#

Track 02 - RAG · Notebook 04 · Runtime: ≈1 min on CPU

Prerequisites: 02_rag/02 (FAISS), 02_rag/03 (BM25 hybrid).

Papers:

  • Khattab & Zaharia 2020, ColBERT (2004.12832).

  • Santhanam et al. 2022, ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction (2112.01488).

What#

Bi-encoder retrieval compresses a document into one vector before the query is known - fine for many cases, but details get lost. Cross encoders score every query-document pair jointly - most accurate, slowest. ColBERT is the middle path:

  • Encode a document into one vector per token at index time.

  • At query time, encode the query into one vector per token.

  • Score with MaxSim: for each query token, find its best matching document token; sum those max similarities.

Formally:

MaxSim(Q, D) = sum_{q in Q} max_{d in D} q · d

This recovers much of the fine-grained interaction a cross-encoder would compute, at O(|Q|·|D|) per pair - way cheaper than re-running the model end-to-end.

We implement MaxSim from scratch on a synthetic token-embedding corpus and verify three properties: MaxSim is always at least as good as the bi-encoder CLS-vector baseline, it gives higher recall on queries with rare matching terms, and the storage cost scales linearly with average document length.

from llm_systems_cookbook.nb import bootstrap

import numpy as np

s = bootstrap("02_rag_04_colbertv2_late_interaction")

Synthetic token-embedding setup#

We simulate a small vocabulary where each word has a 64-dim embedding. Documents are sequences of tokens; their “encoder output” is just the stacked word embeddings. Queries share the same word space so a query token that matches a document token gets a high dot product.

This lets us study MaxSim’s behaviour without training a real ColBERT encoder.

rng = np.random.default_rng(0)
VOCAB = [
    # Photosynthesis cluster
    "photosynthesis", "chlorophyll", "carbon", "calvin", "chloroplast", "plant", "light",
    # Mitochondria cluster
    "mitochondria", "atp", "oxidative", "phosphorylation", "membrane", "dna", "crista",
    # Neuron cluster
    "neuron", "action", "potential", "sodium", "channel", "axon", "myelin", "saltatory",
    # CRISPR cluster
    "crispr", "cas9", "rna", "pam", "sequence", "editing", "genome",
    # Other vocabulary
    "the", "of", "and", "a", "is", "in", "to", "by", "with", "that",
    "cell", "protein", "energy", "molecule", "reaction", "process",
]

DIM = 64
word_emb = {w: rng.normal(0, 1, DIM).astype(np.float32) for w in VOCAB}
for w in word_emb:
    word_emb[w] /= np.linalg.norm(word_emb[w]) + 1e-9


def encode_tokens(tokens: list[str]) -> np.ndarray:
    return np.stack([word_emb.get(t, np.zeros(DIM, dtype=np.float32)) for t in tokens])


# 20 documents, each 20-40 tokens
DOCS = [
    {"id": "d0", "topic": "photosynthesis", "tokens": (
        "photosynthesis is the process by which plants convert light to energy "
        "chlorophyll absorbs light and the calvin cycle fixes carbon in the chloroplast"
    ).split()},
    {"id": "d1", "topic": "mitochondria", "tokens": (
        "mitochondria are organelles that generate atp via oxidative phosphorylation "
        "they have a double membrane and contain dna the crista folds increase surface area"
    ).split()},
    {"id": "d2", "topic": "neuron", "tokens": (
        "a neuron action potential is a rapid change in membrane voltage "
        "sodium channels open at threshold saltatory conduction along myelin axon is fast"
    ).split()},
    {"id": "d3", "topic": "crispr", "tokens": (
        "crispr cas9 uses a guide rna to find a dna sequence flanked by a pam motif "
        "the resulting cut enables genome editing"
    ).split()},
]
# Add noise docs (rotations of the same topics with extra stopwords)
for i in range(16):
    base = DOCS[i % 4]
    noisy = base["tokens"] + ["the", "is", "of", "in", "a", "protein", "cell", "process"]
    rng.shuffle(noisy)
    DOCS.append({"id": f"d{4 + i}", "topic": base["topic"], "tokens": noisy[:35]})


QUERIES = [
    {"id": "q0", "text": "plants convert light energy", "topic": "photosynthesis"},
    {"id": "q1", "text": "organelle generates atp", "topic": "mitochondria"},
    {"id": "q2", "text": "saltatory conduction axon", "topic": "neuron"},
    {"id": "q3", "text": "cas9 finds dna using guide rna", "topic": "crispr"},
    {"id": "q4", "text": "calvin cycle fixes carbon", "topic": "photosynthesis"},
    {"id": "q5", "text": "pam motif genome editing", "topic": "crispr"},
]
print(f"docs={len(DOCS)}  queries={len(QUERIES)}")

Bi-encoder baseline#

CLS-style baseline: represent each document as the mean of its token embeddings; score queries by dot product. This is the “one vector per document” regime that ColBERT tries to beat.

def bi_encoder_score(query_tokens: list[str], doc_tokens: list[str]) -> float:
    q = encode_tokens(query_tokens).mean(axis=0)
    d = encode_tokens(doc_tokens).mean(axis=0)
    return float(q @ d)


def rank_with_scorer(scorer) -> dict[str, list[str]]:
    out = {}
    for q in QUERIES:
        q_tokens = q["text"].split()
        scores = [(d["id"], scorer(q_tokens, d["tokens"])) for d in DOCS]
        out[q["id"]] = [doc_id for doc_id, _ in sorted(scores, key=lambda x: -x[1])]
    return out


def recall_at_k(ranking: dict[str, list[str]], k: int = 5) -> float:
    hits = 0
    for q in QUERIES:
        top = set(ranking[q["id"]][:k])
        for d in DOCS:
            if d["id"] in top and d["topic"] == q["topic"]:
                hits += 1
                break
    return hits / len(QUERIES)


bi_rank = rank_with_scorer(bi_encoder_score)
bi_recall = recall_at_k(bi_rank)
print(f"bi-encoder Recall@5 = {bi_recall:.3f}")

ColBERT MaxSim#

\[ \text{MaxSim}(Q, D) = \sum_{q \in Q} \max_{d \in D} q \cdot d \]

Per query token, find the best matching document token. Sum. Because each query token independently finds its home in the document, partial matches on rare tokens (like “calvin” or “pam”) give a bigger score boost than they would under a mean-pooled bi-encoder.

def colbert_score(query_tokens: list[str], doc_tokens: list[str]) -> float:
    q = encode_tokens(query_tokens)
    d = encode_tokens(doc_tokens)
    sim = q @ d.T
    return float(sim.max(axis=1).sum())


colbert_rank = rank_with_scorer(colbert_score)
colbert_recall = recall_at_k(colbert_rank)
print(f"ColBERT   Recall@5 = {colbert_recall:.3f}")

s.check(
    "colbert_at_least_as_good_as_bi_encoder",
    lambda: colbert_recall >= bi_recall,
    msg=f"bi={bi_recall:.3f}  colbert={colbert_recall:.3f}",
)

# Rare-term query: check that MaxSim ranks the right doc first on a
# query dominated by a rare token.
rare_q = ["calvin", "cycle"]
rare_scores_bi = [(d["id"], bi_encoder_score(rare_q, d["tokens"])) for d in DOCS]
rare_scores_cb = [(d["id"], colbert_score(rare_q, d["tokens"])) for d in DOCS]
top_bi = sorted(rare_scores_bi, key=lambda x: -x[1])[0][0]
top_cb = sorted(rare_scores_cb, key=lambda x: -x[1])[0][0]

def is_photosynthesis(doc_id: str) -> bool:
    for d in DOCS:
        if d["id"] == doc_id:
            return d["topic"] == "photosynthesis"
    return False

s.check("colbert_retrieves_rare_term_doc",
        lambda: is_photosynthesis(top_cb),
        msg=f"top = {top_cb} (topic shown above)")

# Storage blow-up bounded: one vector per token vs one per doc.
avg_doc_len = float(np.mean([len(d["tokens"]) for d in DOCS]))
blow_up = avg_doc_len  # per token vs 1 vector per doc
print(f"ColBERT storage vs bi-encoder: {blow_up:.1f}x  (should be ~average doc length)")
s.check(
    "colbert_storage_linear_in_doc_length",
    lambda: blow_up <= 50.0,
    msg=f"blow-up = {blow_up:.1f}x  (expect = avg doc length)",
)
s.check(
    "bi_encoder_is_nontrivial",
    lambda: bi_recall >= 0.33,
    msg=f"bi-encoder recall = {bi_recall:.3f}",
)

Quality, latency, and storage side by side#

Bi-encoder is cheap and compact (one vector per doc). ColBERT adds per-token vectors and per-query-token MaxSim - more accurate on rare-term queries, linearly heavier in both compute and storage. Three bars make the tradeoff legible: recall, per-query latency, and storage multiplier over the bi-encoder.

import matplotlib.pyplot as plt
import time

def bench(scorer, n_iter: int = 200) -> float:
    qt = QUERIES[0]["text"].split()
    for _ in range(5):
        [scorer(qt, d["tokens"]) for d in DOCS]  # warm up
    t0 = time.perf_counter()
    for _ in range(n_iter):
        [scorer(qt, d["tokens"]) for d in DOCS]
    return (time.perf_counter() - t0) * 1000 / n_iter

methods = ["bi-encoder", "ColBERT"]
recalls = [bi_recall, colbert_recall]
lats    = [bench(bi_encoder_score), bench(colbert_score)]
storage = [1.0, avg_doc_len]
colors  = ["tab:gray", "tab:blue"]

fig, axes = plt.subplots(1, 3, figsize=(10, 3.2))
for ax, vals, title, ylab, fmt in [
    (axes[0], recalls, "quality",  "Recall@5",             "{:.2f}"),
    (axes[1], lats,    "latency",  "ms / query",           "{:.2f}"),
    (axes[2], storage, "storage",  "vectors per doc (rel)", "{:.1f}x"),
]:
    ax.bar(methods, vals, color=colors)
    ax.set_ylabel(ylab); ax.set_title(title)
    for i, v in enumerate(vals):
        ax.text(i, v, fmt.format(v), ha="center", va="bottom", fontsize=9)
axes[0].set_ylim(0, 1.02)
fig.tight_layout(); plt.show()

Exercises#

  1. PLAID centroid prefilter. Real ColBERTv2 reduces serve latency by clustering all document tokens into centroids and only interacting a query with documents whose nearest centroids show up. Implement a simple version: run k-means on all token embeddings, cache each document’s centroid assignments, filter candidates to documents sharing any query centroid.

  2. Add score normalisation. Divide MaxSim by |Q| to get a length-normalised score; compare behaviour on single-token vs long queries.

  3. Real ColBERT. pip install pylate and swap the stub embedder for a real trained ColBERTv2 checkpoint on BEIR/scifact dev. Target: Recall@10 ≥ 0.82 (spec §45).

References#

  • ColBERTv2 paper for the PLAID prefilter details.

  • pylate package for a production ColBERT-v2 implementation.

s.summary()
s.save()