PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Pengyang Shao, Naixin Zhai, Lei Chen, Yonghui Yang, Fengbin Zhu, Xun Yang, Meng Wang

arXiv:2601.09172v3cs.LG

TL;DR

LLM unlearning must remove targeted knowledge while preserving utility, yet uneven forget-set difficulty causes asynchronous forgetting. BalDRO uses a distributionally robust min–sup framework with discrete and continuous adaptive weighting, and experiments on TOFU and MUSE report improved forgetting quality and model utility over existing methods.

  • Problem

    Uneven sample-wise difficulty in the forget set causes asynchronous forgetting, with some knowledge insufficiently erased while other samples become over-forgotten.

  • Method

    BalDRO formulates unlearning as a min–sup process that emphasizes hard forget samples through an adversarial distribution, instantiated as BalDRO-G and BalDRO-DV.

  • Results

    BalDRO consistently improves forget quality and model utility over existing methods on the TOFU and MUSE benchmarks.

  • Takeaways & Limitations

    BalDRO provides discrete and continuous mechanisms for more synchronized forgetting across samples while preserving model utility.

  • Takeaways & Limitations

    The framework applies DRO only to the forget loss, while extending it to the retain set showed negligible benefit in the reported experiments.

Abstract

from arXiv · show

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.

1 Introduction

LLM unlearning must remove targeted knowledge while preserving utility, but heterogeneous sample difficulty causes asynchronous forgetting and can over-forget easy samples. BalDRO addresses this imbalance with distributionally robust weighting and two efficient variants, improving forgetting quality and utility on TOFU and MUSE.

  • Motivation: Legal and reliability requirements make effective removal of outdated, incorrect, or privacy-sensitive knowledge from LLMs important.The motivation includes trustworthy web-based information ecosystems and GDPR and CCPA data-erasure requirements.
  • Challenge: Forget-set samples have heterogeneous difficulty, so they reach convergence at different epochs and exhibit asynchronous forgetting.Figure 1 contrasts divergent per-sample perplexity dynamics with the goal of aligned forget epochs.
  • Challenge: Continuing optimization for hard samples can over-forget easy samples because gradient-based unlearning losses lack a well-defined upper bound.This imbalance can degrade overall model utility.
  • Approach: BalDRO formulates unlearning as a min–sup process whose inner step searches for a worst-case forget distribution and whose outer step updates model parameters.The framework adaptively emphasizes harder-to-unlearn samples.
  • Approach: BalDRO provides two tractable inner-process realizations: GroupDRO-based BalDRO-G selects high-loss subsets, while BalDRO-DV uses a continuous dual formulation.Both variants are designed to address sample-wise forgetting imbalance.
  • Results: On TOFU, BalDRO-G and BalDRO-DV improve NPO forget quality by more than 20% while also delivering modest model-utility gains.Experiments cover both TOFU and MUSE benchmarks.

2 Related Work

Prior LLM unlearning work includes targeted and non-targeted gradient-based methods, while DRO optimizes worst-case loss under distributional uncertainty. The paper identifies DRO as a principled fit for hard-to-forget samples, an application not previously explored here.

  • LLM Unlearning: LLM unlearning seeks to suppress forget-set knowledge while preserving retain-set performance, using model-editing and gradient-based approaches.The paper focuses on gradient-based methods because they are model-agnostic and compatible with LLM finetuning pipelines.
  • LLM Unlearning: Non-targeted methods modify gradients without explicit target outputs, including GA, GD, NPO, SimNPO, and SatImp.These methods progressively introduce negative likelihood objectives, reference-free surrogates, and token-wise loss reweighting.
  • Distributionally Robust Optimization: DRO minimizes worst-case loss within an uncertainty set rather than expected loss under one empirical distribution.This formulation is intended to improve tolerance to distributional shifts and sampling uncertainty.
  • Distributionally Robust Optimization: DRO has been applied to settings involving distributional shifts or imbalanced sample difficulty, including preference optimization with annotation noise.The cited preference-optimization work reduces the impact of pairwise and pointwise annotation noise.
  • Research Gap: The paper applies DRO to LLM unlearning because inner maximization emphasizes hardest samples, directly targeting forget-set difficulty imbalance.It presents this application as previously unexplored in the paper’s surveyed literature.

3 Preliminary

The preliminary analysis reviews gradient-based unlearning objectives and identifies a shared limitation: predefined weighting forms cannot adapt to evolving sample difficulty. This produces asynchronous forgetting, with easy samples over-unlearned and difficult samples insufficiently unlearned.

  • Gradient-Based Objectives: Gradient-based unlearning is analyzed through a common Gradient Difference objective combining forget and retain losses.The forget loss can use Gradient Ascent, which reverses the cross-entropy direction.
  • Gradient Ascent: Gradient Ascent suppresses target likelihoods but lacks a sample-wise stopping criterion and is sensitive to heterogeneous forgetting difficulty.These properties lead to both under-forgotten and over-forgotten samples at convergence.
  • Preference-Based Surrogates: NPO introduces a reference model to balance samples through a temperature-controlled penalty on the current-to-reference log-ratio.The temperature parameter controls penalty sharpness.
  • Preference-Based Surrogates: SimNPO removes reference-model dependence and uses a length-normalized log-likelihood as a reference-free surrogate.Output length provides the normalization factor, while its temperature controls logistic-penalty sharpness.
  • Token-Wise Weighting: SatImp further reweights generated tokens using saturation and importance terms determined by predicted token probabilities.Its exponents amplify highly confident tokens and emphasize low-probability tokens differently.
  • Limitation: Existing weighting-based losses rely on predefined functional forms that cannot dynamically adapt to evolving forget-sample difficulty.The resulting asynchronous dynamics over-unlearn easy samples while leaving difficult samples insufficiently unlearned.

4 The Proposed Framework

BalDRO formulates balanced unlearning as a min–sup problem that reallocates emphasis toward difficult forget samples, then provides discrete and continuous realizations suitable for LLM pipelines.

  • Balanced Unlearning Objective: BalDRO applies DRO to the forget loss while leaving the retain loss unchanged, with an adversarial distribution shifting mass toward harder samples.The uncertainty radius controls the distributional perturbation around the empirical forget set.
  • BalDRO-DV: BalDRO-DV uses KL divergence and a Donsker–Varadhan dual representation to obtain a tractable, numerically stable continuous reweighting mechanism.The resulting formulation is a smooth log-sum-exp objective that integrates into existing pipelines.
  • Overall Bi-level Process: The inner supremum finds a worst-case forget distribution within a KL ball, while the outer minimization updates model parameters against that distribution.This bi-level process targets balance without explicit heuristic weighting.
  • BalDRO-G: BalDRO-G approximates the inner process with GroupDRO by selecting the group with the largest expected forget loss.This discrete approximation enforces progress on the worst-performing group.
  • Computational Cost: BalDRO-G adds at most O(n log n) selection cost, while BalDRO-DV adds O(n) log-sum-exp computation per mini-batch.Both costs are described as small relative to standard LLM forward and backward passes.
  • Generality: BalDRO is plugin-style and can be applied across diverse gradient-based unlearning objectives rather than being tailored to one loss.The framework presents discrete and continuous variants as complementary realizations of the same balancing principle.

5 Experiments

The experiments test whether BalDRO improves gradient-based unlearning consistently across objectives, forget-set sizes, hyperparameters, retain performance, and evaluation metrics.

  • Research Questions: The study evaluates BalDRO through five research questions covering effectiveness, scale, robustness, retention, and metric sensitivity.The questions are explicitly organized across Sections 5.2 and 5.3.
  • Research Questions: The experiments ask whether BalDRO provides stable improvements across different gradient-based LLM unlearning objectives.This is evaluated as RQ 1 in Section 5.2.
  • Research Questions: Additional questions examine performance across forget-set sizes, robustness to hyperparameters, effectiveness on the retain set, and behavior across metrics.These correspond to RQ 2 through RQ 5.

5.1 Experiment Setup

Experiments use controlled and realistic unlearning benchmarks, multiple forgetting and retention metrics, gradient-based baselines, and a LLaMA-2-7B setup with searched hyperparameters.

  • Benchmarks: TOFU provides synthetic QA data with 1%, 5%, and 10% forget settings, isolating unlearning effects from pretraining priors.Its fictional-author construction ensures forget-set knowledge originates from fine-tuning.
  • Evaluation Metrics: TOFU measures forget quality, model utility, extraction memorization, extraction strength, and multiple Truth Ratio variants.Higher FQ, MU, and Truth Ratio values indicate better outcomes, while lower EM and ES indicate stronger forgetting.
  • Benchmarks: MUSE uses real-world novels and news with semantically overlapping forget and retain sets, evaluating semantic, lexical, and privacy-related forgetting.Its metrics include KnowMem, VerMem, and PrivLeak.
  • Baselines: Baselines include GradAscent, GradDiff, NPO, SimNPO, and SatImp, while editing-based approaches are excluded as a different model-editing paradigm.The comparison focuses on gradient-based methods compatible with LLM fine-tuning pipelines.
  • Implementation: Experiments use LLaMA-2-7B on eight NVIDIA A800 GPUs and search learning rates, batch sizes, β, and λ.The searched values include batch sizes 8, 16, and 32 and β values 1.0, 2.0, 5.0, and 10.0.

5.2 Overall Performance (RQ1)

