Losses

sentence_transformers.multi_vector_encoder.losses defines the loss functions for training MultiVectorEncoder models. See Training Overview for guidance on which loss to pick for which dataset format.

MultiVectorMultipleNegativesRankingLoss

class sentence_transformers.multi_vector_encoder.losses.MultiVectorMultipleNegativesRankingLoss(model: MultiVectorEncoder, *, scale: float = 1.0, similarity_fct: Callable | None = None, mini_batch_size: int | None = None, score_mini_batch_size: int | None = None, gather_across_devices: bool = False)[source]

In-batch negatives contrastive loss for MultiVectorEncoder models.

For each query in the batch, the matched positive document is treated as the positive sample, and all other in-batch documents (plus any explicitly-provided hard negatives) serve as negatives. Scoring uses MaxSim by default. Pass a different similarity_fct (e.g. XTRScores) to switch scoring strategies without changing the loss.

Parameters:

model – A MultiVectorEncoder model.
scale – 1 / temperature. Scores are multiplied by scale before cross-entropy. Defaults to 1.0 (temperature=1.0), matching PyLate. Unlike cosine similarity (bounded to [-1, 1], where ST’s dense MultipleNegativesRankingLoss uses scale=20.0 to amplify the narrow range), MaxSim is an unbounded sum over query-token similarities (range ~[0, num_query_tokens]), so a scale near 1.0 makes sense, the same reason the dense loss recommends scale=1 for dot-product similarity. With MeanMaxSim scoring (mean_colbert_scores()), each score is divided by its query’s token count, so a scale of roughly the average query length is a reasonable start.
similarity_fct – Scoring callable. Receives queries (Q, q_tokens, dim) and stacked documents (Q, N, d_tokens, dim) and returns (Q, Q*N) with query-major ordering. Defaults to colbert_scores(). Pass XTRScores for XTR-style scoring.
mini_batch_size – Maximum number of rows per model forward during the embedding phase. The document columns are merged into one batch_size * N-row batch and split into row chunks, each embedded at its own width, so a single long outlier document only widens its own chunk. Chunking is exact (no cross-row interactions), but every chunk’s activations stay live for the backward pass, so this bounds padding waste rather than total memory: reach for CachedMultiVectorMultipleNegativesRankingLoss to bound memory by chunk size. None runs the merged batch in one forward.
score_mini_batch_size – If set, queries are processed in chunks of this size during the scoring phase. Useful to bound transient scoring memory for large effective batch sizes. Gradients still flow through a single backward. It chunks the query axis: to chunk the document axis as well, bind a budget into the scorer with similarity_fct=partial(colbert_scores, chunk_elements=...).
gather_across_devices – If True, AllGather document embeddings (and masks) across DDP ranks so that every rank’s queries see the global batch of documents. Useful for very large effective batches.

Requirements:

(anchor, positive) pairs, (anchor, positive, negative) triplets, or (anchor, positive, negative_1, …, negative_n) n-tuples

Inputs:

Inputs	Labels
(anchor, positive) pairs	none
(anchor, positive, negative) triplets	none
(anchor, positive, negative_1, …, negative_n)	none

Recommendations:

Use BatchSamplers.NO_DUPLICATES (docs) to ensure that no in-batch negatives are duplicates of the anchor or positive samples.

Relations:

CachedMultiVectorMultipleNegativesRankingLoss is equivalent to this loss, but it caches the embedding pass so that much higher batch sizes fit in the same memory, at the cost of being slower.
The dense counterpart is MultipleNegativesRankingLoss, which scores pooled embeddings instead of per-token MaxSim.

Example

from datasets import Dataset

from sentence_transformers import MultiVectorEncoder, MultiVectorEncoderTrainer
from sentence_transformers.multi_vector_encoder.losses import MultiVectorMultipleNegativesRankingLoss

model = MultiVectorEncoder("answerdotai/ModernBERT-base")
train_dataset = Dataset.from_dict(
    {
        "query": ["What is the capital of France?", "Who painted the Mona Lisa?"],
        "positive": ["Paris is the capital of France.", "Leonardo da Vinci painted the Mona Lisa."],
        "negative": ["Berlin is the capital of Germany.", "Van Gogh painted The Starry Night."],
    }
)
loss = MultiVectorMultipleNegativesRankingLoss(model)

