Add new unlearning method UNDIAL

Author: dong-river

configs/trainer/UNDIAL.yaml

@@ -0,0 +1,12 @@
+handler: UNDIAL # corresponds to the class defined in src/trainer/unlearn/grad_diff.py
+args: # HuggingFace TrainingArguments
+  per_device_train_batch_size: 2
+  per_device_eval_batch_size: 16
+  gradient_accumulation_steps: 4
+  learning_rate: 1e-5
+  num_train_epochs: 10
+method_args: # Your own method-specific arguments
+  gamma: 1.0
+  alpha: 1.0
+  beta: 10.0 # the strength of penalty for memorized tokens
+  retain_loss_type: NLL

src/trainer/__init__.py

@@ -10,6 +10,7 @@
 from trainer.unlearn.dpo import DPO
 from trainer.unlearn.simnpo import SimNPO
 from trainer.unlearn.rmu import RMU
+from trainer.unlearn.undial import UNDIAL
 
 TRAINER_REGISTRY: Dict[str, Any] = {}
 
@@ -20,6 +21,7 @@ def _register_trainer(trainer_class):
 
 def load_trainer_args(trainer_args: DictConfig, dataset):
     trainer_args = dict(trainer_args)
+    trainer_args["output_dir"] = trainer_args.pop("output_dir", "./output")
     warmup_epochs = trainer_args.pop("warmup_epochs", None)
     if warmup_epochs:
         batch_size = trainer_args["per_device_train_batch_size"]
@@ -81,3 +83,4 @@ def load_trainer(
 _register_trainer(DPO)
 _register_trainer(SimNPO)
 _register_trainer(RMU)
+_register_trainer(UNDIAL)

src/trainer/unlearn/undial.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn.functional as F
+from torch.nn import CrossEntropyLoss
+from trainer.unlearn.grad_diff import GradDiff
+
+class UNDIAL(GradDiff):
+    def __init__(self, beta=1.0, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.beta = beta
+        if self.ref_model is None:
+            self.ref_model = self._prepare_ref_model(self.model)
+
+    def compute_loss(self, model, inputs, return_outputs=False):
+        forget_inputs = inputs["forget"]
+        forget_loss, forget_outputs = self.compute_undial_loss(model, forget_inputs)
+
+        retain_inputs = inputs["retain"]
+        retain_inputs = {
+            "input_ids": retain_inputs["input_ids"],
+            "attention_mask": retain_inputs["attention_mask"],
+            "labels": retain_inputs["labels"],
+        }
+        retain_loss = self.compute_retain_loss(model=model, retain_inputs=retain_inputs)
+
+        loss = self.gamma * forget_loss + self.alpha * retain_loss
+        return (loss, forget_outputs) if return_outputs else loss
+    
+    def compute_undial_loss(self, model, inputs):
+        # Forward pass on the student (trainable) model
+        outputs = model(**inputs)
+        logits = outputs.logits
+        labels = inputs["labels"]
+        
+        shift_labels = labels[..., 1:].contiguous()
+        shift_logits = logits[..., :-1, :].contiguous()
+        
+        # Forward pass on the teacher model (no grad)
+        with torch.no_grad():
+            teacher_logits = self.ref_model(**inputs).logits
+        shift_teacher_logits = teacher_logits[..., :-1, :].contiguous()
+        
+        # Build the mask that identifies the tokens need to be unlearned
+        mask = torch.zeros_like(shift_teacher_logits)
+        batch_idx = torch.arange(mask.shape[0]).view(-1, 1, 1)
+        seq_idx = torch.arange(mask.shape[1]).view(1, -1, 1)
+        mask[batch_idx, seq_idx, shift_labels.unsqueeze(-1)] = 1.0
+        
+        # Adjust teacher logits: subtract di_strength on the correct token
+        pre_softmax = shift_teacher_logits - mask * self.beta
+        soft_label = F.softmax(pre_softmax, dim=-1)
+
+        loss_fct = CrossEntropyLoss(reduction='none')
+        loss = loss_fct(
+            shift_logits.view(-1, shift_logits.size(-1)),
+            soft_label.view(-1, soft_label.size(-1)),
+        )
+        return loss.mean(), outputs