Repository: deepset-ai/haystack

Issue:
Add MultipleNegativesRankingLoss for EmbeddingRetriever training with sentence-transformers
Currently training `EmbeddingRetriever` with sentence-transformer uses the `MarginMSE` loss. We want to also support `MultipleNegativesRankingLoss` which can work with a simpler data requirement, particularly it doesn't necessitate a "score" for the data pairs/tuples.

Related discussion: deepset-ai/haystack-tutorials#35 

## Background and Related work

deepset-ai/haystack#2388 added the ability to train `EmbeddingRetriever` (sentence-transformer variant) with Generative Pseudo Labeling (GPL). It uses `MarginMSE` loss with labels coming from a soft pseudo-labeling process. [Example Colab notebook](https://colab.research.google.com/drive/1Tz9GSzre7JfvXDDKe7sCnO0FMuDViMnN?usp=sharing).

Input data is then of the format:

```
[
{”question”: …, “pos_doc”: …, “neg_doc”: …, “score”: …}, 
... 
]
```

It works well. However, there can be cases where users want to directly move onto the retriever training from their data without the pseudo-labeling or some other intermediate process to come up with the scores. 

Supporting the use of [MultipleNegativesRankingLoss](https://www.sbert.net/docs/package_reference/losses.html#multiplenegativesrankingloss)(MNRL) can provide such an option.

## Proposal

Provide an argument (maybe a string or directly the loss class from `sentence-transformers`) to the `EmbeddingRetriever#train` method to allow for selection between the two losses. Data checks could be added to make sure the loss choice and data formats are compatible. 

For MNRL, the format would be as below (with the `neg_doc` also being optional).

```
[
{”question”: …, “pos_doc”: …, “neg_doc”: …}, 
... 
]
```

## Next up

I'll open a draft PR and start working on this. Fine implementation details can be worked out there. In the meantime, if there are any thoughts, please drop them here. 
cc: @mkkuemmel @mathislucka @vblagoje 


Discussion:

Resolving patch:
diff --git a/haystack/nodes/retriever/_embedding_encoder.py b/haystack/nodes/retriever/_embedding_encoder.py
--- a/haystack/nodes/retriever/_embedding_encoder.py
+++ b/haystack/nodes/retriever/_embedding_encoder.py
@@ -5,7 +5,7 @@
 
 import numpy as np
 import torch
-from sentence_transformers import InputExample, losses
+from sentence_transformers import InputExample
 from torch.utils.data import DataLoader
 from torch.utils.data.sampler import SequentialSampler
 from tqdm.auto import tqdm
@@ -14,6 +14,7 @@
 from haystack.modeling.data_handler.dataloader import NamedDataLoader
 from haystack.modeling.data_handler.dataset import convert_features_to_dataset, flatten_rename
 from haystack.modeling.infer import Inferencer
+from haystack.nodes.retriever._losses import _TRAINING_LOSSES
 from haystack.schema import Document
 
 if TYPE_CHECKING:
@@ -195,14 +196,34 @@ def train(
         n_epochs: int = 1,
         num_warmup_steps: int = None,
         batch_size: int = 16,
+        train_loss: str = "mnrl",
     ):
 
-        train_examples = [
-            InputExample(texts=[i["question"], i["pos_doc"], i["neg_doc"]], label=i["score"]) for i in training_data
-        ]
-        logger.info(f"GPL training/adapting {self.embedding_model} with {len(train_examples)} examples")
+        if train_loss not in _TRAINING_LOSSES:
+            raise ValueError(f"Unrecognized train_loss {train_loss}. Should be one of: {_TRAINING_LOSSES.keys()}")
+
+        st_loss = _TRAINING_LOSSES[train_loss]
+
+        train_examples = []
+        for train_i in training_data:
+            missing_attrs = st_loss.required_attrs.difference(set(train_i.keys()))
+            if len(missing_attrs) > 0:
+                raise ValueError(
+                    f"Some training examples don't contain the fields {missing_attrs} which are necessary when using the '{train_loss}' loss."
+                )
+
+            texts = [train_i["question"], train_i["pos_doc"]]
+            if "neg_doc" in train_i:
+                texts.append(train_i["neg_doc"])
+
+            if "score" in train_i:
+                train_examples.append(InputExample(texts=texts, label=train_i["score"]))
+            else:
+                train_examples.append(InputExample(texts=texts))
+
+        logger.info(f"Training/adapting {self.embedding_model} with {len(train_examples)} examples")
         train_dataloader = DataLoader(train_examples, batch_size=batch_size, drop_last=True, shuffle=True)
-        train_loss = losses.MarginMSELoss(self.embedding_model)
+        train_loss = st_loss.loss(self.embedding_model)
 
         # Tune the model
         self.embedding_model.fit(
diff --git a/haystack/nodes/retriever/_losses.py b/haystack/nodes/retriever/_losses.py
new file mode 100644
--- /dev/null
+++ b/haystack/nodes/retriever/_losses.py
@@ -0,0 +1,12 @@
+from collections import namedtuple
+from typing import Dict
+
+from sentence_transformers import losses
+
+
+SentenceTransformerLoss = namedtuple("SentenceTransformerLoss", "loss required_attrs")
+
+_TRAINING_LOSSES: Dict[str, SentenceTransformerLoss] = {
+    "mnrl": SentenceTransformerLoss(losses.MultipleNegativesRankingLoss, {"question", "pos_doc"}),
+    "margin_mse": SentenceTransformerLoss(losses.MarginMSELoss, {"question", "pos_doc", "neg_doc", "score"}),
+}
diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py
--- a/haystack/nodes/retriever/dense.py
+++ b/haystack/nodes/retriever/dense.py
@@ -1863,6 +1863,7 @@ def train(
         n_epochs: int = 1,
         num_warmup_steps: int = None,
         batch_size: int = 16,
+        train_loss: str = "mnrl",
     ) -> None:
         """
         Trains/adapts the underlying embedding model.
@@ -1885,6 +1886,10 @@ def train(
         :type num_warmup_steps: int
         :param batch_size: The batch size to use for the training, defaults to 16
         :type batch_size: int (optional)
+        :param train_loss: The loss to use for training.
+                           If you're using sentence-transformers as embedding_model (which are the only ones that currently support training),
+                           possible values are 'mnrl' (Multiple Negatives Ranking Loss) or 'margin_mse' (MarginMSE).
+        :type train_loss: str (optional)
         """
         self.embedding_encoder.train(
             training_data,
@@ -1892,6 +1897,7 @@ def train(
             n_epochs=n_epochs,
             num_warmup_steps=num_warmup_steps,
             batch_size=batch_size,
+            train_loss=train_loss,
         )
 
     def save(self, save_dir: Union[Path, str]) -> None:


Test patch: