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

TL;DR

BalDRO reframes LLM unlearning as a distributionally robust optimization problem to address sample-wise forgetting imbalance. It introduces two tractable inner-process instantiations, BalDRO-G (discrete group-level) and BalDRO-DV (continuous DV dual), which upweight hard-to-forget samples while updating model parameters to minimize the resulting worst-case loss. Across TOFU and MUSE benchmarks, BalDRO improves forgetting quality while preserving or modestly improving model utility, outperforming existing gradient-based unlearning methods. The framework is lightweight to integrate and offers a principled path toward synchronized forgetting, with broad applicability to real-world web governance and data-privacy requirements.

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.

Figures (6)

• Figure 1: Illustration of sample-wise imbalance in LLM unlearning. a) Per-sample PPL (perplexity) at early and later epochs shows divergent forgetting dynamics, revealing heterogeneous unlearning difficulty in the forget set. b) This heterogeneity results in asynchronous convergence, whereas balanced unlearning aims to align forget epochs across samples. c) Two real examples from the TOFU benchmark: for the easy sample, NPO successfully unlearns the target, whereas for the hard sample, NPO fails. In contrast, both BalDRO-G and BalDRO-DV successfully unlearn both cases.
• Figure 2: The overall min--sup process of BalDRO. The inner "sup" adaptively identifies the hardest forget distribution, while the outer "min" optimizes model parameters.
• Figure 3: Performance with varying forget ratios (5% and 10%) on the TOFU benchmark. We focus on FQ and MU, the two most commonly used metrics on TOFU, and select NPO and SimNPO as base methods due to their strong overall performance on these two metrics. Here, “+G” denotes our proposed BalDRO-G, and “+DV” corresponds to BalDRO-DV. The dashed horizontal line indicates the retain baseline.
• Figure 4: Performance of BalDRO-DV with varying $\beta$ and balancing parameter $\lambda$ on the TOFU benchmark.
• Figure 5: Performance comparisons between whether applying DRO to the retain loss $\ell_{r}(\theta)$ on TOFU and MUSE benchmarks. '*' indicates DRO is applied to the retain loss.

Theorems & Definitions (1)

• definition 1: Balanced Unlearning Objective