Repository: huggingface/trl

Issue:
Add margin to reward trainer, similar to LLAMA-2
In the LLAMA-2 paper, an additional margin term is added to the loss function of the reward model. This allows to add a higher weight to pairs with a higher score difference. 

![image](https://github.com/huggingface/trl/assets/9884254/6c74c655-d0a5-4bd6-a82d-fa05db9c5c3f)

This should also be possible with the TRL RewardTrainer. Adding that functionality probably requires to rebuild the trainer and collator. 

Does this make sense? 


Discussion:

Resolving patch:
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -47,6 +47,9 @@ class RewardTrainer(Trainer):
     - `input_ids_rejected`
     - `attention_mask_rejected`
 
+    Optionally, you can also pass a `margin` entry to the dataset. This entry should contain the margin used to modulate the
+    loss of the reward model as outlined in https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/.
+    If you don't pass a margin, no margin will be used.
     """
 
     def __init__(
@@ -209,7 +212,12 @@ def compute_loss(
             input_ids=inputs["input_ids_rejected"],
             attention_mask=inputs["attention_mask_rejected"],
         )[0]
-        loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+        # calculate loss, optionally modulate with margin
+        if "margin" in inputs:
+            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean()
+        else:
+            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
+
         if return_outputs:
             return loss, {
                 "rewards_chosen": rewards_chosen,
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -195,6 +195,9 @@ class RewardDataCollatorWithPadding:
     def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
         features_chosen = []
         features_rejected = []
+        margin = []
+        # check if we have a margin. If we do, we need to batch it as well
+        has_margin = "margin" in features[0]
         for feature in features:
             # check if the keys are named as expected
             if (
@@ -219,6 +222,8 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
                     "attention_mask": feature["attention_mask_rejected"],
                 }
             )
+            if has_margin:
+                margin.append(feature["margin"])
         batch_chosen = self.tokenizer.pad(
             features_chosen,
             padding=self.padding,
@@ -240,6 +245,9 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
             "attention_mask_rejected": batch_rejected["attention_mask"],
             "return_loss": True,
         }
+        if has_margin:
+            margin = torch.tensor(margin, dtype=torch.float)
+            batch["margin"] = margin
         return batch
 
 


Test patch:
diff --git a/tests/test_reward_trainer.py b/tests/test_reward_trainer.py
--- a/tests/test_reward_trainer.py
+++ b/tests/test_reward_trainer.py
@@ -260,3 +260,55 @@ def test_reward_trainer_assert_value_error(self):
                     tokenizer=self.tokenizer,
                     train_dataset=dummy_dataset,
                 )
+
+    def test_reward_trainer_margin(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = RewardConfig(
+                output_dir=tmp_dir,
+                per_device_train_batch_size=2,
+                max_steps=3,
+                remove_unused_columns=False,
+                gradient_accumulation_steps=4,
+                learning_rate=9e-1,
+                evaluation_strategy="steps",
+            )
+
+            # fmt: off
+            dummy_dataset_dict = {
+                "input_ids_chosen": [
+                    torch.LongTensor([0, 1, 2,]),
+                ],
+                "attention_mask_chosen": [
+                    torch.LongTensor([1, 1, 1]),
+                ],
+                "input_ids_rejected": [
+                    torch.LongTensor([0, 2,]),
+                ],
+                "attention_mask_rejected": [
+                    torch.LongTensor([1, 1]),
+                ],
+                "margin": [
+                    torch.FloatTensor([1.0]),
+                ]
+            }
+            # fmt: on
+            dummy_dataset = Dataset.from_dict(dummy_dataset_dict)
+
+            trainer = RewardTrainer(
+                model=self.model,
+                args=training_args,
+                tokenizer=self.tokenizer,
+                train_dataset=dummy_dataset,
+                eval_dataset=dummy_dataset,
+            )
+
+            batch = [dummy_dataset[0]]
+            batch = trainer.data_collator(batch)
+            loss, outputs = trainer.compute_loss(trainer.model, batch, return_outputs=True)
+
+            self.assertAlmostEqual(
+                loss,
+                -torch.nn.functional.logsigmoid(
+                    outputs["rewards_chosen"] - outputs["rewards_rejected"] - batch["margin"]
+                ).mean(),
+            )