Scoring

sentence_transformers.multi_vector_encoder.scoring provides the late-interaction scoring functions used by training losses. Pass one of these (or a configured callable) as the similarity_fct parameter on the multi-vector losses to switch between ColBERT-style MaxSim and XTR-style global top-k scoring.

ColBERT scoring

ColBERT (MaxSim) contrastive scoring for in-batch negatives.

Takes (Q_query, q_tokens, dim) query embeddings and (Q_doc, N, d_tokens, dim) stacked per-query document groups and returns the full (Q_query, Q_doc * N) score matrix (float32 regardless of the input dtype) with query-major ordering: scores[i, j*N + n] is the MaxSim of query i against the n-th document in doc-group j. When called with matched Q_query == Q_doc, the positive for query i sits at column i*N.

The document axis is iterated group-by-group so that only one (Q_query, Q_doc, q_tokens, d_tokens) intermediate is live at a time. Pass this as similarity_fct to a losses loss (the default), or xtr_scores() for XTR-style scoring. length_normalize=True divides each score by the real query token count (MeanMaxSim), removing the query-length dependence of the score scale. chunk_elements budgets the padded documents plus the 4D scoring intermediate of one group, as in maxsim() (whose docstring gives the exact shapes): lower it to cut training memory. The groups are scored one at a time, so a per-group budget bounds the peak, and counting only the intermediate under-provisions when the query side is small.

Pairwise ColBERT (MaxSim) scoring for matched (query_i, document_i) pairs.

Takes (batch_size, q_tokens, dim) query embeddings and (batch_size, d_tokens, dim) document embeddings and returns a (batch_size,) float32 score vector, one MaxSim score per pair. A thin delegation to maxsim_pairwise() with the scoring package’s keyword convention, interchangeable with xtr_scores_pairwise() as the similarity_fct of MultiVectorMarginMSELoss. length_normalize=True divides each score by the real query token count (MeanMaxSim), and chunk_elements budgets the padded queries and documents plus the scoring intermediate (both as in maxsim_pairwise(), whose docstring gives the exact shapes). Counting only the intermediate under-provisions, since the padding dominates whenever the query side is small.

Compute MaxSim scores for knowledge distillation.

The query embeddings have shape (batch_size, q_tokens, dim). The document embeddings have the stacked-per-query shape (batch_size, n_ways, d_tokens, dim): for each query, n_ways candidate documents (typically a positive plus several negatives) were retrieved and scored by a teacher. This function returns (batch_size, n_ways) MaxSim scores suitable for KL-distillation against the teacher scores.

Parameters:

queries_embeddings – (batch_size, q_tokens, dim).
documents_embeddings – (batch_size, n_ways, d_tokens, dim).
queries_mask – optional (batch_size, q_tokens) mask.
documents_mask – optional (batch_size, n_ways, d_tokens) mask.
chunk_elements – element budget for the padded (chunk, q_tokens, dim) queries and (chunk, d_tokens, dim) documents plus the (chunk, q_tokens, d_tokens) scoring intermediate, forwarded to maxsim_pairwise() per n_ways column. The padding is the bigger half whenever the query side is small, so budgeting for the intermediate alone under-provisions. Lower it to cut training memory. Defaults to None (maxsim_pairwise’s 100M-element budget).
length_normalize – divide each score by the real query token count (MeanMaxSim). Defaults to False.

Returns:

(batch_size, n_ways) score tensor, float32 regardless of the input dtype.

MeanMaxSim scoring

MaxSim divided by each query’s real token count, so scores are comparable across query lengths. Rankings within a query are unchanged. Pair these with model.similarity_fn_name = "meanmaxsim" so evaluation and the model card score the way training did.

MeanMaxSim contrastive scoring: colbert_scores() divided by each query’s real token count.

Pair it with model.similarity_fn_name = "meanmaxsim" so evaluation scores the way training did.

length_normalize defaults to True here (the Mean in the name): False recovers plain colbert_scores(), so the whole ColBERT family accepts the same keywords.