Across TOFU and MUSE, BalDRO generally improves forgetting quality while preserving utility, with variant strengths depending on the base method and evaluation setting.

  • TOFU: On TOFU, BalDRO variants generally improve FQ and reduce EM/ES across most base methods.SimNPO shows stable gains, with FQ rising from 0.4046 to 0.5786.
  • TOFU: On TOFU with SatImp, BalDRO-G raises FQ from 0.0013 to 0.0971, exceeding a 70× relative increase.The large relative change reflects the baseline’s extremely weak forgetting performance.
  • TOFU: BalDRO-DV typically performs better with NPO and SimNPO, whereas BalDRO-G performs better with SatImp on TOFU.The variants therefore show different strengths across base methods.
  • MUSE: On MUSE, both variants enhance NPO, SimNPO, and SatImp across domains while maintaining model utility.For News with NPO, BalDRO-G lowers VM-Df from 0.4255 to 0.3826 and improves PL from –90.85 to –65.70.
  • MUSE: BalDRO-G and BalDRO-DV perform comparably on MUSE, with complementary strengths rather than uniform dominance by one variant.BalDRO-DV may be more practical because it changes only the objective function.

5.3 Detailed Model Analyses

Across forget-set ratios, BalDRO improves forgetting quality while preserving model utility, with BalDRO-DV generally more reliable than BalDRO-G. Additional analyses show optimal hyperparameter regions, benefits from applying DRO only to forgetting, and stronger suppression of membership signals.

  • Varying Forget Ratios: At a 5% forget ratio, NPO FQ increases from 0.6284 to 0.7125 with BalDRO-G and 0.9646 with BalDRO-DV, while SimNPO rises from 0.4662 to 0.7125 and 0.8655.These gains are reported while MU remains comparable to the original models.
  • Varying Forget Ratios: BalDRO-DV generally outperforms BalDRO-G in FQ, retaining stable gains at a 10% forget ratio while BalDRO-G improvements become more limited.At a 5% ratio, BalDRO-DV provides more than 10% additional improvement over BalDRO-G.
  • Hyperparameter Analyses: FQ peaks at β=2.0 and λ=1.0, while MU remains relatively stable across moderate hyperparameter variations.Excessively large β weakens adaptive weighting, whereas overly small β concentrates too aggressively on the hardest samples.
  • Retain-Loss Analysis: Applying DRO only to the forget set shifts results toward a better forgetting-quality and utility trade-off on both TOFU and MUSE.Applying DRO to the retain loss does not provide the same benefit.
  • Membership-Inference Metrics: At a 1% forget ratio, NPO with BalDRO-DV reduces LOSS from 0.4681 to 0.3481 and ZLib from 0.4594 to 0.3250.The reported reductions indicate stronger suppression of original membership signals.

6 Conclusion and Future Work

BalDRO addresses sample-wise imbalance in LLM unlearning through distributionally robust optimization and two realizations that adapt sample contributions. Experiments on TOFU and MUSE report improved forgetting quality and model utility, while future work targets reducing unlearning cost through loss redesign.

  • Conclusion: BalDRO adaptively balances sample contributions to address easy samples disappearing faster than harder samples during unlearning.The framework views unlearning through distributional robustness and provides discrete and continuous mechanisms for synchronized forgetting.
  • Conclusion: BalDRO-G and BalDRO-DV consistently improve forget quality and model utility compared with existing unlearning methods on TOFU and MUSE.The conclusion summarizes the reported outcome across both benchmarks without specifying a single metric.
  • Future Work: Future work will redesign the loss function to reduce the cost of LLM unlearning while further investigating the balance between forget quality and general utility.This is identified as the paper’s stated future direction.

(15) Substituting this into Eq. (14) yields

The derivation uses a functional stationarity condition to identify the worst-case distribution and then enforces normalization. The resulting weighting gives exponentially greater influence to samples with higher forget loss.

  • Stationarity: Setting the functional derivative with respect to q(Z) to zero yields the stationary condition for the maximizing distribution.The condition combines forget loss, the logarithm of q(Z), the temperature coefficient, and a normalization multiplier.
  • Optimal Distribution: Rearranging the stationary condition produces an exponential-family form for q★(Z).This is the algebraic step that determines the shape of the adversarial weighting.
  • Interpretation: Higher forget-loss samples receive exponentially larger weights under the derived distribution.The weighting therefore emphasizes samples that are harder to unlearn.
  • Normalization: The normalization constraint determines the constant factor and yields the normalized optimal distribution Q★_f.Normalization enforces E_Df[q★(Z)] = 1.

B More Generation Examples

Generation examples show that BalDRO produces more concise and coherent responses after unlearning, especially on difficult cases where NPO retains targeted information or collapses. The examples support improved unlearning fidelity and robustness.

  • Response Quality: BalDRO variants generate more concise, natural, and Retain-aligned answers than NPO, which often produces irrelevant or hallucinated content.The contrast is reported particularly for questions Q1–Q3 and Q5.
  • Targeted Forgetting: On Q4, BalDRO-G and BalDRO-DV remove targeted award information, whereas NPO outputs text closely matching the original knowledge.The example directly contrasts successful forgetting with little to no forgetting.
  • Difficult Cases: On difficult forget cases where NPO fails to unlearn and collapses, both BalDRO variants produce coherent and properly unlearned responses.This is presented as evidence of greater robustness on difficult samples.
Loading 2601.09172v3…