# Comparing MaxSim with Other Similarity Functions in PyLate

> Explore MaxSim the PyLate similarity function and compare it with cosine similarity dot product and L2 distance Learn how MaxSim excels in late interaction for enhanced text analysis

- Repository: [LightOn/pylate](https://github.com/lightonai/pylate)
- Tags: deep-dive
- Published: 2026-03-06

---

**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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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:

```python
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`](https://github.com/lightonai/pylate/blob/main/pylate/scores/similarity_functions.py):

```python
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:

```python
@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:

```python
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:

```python
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:

```python
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:

```python
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`](https://github.com/lightonai/pylate/blob/main/pylate/scores/similarity_functions.py) and implemented in [`pylate/scores/scores.py`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/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`](https://github.com/lightonai/pylate/blob/main/pylate/scores/similarity_functions.py). The `ColBERT` model class in [`pylate/models/colbert.py`](https://github.com/lightonai/pylate/blob/main/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`.