trainer = MultiVectorEncoderTrainer(model=model, train_dataset=train_dataset, loss=loss)
trainer.train()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

CachedMultiVectorMultipleNegativesRankingLoss

class sentence_transformers.multi_vector_encoder.losses.CachedMultiVectorMultipleNegativesRankingLoss(model: MultiVectorEncoder, *, scale: float = 1.0, similarity_fct: Callable | None = None, mini_batch_size: int = 32, mini_batch_num_tokens: int | None = None, score_mini_batch_size: int | None = None, gather_across_devices: bool = False, show_progress_bar: bool = False)[source]

A GradCache version of MultiVectorMultipleNegativesRankingLoss.

Enables much larger effective batch sizes than the non-cached loss at the cost of being slightly slower: embeddings are computed in chunks of mini_batch_size (under torch.no_grad), the cross-entropy loss is computed once over the full batch, the per-embedding gradients are cached, and a second pass re-runs the embedding forward in chunks with gradients enabled, feeding the cached gradients into the final loss.backward().

Parameters:

model – A MultiVectorEncoder.
scale – 1 / temperature. Scores are multiplied by scale before cross-entropy. Defaults to 1.0 (temperature=1.0). Unlike cosine similarity (bounded to [-1, 1], where ST’s dense MultipleNegativesRankingLoss uses scale=20.0 to amplify the narrow range), MaxSim is an unbounded sum over query-token similarities (range ~[0, num_query_tokens]), so a scale near 1.0 makes sense, the same reason the dense loss recommends scale=1 for dot-product similarity. With MeanMaxSim scoring (mean_colbert_scores()), each score is divided by its query’s token count, so a scale of roughly the average query length is a reasonable start.
similarity_fct – Scoring callable. Defaults to colbert_scores(). Pass XTRScores for XTR-style scoring.
mini_batch_size – Chunk size for the embedding forward / backward pass. Keep small enough that a single chunk fits in GPU memory.
mini_batch_num_tokens – If set, the embedding mini-batches are packed by total (non-padding) token count instead of by mini_batch_size sequences, which speeds up training on variable-length data. Most effective for models that avoid padded compute, e.g. flash attention with input flattening. See the Speeding up Inference documentation for details.
score_mini_batch_size – Chunk size for the scoring phase (independent of mini_batch_size). Smaller values trim transient scoring intermediates (Q, Q*N, q_tokens, d_tokens) which are usually the bottleneck at large effective batch sizes. Defaults to mini_batch_size.
gather_across_devices – If True, AllGather document embeddings across DDP ranks.
show_progress_bar – If True, show a TQDM progress bar for the embedding / scoring steps.

References

GradCache, the caching strategy this loss implements: https://github.com/luyug/GradCache
Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup: https://huggingface.co/papers/2101.06983

Requirements:

(anchor, positive) pairs, (anchor, positive, negative) triplets, or (anchor, positive, negative_1, …, negative_n) n-tuples
Should be used with a large per_device_train_batch_size and a low mini_batch_size for superior performance, at the cost of slower training than MultiVectorMultipleNegativesRankingLoss.

Inputs:

Inputs	Labels
(anchor, positive) pairs	none
(anchor, positive, negative) triplets	none
(anchor, positive, negative_1, …, negative_n)	none

Recommendations:

Use BatchSamplers.NO_DUPLICATES (docs) to ensure that no in-batch negatives are duplicates of the anchor or positive samples.

Relations:

Equivalent to MultiVectorMultipleNegativesRankingLoss, but with caching that allows for much higher batch sizes without extra memory usage, at the cost of slower training.
The dense counterpart is CachedMultipleNegativesRankingLoss.

Example

from datasets import Dataset

