# ColBERT (Late Interaction Retrieval)

> Khattab and Zaharia 2020 - the document keeps its tokens until scoring time. Each query token takes its best-matching document token and those maxima are summed.

Pick when one vector per document loses too much and a cross-encoder is too slow. Costs a query-by-document matrix per candidate and an index many times larger, so pair it with a cheap first stage.

- Category: Multimodal
- Paper: Khattab and Zaharia (2020). ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. SIGIR 2020 (https://arxiv.org/abs/2004.12832)
- Layers: 14
- Parameters: 75.43M
- Input shape (batchless): 1 × 32
- Output shape: 1
- Verifier verdict: pass
- Graph JSON: https://neurarch.com/templates/colbert/model.json
- Open on the canvas: https://neurarch.com/?template=colbert

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | Query tokens | Input | shape=[1, 32] | 1 × 32 |
| 2 | Token Embed | Embedding | vocabSize=30522 | 1 × 32 × 768 |
| 3 | Pos Encoding | Positional Encoding | maxLen=512 | 1 × 32 × 768 |
| 4 | Query Encoder 1 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 32 × 768 |
| 5 | Query Encoder 2 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 32 × 768 |
| 6 | Down-project to 128 | Linear | outFeatures=128, inFeatures=768 | 1 × 32 × 128 |
| 7 | Document tokens | Input | shape=[1, 180] | 1 × 180 |
| 8 | Token Embed (shared) | Embedding | vocabSize=30522 | 1 × 180 × 768 |
| 9 | Pos Encoding | Positional Encoding | maxLen=512 | 1 × 180 × 768 |
| 10 | Doc Encoder 1 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 180 × 768 |
| 11 | Doc Encoder 2 | Transformer Block | embedDim=768, numHeads=12, ffDim=3072 | 1 × 180 × 768 |
| 12 | Down-project to 128 | Linear | outFeatures=128, inFeatures=768 | 1 × 180 × 128 |
| 13 | MaxSim | Late Interaction (ColBERT MaxSim) | embedDim=128 | 1 |
| 14 | relevance score | Output |  | 1 |

## Verifier findings

- **info** `deep-no-norm`: 11 layers with no BatchNorm, LayerNorm, or GroupNorm. Without normalization, activations can explode or vanish across layers, causing slow or unstable training. Fix: Add BatchNorm after Conv2d (CV tasks), LayerNorm after attention/FFN (NLP/LLM), or GroupNorm for small batch sizes.

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class LateInteraction(nn.Module):
    """ColBERT MaxSim: the best-matching document token for each query token,
    summed. Both sequence axes are consumed and one score comes out. The
    document keeps its tokens until scoring time, which is the whole difference
    from a single-vector retriever, and also the whole cost."""

    def __init__(self, embed_dim: int = 128, similarity: str = "cosine"):
        super().__init__()
        self.embed_dim = embed_dim
        self.similarity = similarity

    def forward(self, query: torch.Tensor, doc=None) -> torch.Tensor:
        if doc is None:
            doc = query
        q, d = query, doc
        if self.similarity == "cosine":
            q, d = F.normalize(q, dim=-1), F.normalize(d, dim=-1)
        sim = q @ d.transpose(-1, -2)          # [.., Lq, Ld]
        return sim.max(dim=-1).values.sum(dim=-1, keepdim=True)


class ColBERTLateInteractionRetrieval(nn.Module):
    def __init__(self):
        super().__init__()

        self.embedding_1 = nn.Embedding(30522, 768)
        self.transformerBlock_1 = nn.TransformerEncoderLayer(d_model=768, nhead=12, dim_feedforward=3072, batch_first=True)
        self.transformerBlock_2 = nn.TransformerEncoderLayer(d_model=768, nhead=12, dim_feedforward=3072, batch_first=True)
        self.linear_1 = nn.Linear(768, 128)
        self.embedding_2 = nn.Embedding(30522, 768)
        self.transformerBlock_3 = nn.TransformerEncoderLayer(d_model=768, nhead=12, dim_feedforward=3072, batch_first=True)
        self.transformerBlock_4 = nn.TransformerEncoderLayer(d_model=768, nhead=12, dim_feedforward=3072, batch_first=True)
        self.linear_2 = nn.Linear(768, 128)
        self.lateInteraction_1 = LateInteraction(embed_dim=128, similarity="cosine")

    def forward(self, src, tgt=None):
        # Query tokens shape: [1,32]
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md