Architectures / Multimodal

🔗 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.

From Khattab and Zaharia (2020). ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. SIGIR 2020. This page is the graph, not the PDF: open it, edit it, verify it, export it.

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

14 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 43 structural checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# 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]

For agents

Also in Multimodal