from sentence_transformers import (
    MultiVectorEncoder,
    MultiVectorEncoderTrainer,
    MultiVectorEncoderTrainingArguments,
)
from sentence_transformers.multi_vector_encoder.losses import (
    CachedMultiVectorMultipleNegativesRankingLoss,
)

model = MultiVectorEncoder("answerdotai/ModernBERT-base")
train_dataset = Dataset.from_dict(
    {
        "query": ["What is the capital of France?", "Who painted the Mona Lisa?"],
        "positive": ["Paris is the capital of France.", "Leonardo da Vinci painted the Mona Lisa."],
        "negative": ["Berlin is the capital of Germany.", "Van Gogh painted The Starry Night."],
    }
)
# The large batch is what this loss buys you, mini_batch_size bounds the memory it costs.
loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=8)
args = MultiVectorEncoderTrainingArguments(
    output_dir="models/multivector-cached-mnrl", per_device_train_batch_size=64
)

trainer = MultiVectorEncoderTrainer(model=model, args=args, train_dataset=train_dataset, loss=loss)
trainer.train()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

MultiVectorDistillKLDivLoss

class sentence_transformers.multi_vector_encoder.losses.MultiVectorDistillKLDivLoss(model: MultiVectorEncoder, *, similarity_fct: Callable | None = None, temperature: float = 1.0, student_temperature: float | None = None, teacher_temperature: float | None = None, mini_batch_size: int | None = None)[source]

KL-divergence distillation loss for MultiVectorEncoder models.

For each query, the dataset provides N candidate documents (a positive plus at least one negative) and teacher scores (N,). This loss computes the model’s MaxSim scores against the same documents and minimises the KL divergence between the softmaxed teacher and student distributions.

Parameters:

model – A MultiVectorEncoder.
similarity_fct – Callable that, given queries (Q, q_tokens, dim) and stacked docs (Q, N, d_tokens, dim), returns (Q, N) scores. Defaults to colbert_kd_scores(). Pass XTRKDScores for XTR-style scoring.
temperature – Temperature applied to both the student and teacher logits before softmax, unless overridden per side. Defaults to 1.0. The loss is multiplied by the student temperature squared, which keeps the gradient magnitude comparable across temperatures (Hinton et al., 2015) in the regime that paper covers, a shared temperature at or above 1.0. Below that the gradient falls off instead, and a sharp student temperature shrinks the reported loss strongly (by 1e-6 at 0.001), so scale the KD weight back up when weighing this loss against another.
student_temperature – Student-side override of temperature. Sharpening it far below the spread of the student’s own scores collapses that distribution to one-hot, and once the teacher’s is one-hot too the loss and its gradient are exactly zero. Defaults to None.
teacher_temperature – Teacher-side override of temperature. Match it to the spread of your teacher’s scores rather than to a fixed value: a float32 softmax underflows to exact zeros once a row’s spread divided by the temperature exceeds about 100, and every candidate that underflows drops out of the KL entirely, which is the ranking information distillation exists to transfer. Defaults to None.
mini_batch_size – Maximum number of rows per model forward. The merged batch_size * n_ways document batch is split into row chunks, each re-trimmed to its own longest document, so a single long outlier document only widens its own chunk. Chunking is exact for this loss (no cross-row interactions). None (default) runs one merged forward.

References

For more details, please refer to https://huggingface.co/papers/2010.11386
Separate student and teacher temperatures follow DenseOn / LateOn: https://huggingface.co/papers/2607.27178

Requirements:

(query, document_1, …, document_N) examples with at least two candidate documents
Labels containing the teacher model’s score for each candidate document, shape (batch_size, N)
resolve_ids() produces this shape from ID-only KD datasets

Inputs:

Inputs	Labels
(query, document_1, …, document_N)	[Teacher(query, document_i) for i in 1..N]

Relations:

MultiVectorMarginMSELoss distills the same teacher scores, but matches margins with MSE rather than matching the whole distribution with KL divergence.
The dense counterpart is DistillKLDivLoss.

Example

from datasets import Dataset

from sentence_transformers import MultiVectorEncoder, MultiVectorEncoderTrainer
from sentence_transformers.multi_vector_encoder.losses import MultiVectorDistillKLDivLoss

