Add BalDRO Unlearning Methods

README.md

@@ -19,7 +19,7 @@
 
 ## 📖 Overview
 
-We provide efficient and streamlined implementations of the TOFU, MUSE and WMDP unlearning benchmarks while supporting 14+ 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 15+ 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.
 
@@ -84,7 +84,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/), [DUET](https://arxiv.org/abs/2602.19612) |
-| **Unlearning Methods** | GradAscent, GradDiff, NPO, SimNPO, DPO, RMU, UNDIAL, AltPO, SatImp, WGA, CE-U, PDU, AdaPop, T3 |
+| **Unlearning Methods** | GradAscent, GradDiff, NPO, SimNPO, DPO, RMU, UNDIAL, AltPO, SatImp, WGA, CE-U, PDU, AdaPop, T3, BalDRO |
 | **Evaluation Metrics** | Verbatim Probability, Verbatim ROUGE, Knowledge QA-ROUGE, Model Utility, Forget Quality, TruthRatio, Extraction Strength, Exact Memorization, 6 MIA attacks, [Unlearning Depth Score](https://arxiv.org/abs/2605.24614), [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, [DUET](https://huggingface.co/datasets/SwetieePawsss/DUET) (rare / popular Wikidata facts) |
 | **Model Families**    | TOFU: Llama-3.2, Llama-3.1, Llama-2; MUSE: Llama-2; DUET: Llama-3.1, Qwen2.5, Gemma; Additional: Phi-3.5, Phi-1.5, Gemma, Zephyr |

community/methods/BalDRO/README.md

@@ -0,0 +1,71 @@
+# BalDRO: A Distributionally Robust Optimization Based Framework for Large Language Model Unlearning
+
+We propose [BalDRO](https://arxiv.org/abs/2601.09172), a distributionally robust optimization framework that adaptively emphasizes hard-to-unlearn samples. BalDRO-G selects hard groups, while BalDRO-DV smoothly reweights individual losses through the Donsker-Varadhan dual. Both variants support NPO, SimNPO, and SatImp.
+
+## Setup
+
+Use the standard Open-Unlearning environment and prepare the TOFU or MUSE benchmark; no additional dependencies are required. The training objective is
+
+$$
+\mathcal{L}(\theta)=\gamma\mathcal{L}_{\mathrm{forget}}(\theta)+\alpha\mathcal{L}_{\mathrm{retain}}(\theta),
+$$
+
+where `gamma` and `alpha` control the forget-retain trade-off.
+
+**BalDRO-G** selects the $k$ forget samples with the lowest length-normalized negative log-likelihood $\bar{\ell}_{\mathrm{NLL}}$, then applies the base unlearning objective to this group:
+
+$$
+G=\operatorname{TopK}_{i\in B_f}\!\left(-\bar{\ell}_{\mathrm{NLL}}(x_i;\theta),k\right),\qquad
+\mathcal{L}_{\mathrm{forget}}^{\mathrm{G}}=\frac{1}{k}\sum_{i\in G}\ell_i(\theta).
+$$
+
+Here $\ell_i$ is the base method's per-sample forget loss and $k=\max(1,\lfloor r|B_f|\rfloor)$. The selection ratio $r$ is `sampling_ratio` (default `0.5`, range $(0,1]$); a random retain group of up to $k$ samples supplies the retain loss.
+
+**BalDRO-DV** replaces the mean forget loss with
+
+$$
+\mathcal{L}_{\mathrm{forget}}^{\mathrm{DV}}=\tau_f\log\left(\frac{1}{|B_f|}\sum_{i\in B_f}\exp\left(\frac{\ell_i(\theta)}{\tau_f}\right)\right).
+$$
+
+The positive temperature $\tau_f$ is `beta_dv_forget` (default `1.0`); smaller values emphasize larger losses more strongly. Defaults are `forget_dro=true` and `retain_dro=false`. Optional retain-side DV uses `beta_dv_retain` (default `1.0`) and supports NLL retention only. Base-method parameters remain `beta` for NPO, `beta` and `delta` for SimNPO, and `beta1` and `beta2` for SatImp.
+
+## Results
+
+The following scripts reproduce the TOFU and MUSE experiment grids. Run from the repository root after preparing the benchmark's retained-model evaluation logs. Scripts use GPU `0` and W&B; authenticate with W&B or set `REPORTTO=none` in the selected script.
+
+### BalDRO-G
+
+| Base method | Trainer | TOFU | MUSE |
+| --- | --- | --- | --- |
+| NPO | `BalDRO/GroupNPO` | [Script](scripts/tofu/train_tofu_groupnpo.sh) | [Script](scripts/muse/train_muse_npo_group.sh) |
+| SimNPO | `BalDRO/GroupSimNPO` | [Script](scripts/tofu/train_tofu_groupsimnpo.sh) | [Script](scripts/muse/train_muse_simnpo_group.sh) |
+| SatImp | `BalDRO/GroupSatImp` | [Script](scripts/tofu/train_tofu_groupsatimp.sh) | [Script](scripts/muse/train_muse_satimp_group.sh) |
+
+### BalDRO-DV
+
+| Base method | Trainer | TOFU | MUSE |
+| --- | --- | --- | --- |
+| NPO | `BalDRO/DrNPO` | [Script](scripts/tofu/train_tofu_drnpo.sh) | [Script](scripts/muse/train_muse_drnpo.sh) |
+| SimNPO | `BalDRO/DrSimNPO` | [Script](scripts/tofu/train_tofu_drsimnpo.sh) | [Script](scripts/muse/train_muse_drsimnpo.sh) |
+| SatImp | `BalDRO/DrSatImp` | [Script](scripts/tofu/train_tofu_drsatimp.sh) | [Script](scripts/muse/train_muse_drsatimp.sh) |
+
+For example:
+
+```bash
+bash community/methods/BalDRO/scripts/tofu/train_tofu_drnpo.sh
+bash community/methods/BalDRO/scripts/muse/train_muse_simnpo_group.sh
+```
+
+Each script runs its full parameter grid. Compared with the source scripts, trainer names use the `BalDRO/` prefix, `run_name` uses Hydra's `+` syntax, and `do_save=True` is omitted because Open-Unlearning saves the final model automatically.
+
+## Citation
+
+```bibtex
+@inproceedings{shao2026baldro,
+  title={Baldro: A distributionally robust optimization based framework for large language model unlearning},
+  author={Shao, Pengyang and Zhai, Naixin and Chen, Lei and Yang, Yonghui and Zhu, Fengbin and Yang, Xun and Wang, Meng},
+  booktitle={Proceedings of the ACM Web Conference 2026},
+  pages={8874--8884},
+  year={2026}
+}
+```

community/methods/BalDRO/run.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Run from the Open-Unlearning repository root.
+if [ "$#" -eq 0 ]; then
+    echo "Usage: $0 <tofu|muse>/<script_name>.sh [script arguments]" >&2
+    exit 2
+fi
+
+script="$1"
+shift
+exec bash "$(dirname "${BASH_SOURCE[0]}")/scripts/${script}" "$@"

community/methods/BalDRO/scripts/muse/train_muse_drnpo.sh

@@ -0,0 +1,77 @@
+#!/usr/bin/env bash
+
+DATE=$(date "+%m%d")
+TIME=$(date "+%H%M%S")
+
+
+export CUDA_VISIBLE_DEVICES=0
+
+REPORTTO="wandb"
+WANDB_PROJECT="BalDRO"
+
+MODEL="Llama-2-7b-hf"
+TRAINER="DrNPO"
+splits=(
+    "News"
+    "Books"
+)
+
+# lr, batchsize, grad_acc, epochs
+lr_set=("1e-5" "2e-5" "3e-5" "4e-5" "5e-5")
+bz_set=("4 4" "4 8")
+epoch_set=(10)
+beta_dv_set=(0.5 1.0 2.0 5.0 10.0)
+
+
+for split in "${splits[@]}"; do
+    for lr in "${lr_set[@]}"; do
+        for bz in "${bz_set[@]}"; do
+            for epochs in "${epoch_set[@]}"; do
+                for beta_dv_forget in "${beta_dv_set[@]}"; do
+                    # Args ========================================
+                    PRETRAINED_PATH="muse-bench/MUSE-${split}_target"
+                    TOKENIZER_PRETRAINED="meta-llama/Llama-2-7b-chat-hf"
+
+                    bsz=$(echo $bz | cut -d' ' -f1)
+                    grad_acc=$(echo $bz | cut -d' ' -f2)
+
+                    # learning_rate, batchsize, grad_acc, epochs
+                    SUFFIX="lr${lr}_b${bsz}_ga${grad_acc}_e${epochs}_betaDV${beta_dv_forget}_day${DATE}_time${TIME}"
+                    TASK_NAME="unlearn_muse_${split}_${MODEL}_${TRAINER}_${SUFFIX}"
+                    OUTPUT_DIR="./saves/unlearn/muse/${split}/${MODEL}/${TRAINER}/${SUFFIX}"
+
+                    # TRAIN COMMAND =================================
+                    export WANDB_PROJECT=${WANDB_PROJECT}
+                    python src/train.py --config-name=unlearn.yaml \
+                        experiment=unlearn/muse/default.yaml  \
+                        trainer=BalDRO/${TRAINER} \
+                        model=${MODEL} \
+                        model.model_args.pretrained_model_name_or_path=${PRETRAINED_PATH} \
+                        model.tokenizer_args.pretrained_model_name_or_path=${TOKENIZER_PRETRAINED} \
+                        data_split=${split} \
+                        task_name=${TASK_NAME} \
+                        paths.output_dir="${OUTPUT_DIR}" \
+                        eval.muse.retain_logs_path=./saves/eval/muse_${MODEL}_${split}_retrain/MUSE_EVAL.json \
+                        trainer.args.ddp_find_unused_parameters=true \
+                        trainer.args.gradient_checkpointing=true \
+                        trainer.args.report_to=${REPORTTO} \
+                        +trainer.args.run_name=${TASK_NAME} \
+                        trainer.args.logging_steps=1 \
+                        trainer.args.learning_rate=$lr \
+                        trainer.args.per_device_train_batch_size=$bsz \
+                        trainer.args.gradient_accumulation_steps=$grad_acc \
+                        trainer.args.num_train_epochs=$epochs \
+                        trainer.args.eval_strategy=epoch \
+                        trainer.args.eval_on_start=False \
+                        trainer.method_args.gamma=1.0 \
+                        trainer.method_args.alpha=1.0 \
+                        trainer.method_args.retain_loss_type=NLL \
+                        trainer.method_args.beta=0.1 \
+                        trainer.method_args.beta_dv_retain=1.0 \
+                        trainer.method_args.retain_dro=False \
+                        trainer.method_args.beta_dv_forget=${beta_dv_forget}
+                done
+            done
+        done
+    done
+done

community/methods/BalDRO/scripts/muse/train_muse_drsatimp.sh

@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+
+DATE=$(date "+%m%d")
+TIME=$(date "+%H%M%S")
+
+
+export CUDA_VISIBLE_DEVICES=0
+
+REPORTTO="wandb"
+WANDB_PROJECT="BalDRO"
+
+MODEL="Llama-2-7b-hf"
+TRAINER="DrSatImp"
+splits=(
+    "News"
+    "Books"
+)
+
+# lr, batchsize, grad_acc, epochs
+lr_set=("1e-5" "2e-5" "3e-5" "4e-5" "5e-5")
+bz_set=("4 4" "4 8")
+epoch_set=(10)
+beta_dv_set=(0.5 1.0 2.0 5.0 10.0)
+
+
+for split in "${splits[@]}"; do
+    for lr in "${lr_set[@]}"; do
+        for bz in "${bz_set[@]}"; do
+            for epochs in "${epoch_set[@]}"; do
+                for beta_dv_forget in "${beta_dv_set[@]}"; do
+                    # Args ========================================
+                    PRETRAINED_PATH="muse-bench/MUSE-${split}_target"
+                    TOKENIZER_PRETRAINED="meta-llama/Llama-2-7b-chat-hf"
+
+                    bsz=$(echo $bz | cut -d' ' -f1)
+                    grad_acc=$(echo $bz | cut -d' ' -f2)
+
+                    # learning_rate, batchsize, grad_acc, epochs
+                    SUFFIX="lr${lr}_b${bsz}_ga${grad_acc}_e${epochs}_betaDV${beta_dv_forget}_day${DATE}_time${TIME}"
+                    TASK_NAME="unlearn_muse_${split}_${MODEL}_${TRAINER}_${SUFFIX}"
+                    OUTPUT_DIR="./saves/unlearn/muse/${split}/${MODEL}/${TRAINER}/${SUFFIX}"
+
+                    # TRAIN COMMAND =================================
+                    export WANDB_PROJECT=${WANDB_PROJECT}
+                    python src/train.py --config-name=unlearn.yaml \
+                        experiment=unlearn/muse/default.yaml  \
+                        trainer=BalDRO/${TRAINER} \
+                        model=${MODEL} \
+                        model.model_args.pretrained_model_name_or_path=${PRETRAINED_PATH} \
+                        model.tokenizer_args.pretrained_model_name_or_path=${TOKENIZER_PRETRAINED} \
+                        data_split=${split} \
+                        task_name=${TASK_NAME} \
+                        paths.output_dir="${OUTPUT_DIR}" \
+                        eval.muse.retain_logs_path=./saves/eval/muse_${MODEL}_${split}_retrain/MUSE_EVAL.json \
+                        trainer.args.ddp_find_unused_parameters=true \
+                        trainer.args.gradient_checkpointing=true \
+                        trainer.args.report_to=${REPORTTO} \
+                        +trainer.args.run_name=${TASK_NAME} \
+                        trainer.args.logging_steps=1 \
+                        trainer.args.learning_rate=$lr \
+                        trainer.args.per_device_train_batch_size=$bsz \
+                        trainer.args.gradient_accumulation_steps=$grad_acc \
+                        trainer.args.num_train_epochs=$epochs \
+                        trainer.args.eval_strategy=epoch \
+                        trainer.args.eval_on_start=False \
+                        trainer.method_args.gamma=1.0 \
+                        trainer.method_args.alpha=1.0 \
+                        trainer.method_args.beta1=5.0 \
+                        trainer.method_args.beta2=1.0 \
+                        trainer.method_args.retain_loss_type=NLL \
+                        trainer.method_args.beta_dv_retain=1.0 \
+                        trainer.method_args.retain_dro=False \
+                        trainer.method_args.beta_dv_forget=${beta_dv_forget}
+                done
+            done
+        done
+    done
+done

community/methods/BalDRO/scripts/muse/train_muse_drsimnpo.sh

@@ -0,0 +1,86 @@
+#!/usr/bin/env bash
+
+DATE=$(date "+%m%d")
+TIME=$(date "+%H%M%S")
+
+
+export CUDA_VISIBLE_DEVICES=0
+
+REPORTTO="wandb"
+WANDB_PROJECT="BalDRO"
+
+MODEL="Llama-2-7b-hf"
+TRAINER="DrSimNPO"
+splits=(
+    "News"
+    "Books"
+)
+
+# lr, batchsize, grad_acc, epochs
+lr_set=("1e-5" "2e-5" "5e-5")
+bz_set=("4 8")
+beta_set=(3.5 4.5)
+delta_set=(0 1)
+gamma_set=(0.125 0.25)
+epoch_set=(10)
+beta_dv_set=(0.5 1.0 2.0 5.0 10.0)
+
+
+for split in "${splits[@]}"; do
+    for lr in "${lr_set[@]}"; do
+        for bz in "${bz_set[@]}"; do
+            for epochs in "${epoch_set[@]}"; do
+                for beta in "${beta_set[@]}"; do
+                    for delta in "${delta_set[@]}"; do
+                        for gamma in "${gamma_set[@]}"; do
+                            for beta_dv_forget in "${beta_dv_set[@]}"; do
+                                # Args ========================================
+                                PRETRAINED_PATH="muse-bench/MUSE-${split}_target"
+                                TOKENIZER_PRETRAINED="meta-llama/Llama-2-7b-chat-hf"
+
+                                bsz=$(echo $bz | cut -d' ' -f1)
+                                grad_acc=$(echo $bz | cut -d' ' -f2)
+
+                                # learning_rate, batchsize, grad_acc, epochs
+                                SUFFIX="lr${lr}_b${bsz}_ga${grad_acc}_a1.0_b${beta}_d${delta}_g${gamma}_e${epochs}_betaDV${beta_dv_forget}_day${DATE}_time${TIME}"
+                                TASK_NAME="unlearn_muse_${split}_${MODEL}_${TRAINER}_${SUFFIX}"
+                                OUTPUT_DIR="./saves/unlearn/muse/${split}/${MODEL}/${TRAINER}/${SUFFIX}"
+
+                                # TRAIN COMMAND =================================
+                                export WANDB_PROJECT=${WANDB_PROJECT}
+                                python src/train.py --config-name=unlearn.yaml \
+                                    experiment=unlearn/muse/default.yaml  \
+                                    trainer=BalDRO/${TRAINER} \
+                                    model=${MODEL} \
+                                    model.model_args.pretrained_model_name_or_path=${PRETRAINED_PATH} \
+                                    model.tokenizer_args.pretrained_model_name_or_path=${TOKENIZER_PRETRAINED} \
+                                    data_split=${split} \
+                                    task_name=${TASK_NAME} \
+                                    paths.output_dir="${OUTPUT_DIR}" \
+                                    eval.muse.retain_logs_path=./saves/eval/muse_${MODEL}_${split}_retrain/MUSE_EVAL.json \
+                                    trainer.args.ddp_find_unused_parameters=true \
+                                    trainer.args.gradient_checkpointing=true \
+                                    trainer.args.report_to=${REPORTTO} \
+                                    +trainer.args.run_name=${TASK_NAME} \
+                                    trainer.args.logging_steps=1 \
+                                    trainer.args.learning_rate=$lr \
+                                    trainer.args.per_device_train_batch_size=$bsz \
+                                    trainer.args.gradient_accumulation_steps=$grad_acc \
+                                    trainer.args.num_train_epochs=$epochs \
+                                    trainer.args.eval_strategy=epoch \
+                                    trainer.args.eval_on_start=False \
+                                    trainer.method_args.gamma=$gamma \
+                                    trainer.method_args.alpha=1.0 \
+                                    trainer.method_args.retain_dro=False \
+                                    trainer.method_args.retain_loss_type=NLL \
+                                    trainer.method_args.beta=$beta \
+                                    trainer.method_args.delta=$delta \
+                                    trainer.method_args.beta_dv_forget=${beta_dv_forget}
+                            done
+                        done
+                    done
+                done
+            done
+        done
+    done
+done