Add UNDIAL (#89)

* Fix hyperlinks in README (#2)

* testing commit

* Fixes

* cleanup

* Fixed DPO command

* download idk

* Revert "Dpo fix"

* download idk data

* fix dpo experiment config

* RMU (#6)

* IdkDPO script fix in tofu_unlearn.sh (#65)

* Fix hyperlinks in README
* Download I don't know data in setup_data.py
* Fix tofu_unlearn.sh for IdkDPO

---------

Co-authored-by: Anmol Mekala <[email protected]>

* overwrite=True

* RMU added

* Fix ref model device

* ruff fix

* RMU updated

* Update rmu.py

* Update README.md: add RMU

* Added references and renamed functions

---------

Co-authored-by: Anmol Mekala <[email protected]>

* Add structure to contributions, setup leaderboard, update documentation (#8)

* docs: updates, small corrections, re-formats

* modified ruff commands

* modified ruff commands

* CI/CD minor updates

* added contributing + leaderboard

* fix minor spelling misatkes

* docs: bunch of minor updates

* docs fixes

---------

Co-authored-by: molereddy <[email protected]>

* UNDIAL

* UNDIAL2

* UNDIAL3

* Ruff quality formatting changes

* fix config

* fix docs and script

* Update readme

---------

Co-authored-by: Anmol Mekala <[email protected]>
Co-authored-by: Dornavineeth <[email protected]>
Co-authored-by: Vineeth <[email protected]>
Co-authored-by: molereddy <[email protected]>

Author: 3 people

README.md

@@ -18,7 +18,7 @@
 
 ## 📖 Overview
 
-We provide efficient and streamlined implementations of the TOFU, MUSE and WMDP unlearning benchmarks while supporting 6 unlearning methods, 5+ datasets, 10+ evaluation metrics, and 7+ LLM architectures. Each of these can be easily extended to incorporate more variants.
+We provide efficient and streamlined implementations of the TOFU, MUSE and WMDP unlearning benchmarks while supporting 7 unlearning methods, 5+ datasets, 10+ evaluation metrics, and 7+ LLM architectures. Each of these can be easily extended to incorporate more variants.
 
 We invite the LLM unlearning community to collaborate by adding new benchmarks, unlearning methods, datasets and evaluation metrics here to expand OpenUnlearning's features, gain feedback from wider usage and drive progress in the field.
 
@@ -62,7 +62,7 @@ We provide several variants for each of the components in the unlearning pipelin
 | **Component**          | **Available Options** |
 |------------------------|----------------------|
 | **Benchmarks**        | [TOFU](https://arxiv.org/abs/2401.06121), [MUSE](https://muse-bench.github.io/), [WMDP](https://www.wmdp.ai/) |
-| **Unlearning Methods** | GradAscent, GradDiff, NPO, SimNPO, DPO, RMU |
+| **Unlearning Methods** | GradAscent, GradDiff, NPO, SimNPO, DPO, RMU, UNDIAL |
 | **Evaluation Metrics** | Verbatim Probability, Verbatim ROUGE, Knowledge QA-ROUGE, Model Utility, Forget Quality, TruthRatio, Extraction Strength, Exact Memorization, 6 MIA attacks, [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) |
 | **Datasets**          | MUSE-News (BBC), MUSE-Books (Harry Potter), TOFU (different splits), WMDP-Bio, WMDP-Cyber |
 | **Model Families**    | TOFU: LLaMA-3.2, LLaMA-3.1, LLaMA-2; MUSE: LLaMA-2; Additional: Phi-3.5, Phi-1.5, Gemma, Zephyr |

community/methods/UNDIAL/README.md

@@ -0,0 +1,24 @@
+# UNDIAL: Self-Distillation with Adjusted Logits for Robust Unlearning in Large Language Models (NAACL 2025)
+
+- Authors: Yijiang River Dong, Hongzhou Lin, Mikhail Belkin, Ramón Huerta, Ivan Vulić
+- Link​: https://arxiv.org/pdf/2402.10052
+
+# Setup
+- Hyperparameters: The original paper uses Llama-2 7B with LoRA to tune the model (rank=8, alpha=16) and learning rate of 1e-4. It's suggested to search the learning rate over [1e-5, 3e-4, 1e-4], and use an effective batch size of 32 (batch_size * gradient_accumulation). The other important hyperparemeter is beta, the strength of penalty, which typically takes a number between [3,10,30]. If we change to other models, adjusting learning rate accordingly.
+
+- Computation Setup: All experiments are run on one A100.
+- Other Details: The original paper does not use the retain set and aims to retain knowledge in all domains, not just on the retain set. So alpha is set to 0. Practionioners could search over the alpha or gamma to better retain the performance on the retain set.
+
+# Results
+Run `run.sh` script.
+
+# Citation
+@misc{dong2024undial,
+      title={UNDIAL: Self-Distillation with Adjusted Logits for Robust Unlearning in Large Language Models}, 
+      author={Yijiang River Dong and Hongzhou Lin and Mikhail Belkin and Ramon Huerta and Ivan Vulić},
+      year={2024},
+      eprint={2402.10052},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2402.10052}, 
+}

community/methods/UNDIAL/run.sh

@@ -0,0 +1,78 @@
+#!/bin/bash
+
+export MASTER_PORT=$(python -c "import socket; s=socket.socket(); s.bind(('', 0)); print(s.getsockname()[1]); s.close()")
+echo "Master Port: $MASTER_PORT"
+
+########################################################################################################################
+########################################### Unlearn TOFU models ########################################################
+########################################################################################################################
+
+models=(
+    "Llama-3.2-1B-Instruct"
+)
+trainers_experiments=(
+    "UNDIAL unlearn/tofu/default.yaml"
+)
+forget_retain_splits=(
+    "forget10 retain90"
+    "forget05 retain95"
+    "forget01 retain99"
+)
+
+per_device_train_batch_size=16
+gradient_accumulation_steps=2
+
+
+lrs=(1e-5 1e-4 3e-4)
+alphas=(1 2 5)
+betas=(3 10 30)
+
+
+for split in "${forget_retain_splits[@]}"; do
+    forget_split=$(echo $split | cut -d' ' -f1)
+    retain_split=$(echo $split | cut -d' ' -f2)
+    for model in "${models[@]}"; do
+        for trainer_experiment in "${trainers_experiments[@]}"; do
+            trainer=$(echo $trainer_experiment | cut -d' ' -f1)
+            experiment=$(echo $trainer_experiment | cut -d' ' -f2)
+            for lr in "${lrs[@]}"; do
+                for beta in "${betas[@]}"; do 
+                    for alpha in "${alphas[@]}"; do          
+                        task_name=tofu_${model}_${forget_split}_${trainer}_lr${lr}_beta${beta}_alpha${alpha}
+                        model_path=open-unlearning/tofu_${model}_full
+                        echo ${task_name}: Unlearning ${model_path} using ${trainer}
+
+                        # Unlearn
+                        CUDA_VISIBLE_DEVICES=0 \
+                        python src/train.py --config-name=unlearn.yaml \
+                        experiment=${experiment} \
+                        trainer=${trainer} \
+                        task_name=${task_name} \
+                        model=${model} \
+                        forget_split=${forget_split} \
+                        retain_split=${retain_split} \
+                        model.model_args.pretrained_model_name_or_path=${model_path} \
+                        retain_logs_path=saves/eval/tofu_${model}_${retain_split}/TOFU_EVAL.json \
+                        trainer.args.per_device_train_batch_size=$per_device_train_batch_size \
+                        trainer.args.gradient_accumulation_steps=$gradient_accumulation_steps \
+                        trainer.args.eval_strategy=no \
+                        trainer.args.eval_on_start=False \
+                        trainer.args.learning_rate=$lr \
+                        trainer.method_args.beta=$beta \
+                        trainer.method_args.alpha=$alpha
+
+                        # Eval
+                        CUDA_VISIBLE_DEVICES=0 python src/eval.py \
+                        experiment=eval/tofu/default.yaml \
+                        forget_split=${forget_split} \
+                        model=${model} \
+                        task_name=${task_name} \
+                        model.model_args.pretrained_model_name_or_path=saves/unlearn/${task_name} \
+                        paths.output_dir=saves/unlearn/${task_name}/evals \
+                        retain_logs_path=saves/eval/tofu_${model}_${retain_split}/TOFU_EVAL.json
+                    done
+                done
+            done
+        done
+    done
+done

configs/trainer/UNDIAL.yaml

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

docs/links.md

@@ -25,6 +25,7 @@ Links to research papers and resources corresponding to implemented features in
 | SimNPO             |  Paper [📄](https://arxiv.org/abs/2410.07163), Code [🐙](https://github.com/OPTML-Group/Unlearn-Simple) |
 | IdkDPO             | TOFU ([📄](https://arxiv.org/abs/2401.06121)) |
 | RMU             | WMDP paper ([🐙](https://github.com/centerforaisafety/wmdp/tree/main/rmu), [🌐](https://www.wmdp.ai/)), later used in G-effect ([🐙](https://github.com/tmlr-group/G-effect/blob/main/dataloader.py)) |
+| UNDIAL             | Paper [📄](https://arxiv.org/pdf/2402.10052), Code [🐙](https://github.com/dong-river/LLM_unlearning/tree/main) |
 
 ---
 

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
 
 import logging
 
@@ -88,3 +89,4 @@ def load_trainer(
 _register_trainer(DPO)
 _register_trainer(SimNPO)
 _register_trainer(RMU)
+_register_trainer(UNDIAL)

src/trainer/unlearn/undial.py

@@ -0,0 +1,27 @@
+from trainer.utils import compute_undial_loss
+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 = compute_undial_loss(
+            model, self.ref_model, forget_inputs, self.beta
+        )
+
+        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

src/trainer/utils.py

@@ -66,3 +66,35 @@ def compute_dpo_loss(model, ref_model, win_inputs=None, lose_inputs=None, beta=1
 
     loss = -2 / beta * F.logsigmoid(beta * (win_log_ratio - lose_log_ratio)).mean()
     return loss, (win_outputs, lose_outputs)
+
+
+def compute_undial_loss(model, ref_model, inputs, beta):
+    # 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 = 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 * beta
+    soft_label = F.softmax(pre_softmax, dim=-1)
+
+    loss_fct = nn.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