sentence_transformers.multi_vector_encoder.scoring.mean_colbert_scores_pairwise(queries_embeddings: list | ndarray | Tensor, documents_embeddings: list | ndarray | Tensor, queries_mask: Tensor | None = None, documents_mask: Tensor | None = None, chunk_elements: int | None = None, length_normalize: bool = True) → Tensor[source]: MeanMaxSim pairwise scoring, the colbert_scores_pairwise() counterpart of mean_colbert_scores(). Use it as MultiVectorMarginMSELoss’s similarity_fct. length_normalize defaults to True here, and False recovers plain colbert_scores_pairwise().

sentence_transformers.multi_vector_encoder.scoring.mean_colbert_kd_scores(queries_embeddings: list | ndarray | Tensor, documents_embeddings: list | ndarray | Tensor, queries_mask: Tensor | None = None, documents_mask: Tensor | None = None, chunk_elements: int | None = None, length_normalize: bool = True) → Tensor[source]: MeanMaxSim listwise KD scoring, the colbert_kd_scores() counterpart of mean_colbert_scores(). Use it as MultiVectorDistillKLDivLoss’s similarity_fct. length_normalize defaults to True here, and False recovers plain colbert_kd_scores().

XTRScores

class sentence_transformers.multi_vector_encoder.scoring.XTRScores(top_k: int = 256, *, chunk_elements: int | None = None)[source]

Configured, reusable xtr_scores() callable for use as a loss similarity_fct.

Stores top_k / chunk_elements so they don’t have to be re-passed on every call (the bare function would otherwise need functools.partial()). See xtr_scores() for the scoring math and shapes.

Parameters:

top_k – Positive number of top token matches to retain per query token across all Q*N documents.
chunk_elements – Element budget for the chunked matmul phase, see xtr_scores().

XTR (eXtendable Token Retrieval) contrastive scoring with global top-k token retrieval.

For each query token, the top-k matches are selected globally across all in-batch document tokens (simulating retrieval from an index). Returns the full (Q, Q*N) cross-product score matrix with query-major ordering: scores[i, j*N + n] is query i against query j’s n-th document. The positive for query i sits at column i*N. As in maxsim(), the matmul runs in the input dtype (integer embeddings are upcast to float32 first) and the per-query-token accumulation and returned scores are float32.

Each score is the sum of the query’s retrieved per-token maxima divided by Z, the number of query tokens that retrieved at least one of the document’s tokens (Lee et al. 2023, eq. 5). This deviates from PyLate / PrimeQA, which divide by the count of positive per-token maxima instead.

Parameters:

queries_embeddings – (Q, q_tokens, dim) query embeddings.
documents_embeddings – (Q, N, d_tokens, dim) stacked per-query document groups.
queries_mask – optional (Q, q_tokens) mask. If None, all-zero query rows are treated as padding.
documents_mask – optional (Q, N, d_tokens) mask. If None, one is derived by treating all-zero document rows as padding (like maxsim()).
top_k – Positive number of top token matches to retain per query token across all Q*N documents.
chunk_elements – Element budget for the matmul + masked_fill phase, matching maxsim()’s parameter of the same name: documents are scored in chunks packed to stay under the budget. The chunks are concatenated before the global top-k, so scoring semantics are unchanged and the full score grid is still materialized: this trims the transient matmul peak, not the overall peak. Defaults to None (single matmul).

Notes

Adapted from PyLate / PrimeQA (Apache 2.0). Pass this (or XTRScores, a configured reusable instance) as similarity_fct to any losses loss to switch from ColBERT-style MaxSim scoring to XTR-style top-k scoring. To compile the hot path, wrap it: similarity_fct=torch.compile(xtr_scores).

Pairwise XTR scoring: compute the XTR score for matched (query_i, document_i) pairs.

Returns a float32 1D tensor of length batch_size. XTR’s top-k runs globally over the batch’s pooled document tokens, so unlike MaxSim a pair’s score shifts with the batch composition whenever top_k is below the pooled token count.

XTR scoring for knowledge distillation.

Same global top-k scoring as xtr_scores(), but returns each query’s own N-way document scores (Q, N) instead of the full (Q, Q*N) cross-product, matching the interface expected by MultiVectorDistillKLDivLoss.

XTRKDScores

class sentence_transformers.multi_vector_encoder.scoring.XTRKDScores(top_k: int = 256, *, chunk_elements: int | None = None)[source]: Configured, reusable xtr_kd_scores() callable (KD (Q, N) output). See XTRScores.