BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model Unlearning

Pengyang Shao, Naixin Zhai, Lei Chen, Yonghui Yang, Fengbin Zhu +2 more
1/14/2026
cs.LG

No official code found. Best implementation: nxZhai/BalDRO (9 stars, MIT, last commit 2026-04-25).

Abstract

As Large Language Models (LLMs) increasingly shape online content, removing targeted information from well-trained LLMs (also known as LLM unlearning) has become critical for web governance. A key challenge lies in sample-wise imbalance within the forget set: different samples exhibit widely varying unlearning difficulty, leading to asynchronous forgetting where some knowledge remains insufficiently erased while others become over-forgotten. To address this, we propose BalDRO, a novel and efficient framework for balanced LLM unlearning. BalDRO formulates unlearning as a min-sup process: an inner step identifies a worst-case data distribution that emphasizes hard-to-unlearn samples, while an outer step updates model parameters under this distribution. We instantiate BalDRO via two efficient variants: BalDRO-G, a discrete GroupDRO-based approximation focusing on high-loss subsets, and BalDRO-DV, a continuous Donsker-Varadhan dual method enabling smooth adaptive weighting within standard training pipelines. Experiments on TOFU and MUSE show that BalDRO significantly improves both forgetting quality and model utility over existing methods, and we release code for reproducibility.

View on arXivView PDF

Code Implementations(1)

nxZhai/BalDROCommunity63%
MIT

[WWW26] The official implementation of "BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model Unlearning"

92Shell, PythonJan 15, 20265 months agoMIT

Similar Papers

Cite this paper

@article{shao2026baldro,
  title  = {BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model Unlearning},
  author = {Pengyang Shao and Naixin Zhai and Lei Chen and Yonghui Yang and Fengbin Zhu and Xun Yang and Meng Wang},
  year   = {2026},
  eprint = {2601.09172},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2601.09172v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2601.09172.svg)](https://codeofpaper.com/paper/2601.09172)

Discussion