Support for using MultipleNegativesRankingLoss with SentenceTransformerTrainer (evaluation, early stopping, eval logging, custom callbacks)

[pinareceaktan]

Hi,

Currently it seems not possible to use MultipleNegativesRankingLoss together with the SentenceTransformerTrainer.

This is the training script i follow (i implemented the same steps for my real world dataset)
https://github.com/UKPLab/sentence-transformers/blob/master/examples/sentence_transformer/training/ms_marco/train_bi-encoder_mnrl.py

And this is an issue because,
When I train with MultipleNegativesRankingLoss using traditional model.fit, the trainer does not return an evaluation loss. I know there is a way to do it but its either too frequent, or reduces the batch size therefore training results in overfit. Here, let me explain:

• In model.fit(), evaluation_steps is defined in terms of training steps (iterations). Since each step = one batch, changing the batch size directly changes the effective evaluation frequency.
• With larger batch sizes (which are recommended for MNRL to stabilize gradients and reduce overfitting), the number of steps per epoch decreases. This means that for the same evaluation_steps value, evaluations become less frequent.
• With smaller batch sizes, you get more frequent evaluations, but the signal from MNRL becomes noisier and training may overfit or become unstable.

So the trade-off is not only about choosing between frequent vs. infrequent evaluations, but also about balancing it with the batch size that MNRL requires.

This makes it very difficult to tune experiments reliably:

If I prioritize large batches (good for MNRL), I lose eval signal.

If I prioritize eval frequency, I am forced into small batches that hurt training quality and efficiency.

A native solution in SentenceTransformerTrainer that decouples evaluation frequency from batch size — and integrates eval loss logging with callbacks — would make MNRL training much more practical and reproducible.

And even more disturbing fact is that; with model.fit, I cannot properly use early stopping callbacks or track metrics through (preferably MLflow) any other logging frameworks.

I am just training blindly without any feedback signal from the validation set, which is very limiting for real-world experiments.

It would be very helpful if there was:

A way to integrate MultipleNegativesRankingLoss into the SentenceTransformerTrainer’s evaluation loop.

Native support for logging the eval loss so that callbacks (early stopping, best checkpoint, etc.) and external loggers (e.g. MLflow, WandB) can be used.

This gap makes it hard to run reproducible and well-controlled experiments — right now the only option is to “just train and hope” with MNRL, which is a bit heartbreaking 🙂

Would it make sense to add an example or official support for this in the examples/ directory?
I would be happy to contribute an implementation (custom collator + compute_loss integration for HF Trainer) if maintainers think this fits.

Thanks a lot!

[tomaarsen]

Hello!

Thanks for opening this issue. As a bit of extra context, though you seem to already know this, the model-fit-style training is being softly deprecated in favor of using the SentenceTransformerTrainer directly since the latter was introduced in the v3.0 release. There's a handful of features only (conveniently) possible by using the SentenceTransformerTrainer directly.

About your suggestions:

A way to integrate MultipleNegativesRankingLoss into the SentenceTransformerTrainer’s evaluation loop.

I do understand the underlying idea that higher batch sizes at the same eval_steps decreases how often evaluation is ran. However, it should be possible to update eval_steps inversely to the batch size, right? E.g. per_device_train_batch_size=64 & eval_steps=100 -> per_device_train_batch_size=256 & eval_steps=25.
This only becomes a bottleneck when you want your eval_steps to be less than 1, i.e. using a batch size of 2048 and wanting to evaluate for every 1000 samples processed. However, it's not feasible to evaluate halfway through processing a training batch. The model weights are updated once per batch, after the batch has been processed. This means that it's not possible/useful to evaluate multiple times per batch, i.e. the lowest reasonable eval_steps is 1.

I'm not sure what behaviour would then be preferred that isn't already supported.

Native support for logging the eval loss so that callbacks (early stopping, best checkpoint, etc.) and external loggers (e.g. MLflow, WandB) can be used.

There's a bit more context in the Training Overview > Evaluator section: The SentenceTransformerTrainer supports 2 forms of evaluation:

1. An evaluation loss with eval_dataset and the loss function
2. Evaluation metrics with an evaluator

Both of these are logged during training, regardless of whether you're using SentenceTransformerTrainer or model.fit, and the former also supports additional callbacks (Training Overview > Trainer > Callbacks) for e.g. W&B, MLFlow, or e.g. early stopping support. It's already compatible with the HF transformers callbacks, because the SentenceTransformerTrainer (like the other trainers in Sentence Transformers) are already a subclass of the HF transformers Trainer.

In conclusion, I believe the functionality that you're proposing should already be possible with the SentenceTransformerTrainer.

• Tom Aarsen

[pinareceaktan]

Hi Tom, thanks a lot for the detailed explanation.

I opened this PR to illustrate what I was aiming for in the issue.
The main motivation is to make MultipleNegativesRankingLoss easier to use with SentenceTransformerTrainer by:

• Allowing evaluation to run reliably independent of batch size,
• Exposing the eval loss so callbacks (early stopping, best checkpoint, etc.) can be used, and
• Supporting logging through the standard HF Trainer interface.

Another important point: datasets like MS MARCO are naturally structured around queries with multiple negatives. This format doesn’t align with the flat datasets.Dataset structure that the HF Trainer expects. Out of the box, it wasn’t really possible to plug MS MARCO into SentenceTransformerTrainer. The PR adds a collator and uses no duplicates sampler(maybe a sampler for this type of task implemented in the future) during training that can handle this case cleanly, so training with real-world retrieval datasets becomes practical.

Experiments
I tested my script on MS MARCO using distilbert-base-uncased as the base encoder.

Training Setup (common across both runs).

Base model: distilbert-base-uncased
Batch size: 128 (per device, A100 GPU)
Trainer and Loss: via SentenceTransformerTrainer with MultipleNegativesRankingLoss
Gradient accumulation: 4 (effective batch size ×4)
FP16: enabled
Early stopping: patience=3, threshold=5e-5
Trainer settings:
logging_steps=100, eval_steps=100, save_steps=100
remove_unused_columns=False
eval_accumulation_steps=10, max_grad_norm=1.0
load_best_model_at_end=True
metric_for_best_model="eval_ms_marco_cosine_accuracy@5"

Run 1: With hard negatives
Collator: Custom FlatMNRLCollator with explicit negatives
Sampler: Batch sampler to distribute multiple negatives per query
Result: Best eval metric (MS MARCO Cosine Accuracy@5) around 0.97 (≈0.970476 at step 1600).
Other metrics: Accuracy@1 ≈ 0.86, Accuracy@10 ≈ 0.98, MAP@100 ≈ 0.9131.
Validation loss converged to ~0.044 (it was 0.075 initialy).

Run 2: Without hard negatives (in-batch negatives only)
Collator: Custom FlatMNRLCollator without negatives
Sampler: i used no duplicates sampler, but its not vital as it was in the 'with negatives training'
Result: Best eval metric (MS MARCO Cosine Accuracy@5) is 0.9711(at step 2100).
Other metrics: Accuracy@1 ≈ 0.85–0.86, Accuracy@10 ≈ 0.98.
Validation loss converged to 0.032 (it was 0.068 initialy).

Would love your feedback on whether this approach makes sense or if there’s a better direction to align with the library’s roadmap.

Thanks again!