Gen AI

Late Interaction Retrieval (ColBERT)

Late interaction models delay the pooling of token embeddings, computing a similarity matrix between every query token and every document token for extreme precision.

In late interaction, instead of comparing two single document vectors, every token in the query is compared against every token in the document.

Why Does This Exist?

Standard dense retrieval uses a "bi-encoder" architecture. A user types a query, the embedding model squashes that query into a single dense vector, and the database calculates the dot product between that vector and millions of pre-squashed document vectors. This is incredibly fast, but it suffers from a massive bottleneck: forcing the meaning of an entire paragraph into a single vector inevitably loses fine-grained details, syntactic relationships, and exact keyword matches.

On the other extreme, "cross-encoders" feed the query and the document together through the transformer model. This allows the self-attention mechanism to perfectly understand how the query relates to the document token-by-token. The recall is phenomenal, but it is brutally slow—you cannot run a cross-encoder against a million documents at search time.

Late Interaction Retrieval (most famously implemented as ColBERT: Contextualized Late Interaction over BERT) was invented to find the perfect middle ground. It achieves near-cross-encoder accuracy with near-bi-encoder speeds by keeping the token-level embeddings intact and delaying the interaction until the very last mathematical step.

Think of It Like This

Summarizing a resume vs. reading line-by-line

Imagine you are a recruiter looking for a "senior Python developer with finance experience."

Standard Dense Retrieval (Bi-encoder): The candidate hands you a one-sentence summary of their entire career. You compare it to a one-sentence summary of your job description. It's very fast to sort through 1,000 resumes this way, but you miss nuance.

Cross-Encoder: You sit down with the candidate and read their entire resume out loud together, discussing how every single bullet point matches the job description. It's incredibly accurate, but it takes an hour per candidate.

Late Interaction (ColBERT): The candidate highlights the 5 most important sentences in their resume. You have 5 requirements in your job description. You simply draw lines matching your 5 requirements to their 5 sentences, keeping only the highest-scoring matches (MaxSim). You didn't compress the resume into one sentence, but you also didn't spend an hour reading it. You interacted late in the process, matching specific points to specific points quickly.

How It Actually Works

The Bi-Encoder Bottleneck (Early Pooling)

In standard embeddings, the outputs of the final transformer layer (one vector per token) are averaged together (mean pooling) or the [CLS] token is used to represent the whole sequence. This is "early pooling." All token-level nuance is destroyed before the search even begins.

The ColBERT Architecture

ColBERT skips the pooling step entirely.

Indexing: When a document is indexed, it is passed through a BERT model. Instead of saving one vector, ColBERT saves a matrix of vectors—one vector for every single token in the document.
Querying: When a query arrives, it is also passed through BERT, resulting in a matrix of token embeddings for the query.

The MaxSim Operation

Because we have a matrix for the query and a matrix for the document, we cannot do a simple dot product. We must compute the similarity between every query token and every document token.

ColBERT uses an operation called MaxSim (Maximum Similarity):

Take the first token of the query (e.g., "Python").
Calculate the dot product between the "Python" vector and every token vector in the document.
Find the maximum value (the document token that best matches "Python").
Repeat this for every token in the query.
Sum up all these maximum values.

This means a document scores highly if it contains some token that strongly matches each token in the query, regardless of word order or distance.

The Storage Tradeoff

Because ColBERT stores a vector for every token rather than every document, the index size is massive. A 500-word document requires 500 vectors instead of 1. To make this practical, ColBERT uses heavy quantization (often compressing the token embeddings down to 1 or 2 bits per dimension), keeping the index size manageable while retaining the token-level routing power.

Show Me the Code

Implementing the MaxSim operation in PyTorch is surprisingly straightforward. This is the core mathematical operation that powers late interaction.

1import torch2import torch.nn.functional as F3
4# Simulate embeddings for a query with 4 tokens and a document with 10 tokens.5# Let's say the embedding dimension is 128.6embed_dim = 1287num_query_tokens = 48num_doc_tokens = 109
10# In reality, these come from a BERT model without pooling11query_embeddings = torch.randn(1, num_query_tokens, embed_dim)12doc_embeddings = torch.randn(1, num_doc_tokens, embed_dim)13
14# 1. Normalize the vectors (so dot product equals cosine similarity)15query_embeddings = F.normalize(query_embeddings, p=2, dim=-1)16doc_embeddings = F.normalize(doc_embeddings, p=2, dim=-1)17
18# 2. Compute the similarity matrix between every query token and every doc token19# Shape: [1, num_query_tokens, num_doc_tokens]20# Using batched matrix multiplication (bmm)21similarity_matrix = torch.bmm(query_embeddings, doc_embeddings.transpose(1, 2))22
23# Let's look at the similarity between query token 0 and all 10 doc tokens24print(f"Similarities for query token 0: {similarity_matrix[0, 0, :].shape}")25
26# 3. The MaxSim operation:27# For each query token, find the maximum similarity score across all doc tokens28# max() returns a tuple of (values, indices), we just want the values29max_similarities = similarity_matrix.max(dim=2).values30
31# 4. Sum the max similarities to get the final document score32final_score = max_similarities.sum(dim=1)33
34print(f"Final ColBERT score for the document: {final_score.item():.4f}")35# -> Similarities for query token 0: torch.Size([10])36# -> Final ColBERT score for the document: 1.4125

Watch Out For

Prohibitive storage costs at scale

Even with aggressive 2-bit quantization, storing token-level embeddings takes 10x to 50x more space than a standard dense index. For databases with billions of documents, a pure ColBERT index is often too expensive to maintain in RAM. This is why late interaction models are frequently used exclusively as rerankers—a fast BM25 or standard dense index retrieves the top 1,000 documents, and ColBERT only computes MaxSim on that small subset.

The Quick Version

Standard embeddings squash documents into a single vector, losing nuance and exact keyword precision.
Cross-encoders compare every token perfectly but are too slow for large-scale retrieval.
Late Interaction (ColBERT) saves a separate vector for every token in the document and the query.
It calculates relevance using MaxSim: finding the best-matching document token for every query token and summing the results.
This preserves token-level accuracy but requires significantly more storage space.