feat: Add option to use MultipleNegativesRankingLoss for EmbeddingRetriever training with sentence-transformers

[bglearning]

Add option to use MultipleNegativesRankingLoss for EmbeddingRetriever training with sentence-transformers.

Related Issues

• Resolves Add MultipleNegativesRankingLoss for EmbeddingRetriever training with sentence-transformers #3136

Proposed Changes:

Expose an option train_loss on EmbeddingRetriever.train. It is propagated to _SentenceTransformersEmbeddingEncoder.train and determines the loss that's used. Also influences the training data format that can be used/is expected.

So the API is:

embedding_retriever = EmbeddingRetriever(...)
training_data = [{"question": ..., "pos_doc": ...}, ...]
# Loss can be specified when calling `train`. Options currently: 'mnrl' or 'margin_mse'
embedding_retriever.train(training_data=training_data, train_loss='mnrl') 

How did you test it?

A Colab notebook with a sample run. Qualitative comparison against the baseline model.

Notes for the reviewer

Some possible aspects up for discussion

• I exposed the loss option as a string.
  • Another option could be to directly accept sentence-transformers losses (might even be able to say can use any loss in there as long as the data format is correct but would be a big jump). Plus might tie us to keep doing so (which may or may not be a problem). With the string we have the flexibility to later internally swap things out.
  • We could also construct our own enums or probably better yet loss classes which are thin wrappers around the sentence-transformers losses. But probably not needed at this stage.
• Slightly tricky to provide users with list of possible losses. I guess best place is in the EmbeddingRetriever docstring, which I went with.
• Any unit or integration tests needed? Any other experiments needed?

Checklist

• I have read the contributors guidelines and the code of conduct
• I have updated the related issue with new insights and changes
• I added tests that demonstrate the correct behavior of the change
• I've used the conventional commit convention for my PR title
• I documented my code
• I ran pre-commit hooks and fixed any issue

[bogdankostic]

LGTM

[vblagoje]

Hey @bglearning, I just realized that mnrl is now the default loss for GPL rather than margin_mse! Why make a change if the original paper used margin_mse?

[bglearning]

Hi @vblagoje,

The default change on this EmbeddingRetriever is primarily motivated to support the more straightforward use-case of directly training the retriever without the pseudo-labeling (through GPL) or some other intermediate process to come up with the scores.

So line of reasoning:

• User has some data (just query and pos-doc) and they want to start training directly (default case). mnrl (the "simpler" loss, margin_mse needs neg-doc and scores)
• User also wants to do GPL (or some other method) to come up with the pseudo-labels before training the model. User sets the loss to margin_mse.

But ya, realized I missed updating the GPL tutorial to explicitly set the loss to margin_mse which is definitely an issue.
And also generally should have highlighted the change in default as a big change as it's an API change.

We could update the GPL tutorial or update the default. Either way seems okay to me.