Comparing MaxSim with Other Similarity Functions in PyLate

MaxSim is currently the only built-in similarity function in PyLate, implementing ColBERT-style late interaction through token-level maximum similarity scoring, while cosine similarity, dot product, and L2 distance are not implemented but can be added via the extensible SimilarityFunction enum.

PyLate is an open-source library by LightOn that implements late interaction retrieval models like ColBERT. When configuring neural search pipelines, understanding the available similarity functions is essential for optimizing retrieval quality. This article examines how MaxSim operates within PyLate and how it conceptually compares to other common similarity measures.

What Is MaxSim in PyLate?

MaxSim (Maximum Similarity) is the default and currently only similarity function available in PyLate. It enables late interaction between query and document token embeddings, preserving fine-grained matching signals that are typically lost in traditional bi-encoder approaches.

Mathematical Definition

For a query (Q = {q_i}) with (n) tokens and a document (D = {d_j}) with (m) tokens, MaxSim is defined as:

[ \text{MaxSim}(Q, D) = \sum_{i=1}^{n} \max_{j=1}^{m} (q_i \cdot d_j) ]

This formula computes the dot product between every query token and every document token, takes the maximum similarity for each query token across all document tokens, and sums these maxima to produce the final relevance score.

Implementation Details

In pylate/scores/scores.py, the colbert_scores function implements this computation using optimized PyTorch operations:

1. Token-wise similarity matrix: An einsum operation (ash,bth -> abst) builds a 4-D tensor containing all dot-product similarities between each query token and each document token.
2. Masking: Query and document masks zero out padded tokens so they don't affect the result.
3. Max-over-tokens: For each query token, the maximum similarity across all document tokens is taken (max(axis=-1)).
4. Summation: The per-token maxima are summed, yielding a single score per query-document pair (sum(axis=-1)).

How MaxSim Compares to Other Similarity Functions

While MaxSim is the only currently implemented option in PyLate, understanding how it differs from classical similarity measures helps clarify its unique properties for neural retrieval.

Measure	Definition	Typical Use Case	Status in PyLate
MaxSim	(\sum_i \max_j (q_i \cdot d_j))	ColBERT-style late interaction retrieval	Implemented (default)
Cosine similarity	(\frac{u \cdot v}{|u||v|})	Dense retrieval, semantic search	Not implemented
Dot product	(u \cdot v)	Bi-encoder ranking where magnitude matters	Not implemented
L2 distance	(|u-v|_2)	Retrieval with metric-learning losses	Not implemented

Granularity Differences

MaxSim operates on token embeddings rather than pooled sentence vectors. While cosine similarity computes the angle between two single vectors representing entire texts, MaxSim calculates the maximum similarity for each query token against all document tokens and sums these maxima. This preserves fine-grained term matching signals that cosine similarity loses during pooling, making MaxSim particularly effective for late interaction retrieval where precise token alignment matters.

Computational Complexity

MaxSim requires (O(|Q| \times |D|)) operations per pair, where (|Q|) and (|D|) are the number of tokens. However, PyLate's implementation uses highly vectorized PyTorch operations that batch efficiently across multiple queries and documents.

Cosine and dot product approaches are (O(1)) per pair after pooling, making them faster for brute-force search but potentially less accurate for precise term matching.

Length Sensitivity

MaxSim sums contributions from all query tokens, meaning longer queries naturally accumulate higher scores. The implementation in colbert_scores uses masks to handle variable lengths, but the raw sum scales with query length.

Cosine similarity normalizes by vector magnitude, making it invariant to text length. Dot product grows with vector norm, which correlates with length in many embedding models.

Where MaxSim Is Implemented in PyLate

The MaxSim functionality is distributed across several key files in the PyLate codebase:

File	Purpose
pylate/scores/similarity_functions.py	Declares the SimilarityFunction enum, which currently only contains MaxSim = "MaxSim". Provides factory methods to_similarity_fn and to_similarity_pairwise_fn to retrieve the concrete scoring functions.
pylate/scores/scores.py	Contains the core implementation: colbert_scores for batch scoring and colbert_scores_pairwise for one-to-one scoring. These functions implement the token-wise max-similarity computation with masking support.
pylate/models/colbert.py	The ColBERT model class stores the selected similarity function name in _similarity_fn_name and resolves it via the enum's factory methods, caching the callables in _similarity and _similarity_pairwise.
pylate/utils/tensor.py	Provides the convert_to_tensor utility used by scoring functions to accept NumPy arrays, lists, or PyTorch tensors uniformly.

Extending PyLate with Custom Similarity Functions

Although only MaxSim is available, the architecture is designed for extensibility. Adding a new similarity function requires three steps:

1. Implement a scoring routine following the expected signature:

def my_similarity_scores(
    queries: torch.Tensor, 
    documents: torch.Tensor,
    queries_mask: torch.Tensor | None = None,
    documents_mask: torch.Tensor | None = None
) -> torch.Tensor:
    """
    Compute similarity between queries and documents.
    
    Args:
        queries: Tensor of shape (batch_size, num_query_tokens, dim)
        documents: Tensor of shape (batch_size, num_doc_tokens, dim)
        queries_mask: Boolean tensor for query tokens (1=keep, 0=mask)
        documents_mask: Boolean tensor for document tokens
    
    Returns:
        Similarity scores of shape (batch_size, batch_size) for batch scoring
        or (batch_size,) for pairwise scoring
    """
    # Implementation here

    pass