model = MultiVectorEncoder("answerdotai/ModernBERT-base")
# One label per document column, holding that document's teacher score.
train_dataset = Dataset.from_dict(
    {
        "query": ["What is the capital of France?", "Who painted the Mona Lisa?"],
        "positive": ["Paris is the capital of France.", "Leonardo da Vinci painted the Mona Lisa."],
        "negative": ["Berlin is the capital of Germany.", "Van Gogh painted The Starry Night."],
        "label": [[9.5, 2.1], [8.8, 1.4]],
    }
)
# Consider sharpening both distributions with a temperature below 1.0 when scoring many candidates.
loss = MultiVectorDistillKLDivLoss(model, temperature=0.25)

trainer = MultiVectorEncoderTrainer(model=model, train_dataset=train_dataset, loss=loss)
trainer.train()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

MultiVectorMarginMSELoss

class sentence_transformers.multi_vector_encoder.losses.MultiVectorMarginMSELoss(model: MultiVectorEncoder, *, similarity_fct: Callable | None = None, mini_batch_size: int | None = None)[source]

Margin-MSE distillation loss for MultiVectorEncoder models.

Adapted from the dense sentence_transformers.sentence_transformer.losses.MarginMSELoss. Given a query, a positive document, and one or more negative documents, plus teacher margins score(q, pos) - score(q, neg), the student’s MaxSim margins are MSE-matched to the teacher’s.

Parameters:

model – A MultiVectorEncoder.
similarity_fct – A pairwise scoring function, called as similarity_fct(queries_embeddings, documents_embeddings, queries_mask=..., documents_mask=...) like in the other multi-vector losses. Defaults to colbert_scores_pairwise(). Pass xtr_scores_pairwise() for XTR-style scoring.
mini_batch_size – Maximum number of rows per model forward. The merged document batch is split into row chunks, each re-trimmed to its own longest document, so a single long outlier document only widens its own chunk. Chunking is exact for this loss (no cross-row interactions). None (default) runs one merged forward.

References

For more details, please refer to https://huggingface.co/papers/2010.02666

Requirements:

(query, positive, negative_1, …, negative_k) examples
Labels holding either the teacher margins, shape (batch_size, k), or the raw teacher scores, shape (batch_size, k + 1), which are converted to margins internally. With a single negative, a flat (batch_size,) margin is also accepted.
Usually used with a finetuned teacher M in a knowledge distillation setup

Inputs:

Inputs	Labels
(query, positive, negative)	M(query, positive) - M(query, negative)
(query, positive, negative)	[M(query, positive), M(query, negative)]
(query, positive, negative_1, …, negative_k)	[M(query, positive) - M(query, negative_i) for i in 1..k]
(query, positive, negative_1, …, negative_k)	[M(query, positive), M(query, negative_1), …, M(query, negative_k)]

Relations:

MultiVectorDistillKLDivLoss distills the same teacher scores, but matches the whole distribution with KL divergence rather than matching margins with MSE.
The dense counterpart is MarginMSELoss.

Example

from datasets import Dataset

from sentence_transformers import MultiVectorEncoder, MultiVectorEncoderTrainer
from sentence_transformers.multi_vector_encoder.losses import MultiVectorMarginMSELoss

model = MultiVectorEncoder("answerdotai/ModernBERT-base")
# label is the teacher margin score(query, positive) - score(query, negative), so a
# positive value means the teacher ranked the positive above the negative.
train_dataset = Dataset.from_dict(
    {
        "query": ["What is the capital of France?", "Who painted the Mona Lisa?"],
        "positive": ["Paris is the capital of France.", "Leonardo da Vinci painted the Mona Lisa."],
        "negative": ["Berlin is the capital of Germany.", "Van Gogh painted The Starry Night."],
        "label": [3.5, 2.8],
    }
)
loss = MultiVectorMarginMSELoss(model)

trainer = MultiVectorEncoderTrainer(model=model, train_dataset=train_dataset, loss=loss)
trainer.train()

Initialize internal Module state, shared by both nn.Module and ScriptModule.