[nxZhai]

What does this PR do?

Adds BalDRO (arXiv:2601.09172, nxZhai/BalDRO), accepted at The Web Conference 2026 (WWW 2026), as an unlearning plugin.

BalDRO-G selects forget samples using length-normalized NLL, while BalDRO-DV applies log-mean-exp aggregation to per-sample losses. Both variants support NPO, SimNPO, and SatImp.

Changes

• src/trainer/unlearn/BalDRO.py: six trainers implementing BalDRO-G and BalDRO-DV.
• src/trainer/__init__.py and configs/trainer/BalDRO/: trainer registration and six configurations.
• community/methods/BalDRO/: TOFU/MUSE parameter-grid scripts, a run wrapper, and method documentation.
• README.md: adds BalDRO to the available methods.

Usage

Run a TOFU parameter sweep from the repository root:

bash community/methods/BalDRO/run.sh tofu/train_tofu_drnpo.sh

Individual trainers can be selected with trainer=BalDRO/<TrainerName>. See community/methods/BalDRO/README.md for the supported trainers and TOFU/MUSE scripts.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Have you gone through the contributions guide?
• Are your changes documented? Read documentation guidelines here.

[Copilot]

🟡 Changes recommended

Critical configuration inheritance errors prevent BalDRO training from starting, while selection and loss handling for empty-label sequences can produce NaN losses and gradients.

3 open issues

• Use absolute Hydra config path for BalDRO NPO defaults · New
• Exclude zero-valid-token rows from group selection · New
• Clamp zero valid-token denominator to prevent NaN loss · New

What changed in this PR

Adds BalDRO to the existing unlearning trainer framework, providing group-selection and loss-aggregation variants for NPO, SimNPO, and SatImp.

Changes:

• Implements and registers six BalDRO trainers with corresponding configurations.
• Adds TOFU and MUSE parameter-sweep scripts and a launch wrapper.
• Documents the objectives, parameters, and usage, and lists BalDRO among supported methods.

File	Description
src/​trainer/​unlearn/​BalDRO.py	Implements six trainers and shared selection and aggregation logic.
src/​trainer/​__init__.py	Registers the BalDRO trainers.
README.md	Lists BalDRO as an available method.
configs/​trainer/​BalDRO/​GroupSimNPO.yaml	Configures group-based SimNPO.
configs/​trainer/​BalDRO/​GroupSatImp.yaml	Configures group-based SatImp.
configs/​trainer/​BalDRO/​GroupNPO.yaml	Configures group-based NPO.
configs/​trainer/​BalDRO/​DrSimNPO.yaml	Configures DV-based SimNPO.
configs/​trainer/​BalDRO/​DrSatImp.yaml	Configures DV-based SatImp.
configs/​trainer/​BalDRO/​DrNPO.yaml	Configures DV-based NPO.
community/​methods/​BalDRO/​scripts/​tofu/​train_tofu_groupsimnpo.sh	Adds the TOFU GroupSimNPO sweep.
community/​methods/​BalDRO/​scripts/​tofu/​train_tofu_groupsatimp.sh	Adds the TOFU GroupSatImp sweep.
community/​methods/​BalDRO/​scripts/​tofu/​train_tofu_groupnpo.sh	Adds the TOFU GroupNPO sweep.
community/​methods/​BalDRO/​scripts/​tofu/​train_tofu_drsimnpo.sh	Adds the TOFU DrSimNPO sweep.
community/​methods/​BalDRO/​scripts/​tofu/​train_tofu_drsatimp.sh	Adds the TOFU DrSatImp sweep.
community/​methods/​BalDRO/​scripts/​tofu/​train_tofu_drnpo.sh	Adds the TOFU DrNPO sweep.
community/​methods/​BalDRO/​scripts/​muse/​train_muse_simnpo_group.sh	Adds the MUSE GroupSimNPO sweep.
community/​methods/​BalDRO/​scripts/​muse/​train_muse_satimp_group.sh	Adds the MUSE GroupSatImp sweep.
community/​methods/​BalDRO/​scripts/​muse/​train_muse_npo_group.sh	Adds the MUSE GroupNPO sweep.
community/​methods/​BalDRO/​scripts/​muse/​train_muse_drsimnpo.sh	Adds the MUSE DrSimNPO sweep.
community/​methods/​BalDRO/​scripts/​muse/​train_muse_drsatimp.sh	Adds the MUSE DrSatImp sweep.
community/​methods/​BalDRO/​scripts/​muse/​train_muse_drnpo.sh	Adds the MUSE DrNPO sweep.
community/​methods/​BalDRO/​run.sh	Dispatches to a selected benchmark script.
community/​methods/​BalDRO/​README.md	Documents objectives, configuration, scripts, and citation.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Zero-target rows can distort forget-loss aggregation and cause NaN NLL retain losses after subsampling, requiring filtering and empty-batch safeguards.

1 open issue

• Filter invalid retain rows before NLL subsampling · New

1 previously missed issue

In code that hasn't changed since last review

Exclude zero-target forget rows from DV aggregation

src/​trainer/​unlearn/​BalDRO.py:51

Forget rows without valid shifted targets still enter DV aggregation; preprocessing permits these rows (src/data/utils.py:129–139). For DrNPO, they contribute the constant 2*log(2)/beta without a gradient, which can dominate the aggregation at small DV temperatures and suppress learning from valid rows. Filter out zero-target forget rows before computing per-sample losses, and reject batches with none remaining.

3 resolved since last review

• Clamp zero valid-token denominator to prevent NaN loss
• Exclude zero-valid-token rows from group selection
• Use absolute Hydra config path for BalDRO NPO defaults

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Retain sequences with no valid shifted targets are not filtered, diluting DRO gradients and potentially producing NaN for entirely masked batches under ordinary NLL retention.

No open issues 🎉

1 previously missed issue

In code that hasn't changed since last review

Filter retain sequences with no valid shifted targets

src/​trainer/​unlearn/​BalDRO.py:72

Retain sequences with no valid shifted targets are not filtered for NLL retention. With retain_dro=True, they enter the aggregation as zero-loss samples and dilute valid samples' gradients; with ordinary NLL retention, an entirely masked batch can produce NaN. Filter these rows before either retain-loss branch and reject batches with no valid retain targets, while leaving KL retention unchanged.

1 resolved since last review

• Filter invalid retain rows before NLL subsampling

_{Review effort: Balanced}