2. Add an entry to the SimilarityFunction enum in pylate/scores/similarity_functions.py:

class SimilarityFunction(str, Enum):
    MaxSim = "MaxSim"
    Cosine = "Cosine"  # New entry

3. Update the factory methods in the same file to map the new enum value to your implementation:

@staticmethod
def to_similarity_fn(value: str):
    if value == SimilarityFunction.MaxSim:
        from .scores import colbert_scores
        return colbert_scores
    elif value == SimilarityFunction.Cosine:
        return my_similarity_scores  # Your new function

    else:
        raise ValueError(f"Unknown similarity function: {value}")

@staticmethod
def to_similarity_pairwise_fn(value: str):
    # Similar mapping for pairwise scoring

    pass

After these changes, users can switch to the new similarity function at runtime:

from pylate.models import ColBERT

model = ColBERT("bert-base-uncased")
model.similarity_fn_name = "Cosine"  # Uses the new implementation

Practical Code Examples

Direct Usage of MaxSim Scoring

You can use the underlying colbert_scores function directly without instantiating the full model:

import torch
from pylate.scores import colbert_scores

# Create dummy token embeddings

# Shape: (batch_size=2, num_tokens=3, dim=4)

queries = torch.randn(2, 3, 4)
documents = torch.randn(2, 5, 4)

# Optional masks to ignore padding (1=keep, 0=mask)

q_mask = torch.tensor([[1, 1, 1], [1, 0, 1]])
d_mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 0, 1, 1]])

# Compute scores

scores = colbert_scores(
    queries, 
    documents, 
    queries_mask=q_mask, 
    documents_mask=d_mask
)
print(scores)  # Output: tensor of shape (2, 2)

Accessing MaxSim Through the ColBERT Model

The ColBERT model class provides a higher-level interface that manages the similarity function internally:

from pylate.models import ColBERT

# Initialize model with default MaxSim

model = ColBERT("bert-base-uncased")

# Check current similarity function

print(model.similarity_fn_name)  # Output: MaxSim

# The model caches the scoring functions

print(model._similarity is not None)        # True: batch scoring

print(model._similarity_pairwise is not None)  # True: pairwise scoring

Pairwise Scoring for Matched Pairs

When you only need to score query-document pairs where query i corresponds to document i (the diagonal), use the pairwise variant for efficiency:

from pylate.scores import colbert_scores_pairwise

# One-to-one scoring: query[i] vs document[i]

pairwise_scores = colbert_scores_pairwise(
    queries, 
    documents,
    queries_mask=q_mask,
    documents_mask=d_mask
)
print(pairwise_scores)  # Output: tensor of shape (2,)

Summary

• MaxSim is the default: PyLate currently ships with only one similarity function—MaxSim—which implements ColBERT-style late interaction scoring.
• Token-level granularity: Unlike cosine similarity or dot product that operate on pooled sentence embeddings, MaxSim computes similarity at the token level, preserving fine-grained matching signals.
• Implementation locations: The function is declared in pylate/scores/similarity_functions.py and implemented in pylate/scores/scores.py via colbert_scores and colbert_scores_pairwise.
• Extensible architecture: Although only MaxSim is available, the SimilarityFunction enum and factory methods are designed to accommodate custom similarity measures.
• Flexible access: Developers can use MaxSim directly through the scoring functions or indirectly through the ColBERT model class, which manages similarity function resolution and caching.

Frequently Asked Questions

Is MaxSim the only similarity function available in PyLate?

Yes, as of the current implementation, MaxSim is the only built-in similarity function available in PyLate. The SimilarityFunction enum in pylate/scores/similarity_functions.py currently only defines MaxSim = "MaxSim". However, the architecture is designed to be extensible, allowing developers to add custom similarity functions by extending the enum and implementing the corresponding scoring logic.

How does MaxSim differ from cosine similarity in neural retrieval?

MaxSim operates on token-level embeddings rather than pooled sentence vectors. While cosine similarity computes the angle between two single vectors representing entire texts, MaxSim calculates the maximum similarity for each query token against all document tokens and sums these maxima. This preserves fine-grained term matching signals that cosine similarity loses during pooling, making MaxSim particularly effective for late interaction retrieval where precise token alignment matters.

Can I use dot product or L2 distance instead of MaxSim in PyLate?

Currently, dot product and L2 distance are not implemented in PyLate. The library focuses on late interaction retrieval using the MaxSim scoring mechanism. However, you can extend PyLate to support these measures by implementing a scoring function with the appropriate signature (accepting query embeddings, document embeddings, and optional masks), adding a new entry to the SimilarityFunction enum in pylate/scores/similarity_functions.py, and updating the factory methods to_similarity_fn and to_similarity_pairwise_fn to return your implementation when the corresponding enum value is requested.

Where is the MaxSim scoring logic actually implemented in the source code?

The MaxSim scoring logic is implemented in pylate/scores/scores.py within the colbert_scores function for batch scoring and colbert_scores_pairwise for one-to-one scoring. The function is declared and exposed through the SimilarityFunction enum located in pylate/scores/similarity_functions.py. The ColBERT model class in pylate/models/colbert.py stores the selected similarity function name in _similarity_fn_name and resolves it via the enum's factory methods, caching the callables in _similarity and _similarity_pairwise.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}