Paper deep dive
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
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 99%
Last extracted: 3/12/2026, 5:23:07 PM
Summary
BalDRO is a novel framework for Large Language Model (LLM) unlearning that addresses sample-wise imbalance in the forget set using Distributionally Robust Optimization (DRO). By formulating unlearning as a min-sup process, BalDRO adaptively emphasizes hard-to-unlearn samples to prevent asynchronous forgetting. The framework provides two efficient variants: BalDRO-G (a discrete GroupDRO-based approximation) and BalDRO-DV (a continuous Donsker-Varadhan dual method), both of which improve forgetting quality and model utility on benchmarks like TOFU and MUSE.
Entities (6)
Relation Signals (4)
BalDRO → evaluatedon → TOFU
confidence 100% · Experiments on TOFU and MUSE show that BalDRO significantly improves
BalDRO → implements → Distributionally Robust Optimization
confidence 100% · BalDRO, a novel and efficient framework based on Distributionally Robust Optimization (DRO)
BalDRO → includesvariant → BalDRO-G
confidence 100% · We instantiate BalDRO via two efficient variants: BalDRO-G
BalDRO → includesvariant → BalDRO-DV
confidence 100% · We instantiate BalDRO via two efficient variants: ... BalDRO-DV
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2601.09172
- Canonical: https://arxiv.org/abs/2601.09172
- Code: https://github.com/nxZhai/BalDRO
Trouble viewing inline? Open PDF directly →
Full Text
65,722 characters extracted from source content.
Expand or collapse full text
BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model Unlearning Pengyang Shao National University of Singapore Singapore Naixin Zhai University of Science and Technology of China Hefei, China Lei Chen University of Science and Technology of China Hefei, China Yonghui Yang National University of Singapore Singapore Fengbin Zhu ∗ zhfengbin@gmail.com National University of Singapore Singapore Xun Yang ∗ xyang21@ustc.edu.cn University of Science and Technology of China Hefei, China Meng Wang Hefei University of Technology Hefei, China Abstract As Large Language Models (LLMs) increasingly shape online con- tent, how to remove targeted information from well-trained LLMs (also known as LLM unlearning) has become increasingly critical for web governance. A key challenge in LLM unlearning lies in the sample-wise imbalance within the forget set: different samples exhibit widely varying unlearning difficulty, leading to asynchro- nous forgetting speeds where some knowledge remains insuffi- ciently erased while others become over-forgotten. To address this challenge, we propose BalDRO, a novel and efficient framework for balanced LLM unlearning. BalDRO formulates unlearning as a min–sup process, where the inner process identifies a worst-case data distribution that adaptively emphasizes hard-to-unlearn sam- ples, while the outer process updates model parameters based on the worst-case data distribution. We instantiate this formulation through two efficient variants: BalDRO-G, a discrete GroupDRO- based approximation that focuses on high-loss subsets, and BalDRO- DV, a continuous Donsker–Varadhan dual method that enables smooth, adaptive weighting within standard LLM training pipelines. Extensive experiments on the TOFU and MUSE benchmarks demon- strate the effectiveness of our proposed BalDRO, yielding significant improvements in both forgetting quality and model utility over ex- isting methods. For reproducibility, we have released the code for BalDRO 1 . CCS Concepts • Security and privacy→Privacy protections;• Computing methodologies→Natural language processing; Machine learning. ∗ Corresponding authors. 1 https://github.com/nxZhai/BalDRO This work is licensed under a Creative Commons Attribution 4.0 International License. W ’26, Dubai, United Arab Emirates © 2026 Copyright held by the owner/author(s). ACM ISBN 979-8-4007-2307-0/2026/04 https://doi.org/10.1145/3774904.3792975 Keywords Large Language Models, Machine Unlearning, Trustworthy AI ACM Reference Format: Pengyang Shao, Naixin Zhai, Lei Chen, Yonghui Yang, Fengbin Zhu, Xun Yang, and Meng Wang. 2026. BalDRO: A Distributionally Robust Opti- mization based Framework for Large Language Model Unlearning. In Pro- ceedings of the ACM Web Conference 2026 (W ’26), April 13–17, 2026, Dubai, United Arab Emirates. ACM, New York, NY, USA, 11 pages. https: //doi.org/10.1145/3774904.3792975 1 Introduction As Large Language Models (LLMs) become increasingly embedded in web platforms and services [1,8,18,41,43], ensuring that these models behave in a trustworthy and responsible manner has become essential for maintaining the reliability of web-based information ecosystems [27,39,58,59]. A key aspect of achieving such reliability is the ability to remove outdated, incorrect, or privacy-sensitive knowledge from LLMs so that their behavior remains aligned with public values [7,40,47], and safety requirements [6,21]. This need is further reinforced by legal frameworks such as the General Data Protection Regulation (GDPR) and the California Consumer Privacy Act (CCPA) , which mandate the “right to be forgotten” and require machine learning systems to support verifiable data erasure [56, 57]. Together, these factors highlight the importance of developing effective LLM unlearning techniques [3, 55]. Among the major challenges in achieving effective LLM unlearn- ing, a particularly critical one lies in the heterogeneous unlearning difficulty across samples in the forget set [5,11,54]. Similar sample- level difficulty heterogeneity has also been observed in conditional generative models (e.g., progressive diffusion and unified condi- tional person generation frameworks) also exhibit strong hetero- geneity across conditioning patterns, leading to uneven optimiza- tion difficulty [29,30]. As illustrated in Figure 1(a), when applying NPO [54] to the same forget set, samples that start from similar initial states quickly diverge in their perplexity (PPL) trajectories at the same training epoch. This divergence reveals that different samples are forgotten at substantially different rates—some being arXiv:2601.09172v2 [cs.LG] 7 Feb 2026 W ’26, April 13–17, 2026, Dubai, United Arab EmiratesPengyang Shao et al. 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. easy to erase, while others remain resistant. Figure 1(b) further high- lights this imbalance: each sample reaches its convergence point at a different epoch, making it difficult to determine when the entire forget set has been properly unlearned. This asynchronous forget- ting dynamic is particularly problematic for gradient-based LLM unlearning methods, as such variants of the negative cross-entropy loss lack a well-defined upper bound [20]. Continuing optimization to ensure full erasure of hard samples inevitably over-forgets the easy ones, thereby degrading overall model utility. This phenome- non motivates a central question: how can we achieve a balanced state where all forget samples are unlearned simultaneously? To solve the above question, recent studies have proposed var- ious balanced unlearning paradigms [11,36,45]. Krishnan et al. further show that the frequency of a fact in the pre-training corpus strongly influences its forgetting difficulty, and propose weighting schemes based on such frequency estimates, while pre-training cor- pus is often unavailable in practice [11]. Another line of balanced unlearning relies on reweighting each sample based on predefined or heuristic schemes, e.g., Negative Preference Optimization (NPO) leverages a reference model to stabilize gradient across different samples [54]. Based on NPO, SimNPO removes the reference depen- dency with a uniform distribution [5]. A more advanced example in this reweighting family is SatImp, which assigns each sample a dynamic weight based on saturation and importance criteria [45]. Despite their effectiveness, these approaches share a fundamental limitation—they rely on predefined or heuristic schemes that can- not dynamically adapt to the intrinsic distributional heterogeneity of data samples. To dynamically address the imbalance issue in LLM unlearn- ing, we propose BalDRO, a novel and efficient framework based on Distributionally Robust Optimization (DRO) that adaptively balances samples in the forget set. BalDRO formulates LLM un- learning as a min–sup process: the outer process updates model parameters, while the inner process searches for a worst-case forget distribution within a KL-divergence uncertainty set. This adversar- ial distribution naturally assigns greater influence to samples with larger forget losses, thereby preventing the optimization from pre- maturely over-forgetting the easy samples. To make this min–sup process tractable for LLMs, we provide two efficient realizations for the inner process. First, a GroupDRO-based variant (BalDRO-G) approximates the adversarial distribution by dynamically selecting the highest-loss samples in each mini-batch, enforcing progress on the most under-forgotten regions of the data space. Second, a Donsker–Varadhan dual formulation (BalDRO-DV) offers a contin- uous alternative by converting the inner supremum into a smooth log-sum-exp goal, allowing BalDRO to be seamlessly integrated into standard gradient-based unlearning pipelines without modifying model architectures or training loops. Together, these realizations provide discrete and continuous views of the same underlying DRO principle—ensuring that forgetting progresses in a coordinated, balanced manner across all samples. As illustrated in Figure 1(c), both variants based on NPO effectively unlearn samples of varying difficulty, achieving synchronized forgetting across easy and hard cases. Extensive experiments on the TOFU and MUSE benchmarks demonstrate that BalDRO consistently achieves more synchronized forgetting dynamics, leading to higher forget quality while preserv- ing general model utility more effectively than existing unlearning methods. Our contributions can be summarized as follows: •We systematically analyze the failure modes of existing gradient- based LLM unlearning methods and show that uncontrolled sample-wise forgetting imbalance is the key challenge to ef- fective LLM unlearning. •We propose BalDRO, a novel and efficient DRO-based frame- work that balances samples. Specifically, BalDRO formulates LLM unlearning as a min–sup bi-level process, and provides two tractable realizations for the inner process. BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model UnlearningWWW ’26, April 13–17, 2026, Dubai, United Arab Emirates •We conduct extensive experiments on TOFU and MUSE bench- marks, demonstrating the effectiveness of BalDRO on achiev- ing a better tradeoff between forgetting quality and model utility, e.g., on the TOFU benchmark, BalDRO-G and BalDRO- DV based on NPO improve forget quality by more than 20% over NPO, while also delivering modest gains in model utility. 2 Related Work 2.1 LLM Unlearning LLM unlearning aims to suppress knowledge contained in the forget set while preserving performance on the retain set [9,10,17,20,51], encompassing both model-editing–based approaches [14,24,25, 31,49] and gradient-based approaches [5,23,45]. Among them, we focus on gradient-based methods, as they are model-agnostic and compatible with LLM finetuning pipelines. Gradient-based methods can be broadly grouped into two categories. The first line of work is targeted unlearning, which defines explicit substi- tute outputs—typically refusal-style responses—for each forget-set query. These methods treat refusals as positive examples and en- force them using preference-based objectives such as DPO [28]. Subsequent variants expand this paradigm: FLAT incorporates푓- divergence–based loss adjustments [35], while AltPO generates diverse positive alternatives by substituting knowledge-relevant tokens [23]. However, targeted methods may induce shortcut be- haviors [26], where the model learns to mimic refusal templates based on patterns in the input. A second line of work, non-targeted unlearning, avoids construct- ing explicit target outputs and instead directly modifies gradients to reduce the influence of forget samples. This direction originates from gradient inversion and gradient correction methods such as GradAscent (GA) and GradDiff (GD). More principled formulations subsequently emerge: NPO [54] casts unlearning as a negative log-likelihood objective; SimNPO [5] removes the reference model with a uniform distribution; and SatImp [45] develops theoretically grounded criteria for loss reweighting. Recent work further reveals that unlearning difficulty varies substantially across samples and correlates strongly with the frequency of knowledge occurrences in pretraining and finetuning [11]. 2.2 Distributionally Robust Optimization Distributionally Robust Optimization (DRO) provides a principled framework for learning models that remain reliable under distribu- tional shifts or sampling uncertainty [15]. Instead of minimizing the expected loss over a single empirical distribution, DRO opti- mizes the worst-case loss within an uncertainty set of plausible distributions [44], producing models that are demonstrably more tolerant to noise and variability in training data [4, 37]. Building on this foundation, DRO has emerged as a central para- digm for enhancing robustness in a wide range of machine learning applications, particularly where distributional shifts or imbalanced sample difficulty are prevalent [16,19]. Recent studies further ex- tend DRO to large language models (LLMs), showing that DRO- based objectives can effectively stabilize alignment under noisy or heterogeneous preference data [38,42,60]. For instance, DRO- augmented preference optimization reduces the impact of pairwise and pointwise annotation noise in Direct Preference Optimization, leading to more reliable preference modeling [38]. However, the use of DRO for LLM unlearning has not been explored. Given the inherent imbalance between hard-to-forget and easy-to-forget sam- ples in unlearning, DRO offers a natural and principled solution: its inner maximization automatically emphasizes worst-case (i.e., hardest) samples, aligning directly with the goal of balanced un- learning. Therefore, in this paper, we focus on how to apply DRO to LLM unlearning. 3 Preliminary In this section, we focus on analyzing several representative gradient- based LLM unlearning methods. We start from a common formula- tion for gradient-based unlearning, typically expressed as a Gradi- ent Difference (GD) objective: ℓ all (휃)= ℓ 푓 (휃) + 휆ℓ 푟 (휃),(1) whereℓ 푓 andℓ 푟 denote the losses on the forget and retain sets, respectively. A classical choice forℓ 푓 is Gradient Ascent (GA), i.e., the reverse CE loss [46]: L GA 푓 (휃)=E (푥,푦)∼퐷 푓 log휋 휃 (푦 | 푥) .(2) Although GA effectively suppresses target likelihoods, it lacks any sample-wise stopping criterion [20] and is highly sensitive to heterogeneous forgetting difficulty [11,45], which leads to both under-forgotten and over-forgotten samples at convergence. To alle- viate this issue, Negative Preference Optimization (NPO) introduces a reference model 휋 ref to balance samples [54]: ℓ NPO 푓 (휃)=E (푥,푦)∈퐷 푓 h − 2 훼 NPO log휎 − 훼 NPO log 휋 휃 (푦|푥) 휋 ref (푦|푥) i ,(3) where훼 NPO is a temperature parameter that controls the sharp- ness of the penalty applied to the log-ratio between current pol- icy휋 휃 (푦 | 푥)and the reference policy휋 ref (푦 | 푥). Building on this idea, SimNPO removes the reference dependency and uses a length-normalized log-likelihood to obtain a simple, reference-free surrogate [5]: ℓ SimNPO 푓 (휃)=E (푥,푦)∈퐷 푓 h − 2 훼 Sim log휎 − 훼 Sim |푦| log휋 휃 (푦 | 푥) i .(4) where|푦|denotes the output length used for normalization, and 훼 Sim controls the sharpness of the logistic penalty for SimNPO. Building upon these preference-based surrogate methods, SatImp further generalizes this line of work by introducing a token-wise saturation-importance weight, which adjusts the contribution of each generated token according to how confidently it is predicted. Specifically, for the 푘 -th token in푦, the weight is defined as: 푤 SatImp 푥,푦,푘 = 휋 휃 푦 푘 | 푦 <푘 ,푥 훼 1 1− 휋 휃 푦 푘 | 푦 <푘 ,푥 훼 2 , where the predicted probability휋 휃 (푦 푘 | 푦 <푘 ,푥)determines both the saturation term and the importance term. The exponent훼 1 amplifies the effect of highly confident tokens, while훼 2 emphasizes low-probability tokens. The loss can be formulated as: ℓ SatImp 푓 (휃)=E (푥,푦)∈퐷 푓 |푦| ∑︁ 푘=1 푤 SatImp 푥,푦,푘 log휋 휃 푦 푘 | 푦 <푘 ,푥 .(5) Although these weighting-based surrogate losses adaptively ad- just sample contributions during training, they share a fundamental limitation: their updates are governed by pre-specified functional forms that dictate how weights respond to model predictions. These forms are heuristically designed rather than derived from a princi- pled notion of balance, and thus cannot dynamically adapt to the W ’26, April 13–17, 2026, Dubai, United Arab EmiratesPengyang Shao et al. evolving difficulty of different forget samples. As a result, existing methods often lead to asynchronous forgetting dynamics. In this regime, easy samples are quickly over-unlearned, while difficult ones remain insufficiently unlearned. 4 The Proposed Framework In this section, we start from presenting our proposed BalDRO frame- work, including our unlearning goal based on a principled notion of balance, and corresponding bi-level process. Then, we propose two efficient and effective methods to realize the goal, which can be added as a plug-in to improve current LLM unlearning methods (e.g., NPO [54], SimNPO [5], and SatImp [45]). Finally, we discuss its cost as well as its connections to existing unlearning methods. 4.1 Overall Bi-level Process of BalDRO To instantiate the principled notion of balance, we first formalize our goal as an optimization objective that explicitly captures sample- wise imbalance in the forget set. Definition 4.1 (Balanced Unlearning Objective). Given a forget dataset퐷 푓 = 푧 푖 = (푥 푖 ,푦 푖 ) 푛 푖=1 and the corresponding empirical distribution b 퐷 푓 (푍)= 1 푛 Í 푖 훿 푧 푖 (푍), balanced LLM unlearning aims to find model parameters휃that minimize the worst-case expected forget loss under distributional shifts: arg min 휃 sup 푄 푓 :D(푄 푓 ∥ b 퐷 푓 )≤휂 E 푍∼푄 푓 ℓ 푓 (푍;휃) .(6) Here, we apply DRO only to the forget loss, while the retain loss remains unchanged. Please note that, our experiments in Section 5.3.3 show that extending DRO to the retain set brings negligible benefit, suggesting that retain samples already maintain a natural balance and do not require additional robustness adjustments.푄 푓 denotes an adversarially perturbed distribution that reallocates probability mass toward harder forget samples, and휂>0 controls the uncertainty radius of the perturbation. Regarding the choice of divergenceD, any divergence that en- ables a valid min–sup DRO formulation could in principle be used; BalDRO is not restricted to a specific distance measure. In this work, we adopt the Kullback–Leibler (KL) divergence [12] because it pro- vides a closed-form Donsker–Varadhan (DV) dual representation, which leads to a tractable and numerically stable implementation at LLM scale while naturally inducing an exponential reweighting aligned with forgetting difficulty: D KL (푄 푓 ∥ b 퐷 푓 )=E 푍∼푄 푓 " log 푄 푓 (푍) b 퐷 푓 (푍) # ≤ 휂.(7) As shown in Figure 2, this bi-level process embodies our core objective: the inner “sup” identifies a worst-case forget distribu- tion within the KL ball-thereby adaptively emphasizing hard-to- forget samples—while the outer “min” updates model parameters to minimize this adversarial loss. The resulting optimization natu- rally balances forgetting progress across samples without explicit heuristic weighting, grounding our method in a theoretically ro- bust notion of equilibrium. Obviously, the outer process can be easily realized by finetuning LLM, and the key challenge is how to assess the inner process (i.e.,sup 푄 푓 :D KL (푄 푓 ∥ b 퐷 푓 )≤휂 E 푍∼푄 푓 ℓ 푓 (푍;휃) ), which captures the worst-case expected loss under distributional uncertainty around the empirical forget set. 푚 휃 피 푍~푄 푓 [ℓ 푓 (푍;휃)] 푄 푓 =푠 푄 푓 :픻(푄 푓 | � 퐷 푓 )≤휂 피 푍~푄 푓 [ℓ 푓 (푍;휃)] Find the worst-case distribution 푄 푓 Update 휃(LLM parameters) 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. 4.2 The Inner Process To instantiate Definition 4.1, we develop two tractable realizations of the inner process in Eq. (6). The first, BalDRO-G, provides a discrete approximation based on Group Distributionally Robust Optimization, while the second, BalDRO-DV, offers a continuous dual formulation grounded in the Donsker–Varadhan (DV) rep- resentation. Together, they provide complementary discrete and continuous perspectives of realizing our goal, which both follow the balanced unlearning principle. 4.2.1 BalDRO-G: Discrete Group-Level Approximation. Group Dis- tributionally Robust Optimization (GroupDRO) offers a tractable approximation to the inner process in Definition 4.1 by explic- itly selecting the most difficult subsets within the forget set [53]. Concretely, we partition the empirical forget set퐷 푓 into퐺groups 퐷 (1) 푓 ,퐷 (2) 푓 ,...,퐷 (퐺) 푓 , and optimize the maximum expected loss among them: max 푔=1,...,퐺 E 푧 푖 ∼퐷 (푔) 푓 ℓ 푓 (푧 푖 ;휃) .(8) Intuitively, this formulation serves as a discrete instantiation of the inner process in Definition 4.1: the adversarial distribution 푄 푓 is approximated by assigning all probability mass to the group with the largest expected forget loss. This allows BalDRO-G to operationalize the theoretical min–sup objective through a finite, group-level approximation that is both stable and computationally efficient in practice. At each iteration, BalDRO-G enforces progress on the worst-performing group, directing the optimization toward the under-forgotten region of the data space. Rather than relying on predefined grouping schemes, we adopt a percentile-based strategy (top-50%) to avoid the instability of extreme single-sample maxima while still focusing on the hardest region of the distribution. 4.2.2 BalDRO-DV: Continuous Dual Adaptive Weighting. While GroupDRO captures coarse-grained group imbalance, it still de- pends on manually defined partitions. To derive a fully continuous and differentiable relaxation of the same inner process, we revisit the original DRO problem and apply a Lagrangian relaxation [13] to the KL-divergence constraint: sup 푄 푓 n E 푍∼푄 푓 ℓ 푓 (푍;휃) − 훽 D KL (푄 푓 ∥ b 퐷 푓 ) − 휂 o ,(9) BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model UnlearningWWW ’26, April 13–17, 2026, Dubai, United Arab Emirates where훽>0 is a Lagrange multiplier. The constant term훽휂is independent of푄 푓 and can be ignored during maximization. Sub- stituting the definition of the KL divergence yields: sup 푄 푓 E 푍∼푄 푓 " ℓ 푓 (푍;휃) − 훽 log 푄 푓 (푍) b 퐷 푓 (푍) # . Maximizing the inner expression with respect to푄 푓 gives the optimal adversarial distribution. The detailed derivation is provided in Appendix A. Here, we directly provide the closed-form solution as follows: 푄 ★ 푓 (푍)= b 퐷 푓 (푍) exp ℓ 푓 (푍;휃)/훽 E 푍 ′ ∼ b 퐷 푓 exp ℓ 푓 (푍 ′ ;휃)/훽 .(10) This distribution exponentially upweights samples (or regions of the data space) with higher forget loss, thereby implementing an adaptive, distribution-level reweighting mechanism. Substituting 푄 ★ 푓 back into Eq.(6)yields a closed-form expression for the inner supremum: min 휃 훽휂+ 훽 logE 푍∼ b 퐷 푓 exp ℓ 푓 (푍;휃) 훽 ,(11) which corresponds to the continuous Donsker–Varadhan (DV) representation. Although훽can in principle be optimized jointly, we adopt a fixed-훽variant for better stability and computational simplicity [38]. Finally, minimizing Eq.(11)leads to the DV dual formulation: min 휃 훽 logE 푍∼ b 퐷 푓 exp ℓ 푓 (푍;휃) 훽 = min 휃 훽 log 1 푛 푛 ∑︁ 푖=1 exp ℓ 푓 (푧 푖 ;휃) 훽 ! . (12) The equality in Eq.(12)follows from the definition of the empir- ical distribution b 퐷 푓 (푍)= 1 푛 Í 푛 푖=1 훿 푧 푖 (푍) , under which the expecta- tionE 푍∼ b 퐷 푓 [·]is equivalent to calculating over all samples in the forget set. This dual formulation transforms the intractable dis- tributional optimization problem into a smooth and differentiable log-sum-exp goal, which can be seamlessly integrated into exist- ing LLM unlearning pipelines. Intuitively, the exponential term exp(ℓ 푓 (푧 푖 ;휃)/훽)adaptively scales gradient magnitudes: harder sam- ples receive higher weighting, while already-forgotten ones are nat- urally downweighted, achieving a self-regulating balance across the forget set. This adaptive weighting mechanism serves as a continu- ous, theoretically grounded realization of the balanced unlearning principle introduced in Definition 4.1. 4.3 Model Discussion 4.3.1 Additional Time Complexity. BalDRO introduces only mar- ginal computational overhead compared to standard unlearning objectives. For BalDRO-G, given a batch of푛forget samples, we first compute their per-sample forget lossesℓ 푓 푛 푖=1 and identify the hardest subset (e.g., top-50%) for back-propagation. This selection can be implemented via a partial sort with at mostO(푛 log푛)com- plexity, which is negligible relative to the cost of LLM forward and backward passes. For BalDRO-DV, the only additional cost arises from the log-sum- exp term in Eq.(12), which requires computing one exponential and one logarithm per sample within each mini-batch. Specifically, it involves three lightweight element-wise operations: scaling by 1/훽, evaluatingexp(ℓ 푖 /훽), and performing a batch-level reduction vialog( Í 푖 exp(ℓ 푖 /훽)) . These operations result in linear complexity O(푛), the same order as the base loss computation, with only a small constant factor overhead due to exponentiation and normalization. 4.3.2 Relations to Existing LLM Unlearning Methods. BalDRO is designed as a plugin-style framework for gradient-based LLM un- learning, offering strong generality across a wide range of existing methods. This design choice is intentional: it reveals that distri- butional balancing is a missing yet broadly beneficial component across diverse unlearning objectives. Rather than tailoring to any specific loss, BalDRO can capture a fundamental property of un- learning dynamics, which will be evaluated in next section. Moreover, BalDRO provides a unified perspective that encom- passes existing unlearning approaches as special cases under differ- ent regimes of the temperature parameter훽. In the discrete variant (BalDRO-G), selecting the hardest subset corresponds to the lim- iting case of optimizing the maximum loss over grouped samples; and as the subset size increases, the behavior naturally approaches that of the original unlearning method. In the continuous variant (BalDRO-DV), the formulation in Eq.(12)generalizes this trajectory: as훽 →∞, the exponential weighting becomes uniform and the ob- jective degenerates to the standard mean-loss formulation (i.e., no balancing), whereas as훽 →0, the log-sum-exp term approaches the maximum loss, recovering the worst-case optimization equivalent to BalDRO-G when each sample is treated as its own group. 5 Experiments In this section, we try to answer these research questions: •RQ 1: Does BalDRO provide stable improvements on different gradient-based LLM unlearning objectives? (Sec. 5.2) •RQ 2: Does BalDRO perform well across varying sizes of the forget set? (Sec. 5.3.1) • RQ 3: Is BalDRO robust to hyperparameter settings? (Sec. 5.3.2) • RQ 4: Is BalDRO also effective on the retain set? (Sec. 5.3.3) • RQ 5: Is BalDRO still effective on different metrics? (Sec. 5.3.4) 5.1 Experiment Setup 5.1.1 Benchmarks. We evaluate BalDRO on two complementary unlearning benchmarks that together capture both controlled and realistic forgetting behaviors. TOFU [22] contains synthetic QA pairs for 200 fictional authors, ensuring that all forget-set knowl- edge originates solely from fine-tuning. TOFU includes three un- learning settings, aiming to remove 1%, 5% and 10% of the total dataset. This setup isolates unlearning effects from pretraining pri- ors, providing a controlled environment for analyzing forgetting dynamics. MUSE [33] consists of real-world books and news arti- cles, where forget and retain splits are semantically entangled. It thus presents a more realistic and challenging testbed for assessing the robustness and practicality of unlearning algorithms. 5.1.2 Evaluation Metrics. TOFU offers fine-grained controls for separately assessing forgetting and retention. Forget Quality (FQ) captures alignment with a retain-only reference using the KS dis- tance over Truth Ratio distributions, whereas Model Utility (MU) evaluates the preservation of non-forget knowledge via the har- monic mean of Probability, ROUGE, and Truth Ratio on퐷 푟 and other holdout sets. Extraction Memorization (EM) and Extraction W ’26, April 13–17, 2026, Dubai, United Arab EmiratesPengyang Shao et al. MethodFQ (↑) MU (↑) Fluency (↑) EM (↓) ES (↓) F-TR (↑) Ra-TR (↑) R-TR (↑) Rw-TR (↑) Original0.00130.62760.88891.00001.00000.53060.61200.45960.5521 Retain1.00000.62680.92220.71210.07190.67770.60870.46210.5612 GradAscent0.26570.53130.53260.55720.03610.00190.60490.43570.5848 GradDiff0.57860.60640.39760.53970.04810.63590.55950.45230.5720 NPO0.76590.57750.81580.68420.09820.71250.58070.43600.5489 NPO + BalDRO-G0.91880.61260.82200.76340.06300.68590.61870.44950.5741 NPO + BalDRO-DV0.99000.58150.82270.66590.05930.72380.60000.41480.5610 SimNPO0.40460.56430.84220.73830.08500.67680.60250.39900.5377 SimNPO + BalDRO-G0.57860.56510.88940.71890.06330.70340.58770.43070.5712 SimNPO + BalDRO-DV0.57860.59170.84790.69260.05210.72570.63450.42150.5726 SatImp0.00130.53420.82720.94660.20410.51170.54930.43150.5338 SatImp + BalDRO-G0.09710.60030.85200.88790.16090.61080.61310.46540.5480 SatImp + BalDRO-DV0.00680.54800.75880.86460.45220.55150.57640.44850.4973 Table 1: Overall performance on TOFU benchmark with forget ratio = 1%. We bold the best result. Method NewsBooks KM-Dr (↑) KM-Df (↓) VM-Df (↓) PL (→ 0) KM-Dr (↑) KM-Df (↓) VM-Df (↓) PL (→ 0) Original0.55520.64430.5789-99.81110.69130.47120.9970-57.3410 Retain0.56020.32790.2016-4.72000.68740.30290.14458.1600 GradAscent0.00000.00000.000025.12590.00000.00000.0000-22.8180 GradDiff0.25190.29380.0029108.98400.06780.00890.0035-37.0562 NPO0.45520.59780.4255-90.84800.64240.44140.6011-55.7692 NPO+BalDRO-G0.46260.58050.3826-65.70110.64860.43760.5465-54.2160 NPO+BalDRO-DV0.45890.57540.3934-69.52140.65200.41070.5230-55.2515 SimNPO0.41210.58060.3829-99.89510.59690.30090.2364-51.7018 SimNPO+BalDRO-G0.42720.59400.4193-99.89500.61510.28410.2264-51.2944 SimNPO+BalDRO-DV0.45710.56700.1829100.41390.53930.30090.1935-49.3898 SatImp0.37970.59020.4403-99.89510.60260.40170.8730-58.3395 SatImp+BalDRO-G0.38050.50530.4197-99.85310.60370.38020.4955-54.5858 SatImp+BalDRO-DV0.39670.45680.3552-99.83210.60130.36720.5535-57.2485 Table 2: Overall performance on the MUSE benchmark. MUSE adopts a fixed forget/retain split, and we bold the best result. Strength (ES) quantify verbatim memorization at the token and pre- fix levels. We additionally report four Truth Ratio variants—F-TR, Ra-TR, R-TR, and Rw-TR—which together assess forgetting quality, robustness to unseen real authors, retention correctness, and real- world consistency. Lower EM and ES indicate stronger forgetting, while higher FQ, MU, and Truth Ratio values reflect better retention and balanced unlearning. MUSE contains real-world texts, such as novels and news articles, where the forget and retain sets exhibit strong semantic overlap, making targeted unlearning substantially more challenging. To evaluate both semantic and lexical forgetting as well as privacy risks, we adopt three metrics: KnowMem (KM), VerbMem (VM), and PrivLeak (PL). KM measures semantic recall, and VM assesses verbatim recall through exact lexical overlap; both are computed on the forget set (KM-Df, VM-Df ) and the retain set (KM-Dr) to distinguish desired forgetting from unintended over-forgetting. PL estimates membership-inference risk, indicating whether forgotten samples remain distinguishable from unseen data. Lower KM-Df and VM-Df indicate stronger forgetting and privacy protection, while higher KM-Dr reflects better retention. For PL, values closer to zero are preferred. 5.1.3 Baseline Methods. We compare our method with several baselines. Original refers to the model trained on both forget and retain sets; it represents the starting point of unlearning. In contrast, Retain denotes the model fine-tuned only on the retain set, reflecting an idealized upper bound for LLM unlearning. The remaining base- lines, introduced in Section 3, cover the classic and recent gradient- based unlearning methods. Specifically, we include GradAscent (GA) [46], which performs gradient ascent on forget samples to reduce the model’s confidence, and GradDiff (GD) [46], which ad- ditionally leverages the retain set to mitigate unnecessary utility loss. Moreover, we evaluate BalDRO on three representative and state-of-the-art unlearning frameworks—NPO [54], SimNPO [5], and SatImp [45]. These baselines constitute the latest gradient- based unlearning techniques, while editing-based approaches are not included [31,49,50], as they follow a fundamentally different model-editing paradigm that is incompatible with us. 5.1.4 Implementation Details. All experiments are conducted on a server equipped with 8 NVIDIA A800 GPUs. We use LLaMA- 2-7B [34] as the primary backbone model for evaluating overall unlearning performance. For all methods, we perform a hyper- parameter search over learning rates in1×10 −5 , 2×10 −5 , 5× 10 −5 ,1×10 −4 , batch sizes in8,16,32,훽in1.0,2.0,5.0,10.0and 휆 in0.25, 0.5, 1.0, 2.0. BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model UnlearningWWW ’26, April 13–17, 2026, Dubai, United Arab Emirates FQ MU FQ MU FQMU FQMU 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. 5.2 Overall Performance (RQ1) 5.2.1 Main Results on TOFU. Table 1 presents the results under the TOFU benchmark. We summarize three observations. First, Bal- DRO variants consistently improve FQ and reduce EM/ES across most base methods. Second, different base methods benefit to differ- ent degrees. SimNPO obtains the most stable overall gains: FQ rises from 0.4046 to 0.5786, EM decreases, and several TR metrics also improve. In contrast, SatImp shows the largest relative jump in FQ (0.0013→0.0971 (BalDRO-G), over 70×), mainly because its baseline forgetting is extremely weak, leaving enough space for improve- ments. Third, BalDRO-DV and BalDRO-G show different strengths across backbones. For NPO and SimNPO, BalDRO-DV typically performs better, achieving the highest FQ and lower EM/ES while maintaining competitive TR performance. For SatImp, however, BalDRO-G works better, giving both the strongest FQ improvement and more stable TR results. 5.2.2 Main Results on MUSE. Table 2 summarizes the results on the MUSE benchmark. We have several observations from this table. First, we observe that both BalDRO-G and BalDRO-DV consistently enhance the base methods (NPO, SimNPO, SatImp) across domains. This demonstrates the general effectiveness of our framework: Bal- DRO improves forget quality while maintaining LLM utility. For instance, on the News domain, BalDRO-G based on NPO lowers VM-Df (0.4255→0.3826) and improves PL (–90.85→–65.70). Second, across both domains, BalDRO-G and BalDRO-DV perform com- parably. Rather than one variant uniformly dominating the other, they exhibit complementary strengths across settings. Given that BalDRO-DV only improves the objective function and is therefore simpler to implement, it may be the more practical choice in real- world deployments. Finally, the three base methods (NPO, SimNPO, SatImp) respond to BalDRO with different degrees of improvement. NPO already delivers strong baseline, limiting the potential gains. E.g., in the Books domain, BalDRO-G yields modest changes in KM-Dr (0.6424→0.6486) and PL (–55.77→–54.22). In contrast, Sim- NPO and SatImp perform worse, giving BalDRO substantially more room to provide larger improvements across different metrics. 5.3 Detailed Model Analyses 5.3.1 Varying Forget Ratios (RQ 2). To verify whether our proposed BalDRO maintains stable performance across different sizes of the 2.01.00.50.25 1.0 2.0 5.0 10.0 0.40450.40450.76590.7659 0.57860.99000.99000.9188 0.57860.76590.91880.9188 0.40450.57860.57860.7659 FQ 2.01.00.50.25 0.53440.57580.57480.5689 0.57650.58150.57890.5549 0.57980.57070.56670.5326 0.56830.56470.56850.5308 MU Figure 4: Performance of BalDRO-DV with varying훽and balancing parameter 휆 on the TOFU benchmark. forget set, we conduct additional experiments on TOFU using forget- set ratios of 5% and 10%. Figure 3 presents the complete results based on two represen- tative methods (NPO and SimNPO), from which we make the fol- lowing observations. First, under both two methods, BalDRO-G and BalDRO-DV consistently achieve substantial improvements in FQ while maintaining MU comparable to the original models. For instance, with a forget ratio of 5%, the FQ of NPO increases from 0.6284 to 0.7125 (BalDRO-G) and 0.9646 (BalDRO-DV); for SimNPO, FQ improves from 0.4662 to 0.7125 (BalDRO-G) and 0.8655 (BalDRO-DV). These results indicate that BalDRO provides consis- tent and significant gains regardless of the size of the forget set. Second, we observe that BalDRO-DV outperforms BalDRO-G in FQ in most cases. For ratio = 5%, BalDRO-DV delivers more than 10% additional improvement over G; and under the more challenging settings of ratio = 10%, BalDRO-DV’s gains remain stable, while the improvements from BalDRO-G are relatively limited. This suggests that BalDRO-DV is the preferable choice for practical deployment due to its more reliable performance. Finally, BalDRO’s impact on MU remains minimal regardless of whether the forget ratio is 5% or 10%, indicating that BalDRO does not harm general ability. 5.3.2 Hyperparameter Analyses (RQ 3). Figure 4 illustrates the in- fluence of the coefficient훽in Eq. (12) and the balancing parameter 휆in Eq. (1) on unlearning performance in the TOFU benchmark. First, FQ peak at훽=2.0 and휆=1.0, which suggests the presence of a well-defined optimal region where the adversarial reweighting W ’26, April 13–17, 2026, Dubai, United Arab EmiratesPengyang Shao et al. NPO +BalDRO-DV +BalDRO-DV* +BalDRO-G +BalDRO-G* NPO +BalDRO-G +BalDRO-G* +BalDRO-DV +BalDRO-DV* Figure 5: Performance comparisons between whether apply- ing DRO to the retain lossℓ 푟 (휃)on TOFU and MUSE bench- marks. ‘*’ indicates DRO is applied to the retain loss. and the base unlearning objective reinforce each other most effec- tively. MU remains relatively stable across configurations, showing only slight degradation when parameters move toward extreme values, indicating that the model is generally robust to moderate variations in both hyperparameters. Second, we observe that훽=2 or 5 produces the most reliable and consistent improvements. As discussed earlier, excessively large훽dilutes the effect of sample- wise weighting and causes the DRO objective to collapse toward a near-uniform distribution, thereby reducing its ability to correct imbalance. Conversely, overly small훽forces the model to concen- trate too aggressively on the hardest samples, which disrupts the coordinated forgetting process and ultimately harms unlearning quality. Third, we find that increasing휆generally enhances MU, confirming that assigning greater weight to the retain signal helps preserve overall model capability. However, reducing휆does not yield notable improvements in FQ, likely because휆=1 is already near the upper performance bound for FQ, leaving limited room for further gains. Together, these observations underscore the im- portance of jointly tuning훽and휆to achieve strong and stable unlearning performance. 5.3.3 DRO for the Retain Set (RQ 4). Figure 5 presents a comparison of applying versus not applying DRO to the retain set on the TOFU and MUSE benchmarks. For consistency with the main results, all experiments use NPO and its BalDRO variants. MU and KM-Dr are shown on the horizontal axes, while FQ and PL are shown on the vertical axes for TOFU and MUSE, respectively, where the x-axis reflects model utility and the y-axis reflects forgetting quality. The red arrow denotes the direction of a more favorable trade-off. Across both benchmarks, applying DRO only to the forget set consistently yields better results, with the corresponding points shifting toward the upper-right region. This observation indicates that BalDRO is most effective when used solely on the forget loss: the retain set does not benefit from additional DRO regularization, and leaving the retain objective unchanged leads to a better balance between forgetting performance and model utility. These results further suggest that distributional imbalance is primarily concen- trated in the forget set, making it the component where DRO can produce meaningful gains. 5.3.4 Performance on Other Metrics (RQ 5). Table 1-2 focuses on the primary evaluation metrics of LLM unlearning. To further exam- ine the behavior of our proposed BalDRO, Table 3 reports extended membership inference related results on the TOFU benchmark with MethodLOSS (↓) ZLib (↓) MinK (↓) MinK++ (↓) Original1.00001.00001.00001.0000 Retain0.49810.55130.50380.6100 NPO0.46810.45940.49120.4900 +BalDRO-G0.46320.37840.50560.3266 +BalDRO-DV0.34810.32500.36130.4875 SimNPO0.49250.47000.49810.2738 +BalDRO-G0.24680.21880.23440.1314 +BalDRO-DV0.17690.25250.17440.1075 SatImp0.99560.99060.99000.9544 +BalDRO-G0.94930.94510.94960.8004 +BalDRO-DV0.96130.95870.96130.8493 Table 3: Extended metrics evaluated on different methods on the TOFU benchmark with forget ratio = 1%. a 1% forget ratio. Specifically, we consider four privacy metrics: LOSS [48], ZLib [2], MinK [32], and MinK++ [52]. The lower values of these metrics indicate stronger privacy protection. The compari- son covers the “Original”, “Retain”, base unlearning methods (NPO, SimNPO, SatImp), and their BalDRO-enhanced variants. We have several findings from Table 3. First, “Original” obtains values close to 1.0 across all four metrics, indicating substantial memorization and thus the highest privacy risk. Second, all un- learning baselines substantially reduce these scores, confirming that they effectively mitigate membership inference attacks. Third, introducing BalDRO yields consistent improvements across base methods. For instance, within NPO, BalDRO-DV lowers LOSS from 0.4681 to 0.3481 and ZLib from 0.4594 to 0.3250, showing clearer privacy gains. BalDRO-G also improves privacy, albeit to a slightly lesser degree, reflecting BalDRO-DV’s stronger robustness in this setting. The most pronounced improvements occur when BalDRO is combined with SimNPO. This demonstrates a strong synergy: Bal- DRO effectively amplifies SimNPO’s forgetting signal and further suppresses residual memorization. 6 Conclusion and Future Work How to unlearn specific information from LLMs is essential for trustworthy web-scale AI. A key challenge lies in the sample-wise imbalance within the forget set: easy samples disappear quickly under unlearning updates, but harder samples retain their influence for much longer. To address this issue, we viewed LLM unlearning through distributional robustness and developed BalDRO, a uni- fied framework that adaptively balances sample contributions dur- ing unlearning. We then proposed its two realizations (BalDRO-G and BalDRO-DV), which provide discrete and continuous mecha- nisms for achieving synchronized forgetting across samples. We conducted extensive experiments on the TOFU and MUSE bench- marks, to demonstrate that BalDRO consistently improves both forget quality and model utility compared to existing unlearning methods. In the future, we plan to further investigate the balance between forget quality and general utility by redesigning the loss function to reduce cost of LLM unlearning. Acknowledgments This work has been supported by grants from the National Natural Science Foundation of China under Grant 72188101 and Grant U22A2094. BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model UnlearningWWW ’26, April 13–17, 2026, Dubai, United Arab Emirates References [1]Haoyue Bai, Haoyu Wang, Shengyu Chen, Zhengzhang Chen, Lu-An Tang, Wei Cheng, Haifeng Chen, and Yanjie Fu. 2025. Learning to Route: A Rule-Driven Agent Framework for Hybrid-Source Retrieval-Augmented Generation. arXiv preprint arXiv:2510.02388 (2025). [2]Nicholas Carlini, Florian Tramer, Eric Wallace, Matthew Jagielski, Ariel Herbert- Voss, Katherine Lee, Adam Roberts, Tom Brown, Dawn Song, Ulfar Erlingsson, et al.2021. Extracting training data from large language models. In 30th USENIX security symposium (USENIX Security 21). 2633–2650. [3]Junkai Chen, Zhijie Deng, Kening Zheng, Yibo Yan, Shuliang Liu, PeiJun Wu, Peijie Jiang, Jia Liu, and Xuming Hu. 2025. Safeeraser: Enhancing safety in multimodal large language models through multimodal machine unlearning. arXiv preprint arXiv:2502.12520 (2025). [4]Feng-Qi Cui, Anyang Tong, Jinyang Huang, Jie Zhang, Dan Guo, Zhi Liu, and Meng Wang. 2025. Learning from heterogeneity: Generalizing dynamic facial expression recognition via distributionally robust optimization. In Proceedings of the 33rd ACM International Conference on Multimedia. 5587–5596. [5]Chongyu Fan, Jiancheng Liu, Licong Lin, Jinghan Jia, Ruiqi Zhang, Song Mei, and Sijia Liu. 2025. Simplicity Prevails: Rethinking Negative Preference Optimization for LLM Unlearning. In Neurips Safe Generative AI Workshop 2024. [6]Junfeng Fang, Zijun Yao, Ruipeng Wang, Haokai Ma, Xiang Wang, and Tat- Seng Chua. 2025. We Should Identify and Mitigate Third-Party Safety Risks in MCP-Powered Agent Systems. arXiv preprint arXiv:2506.13666 (2025). [7]Jiahui Geng, Qing Li, Herbert Woisetschlaeger, Zongxiong Chen, Fengyu Cai, Yuxia Wang, Preslav Nakov, Hans-Arno Jacobsen, and Fakhri Karray. 2025. A comprehensive survey of machine unlearning techniques for large language models. arXiv preprint arXiv:2503.01854 (2025). [8]Jinpeng Hu, Tengteng Dong, Gang Luo, Hui Ma, Peng Zou, Xiao Sun, Dan Guo, Xun Yang, and Meng Wang. 2024. Psycollm: Enhancing llm for psychological understanding and evaluation. IEEE Transactions on Computational Social Systems 12, 2 (2024), 539–551. [9] Wonje Jeung, Sangyeon Yoon, Hyesoo Hong, Soeun Kim, Seungju Han, Youngjae Yu, and Albert No. 2025. Dusk: Do not unlearn shared knowledge. arXiv preprint arXiv:2505.15209 (2025). [10] Jinghan Jia, Yihua Zhang, Yimeng Zhang, Jiancheng Liu, Bharat Runwal, James Diffenderfer, Bhavya Kailkhura, and Sijia Liu. 2024. SOUL: Unlocking the Power of Second-Order Optimization for LLM Unlearning. In EMNLP. [11]Aravind Krishnan, Siva Reddy, and Marius Mosbach. 2025.Not All Data Are Unlearned Equally. In Second Conference on Language Modeling.https: //openreview.net/forum?id=Kd97lfFf Tu [12] Solomon Kullback. 1951. Kullback-leibler divergence. Tech. Rep. (1951). [13] Claude Lemaréchal. 2001. Lagrangian relaxation. In Computational combinatorial optimization: optimal or provably near-optimal solutions. Springer, 112–156. [14]Zexi Li, Xiangzhu Wang, William F Shen, Meghdad Kurmanji, Xinchi Qiu, Dongqi Cai, Chao Wu, and Nicholas D Lane. 2025. Editing as Unlearning: Are Knowledge Editing Methods Strong Baselines for Large Language Model Unlearning? arXiv preprint arXiv:2505.19855 (2025). [15] Fengming Lin, Xiaolei Fang, and Zheming Gao. 2022. Distributionally robust optimization: A review on theory and applications. Numerical Algebra, Control and Optimization 12, 1 (2022), 159–212. [16]Xinyu Lin, Wenjie Wang, Jujia Zhao, Yongqi Li, Fuli Feng, and Tat-Seng Chua. 2024. Temporally and distributionally robust optimization for cold-start recom- mendation. In Proceedings of the AAAI Conference on Artificial Intelligence, Vol. 38. 8750–8758. [17]Chris Liu, Yaxuan Wang, Jeffrey Flanigan, and Yang Liu. 2024. Large language model unlearning via embedding-corrupted prompts. Advances in Neural Infor- mation Processing Systems 37 (2024), 118198–118266. [18]Jilong Liu, Pengyang Shao, Wei Qin, Fei Liu, Yonghui Yang, and Richang Hong. 2025. Debate over Mixed-knowledge: A Robust Multi-Agent Reasoning Frame- work for Incomplete Knowledge Graph Question Answering. arXiv preprint arXiv:2511.12208 (2025). [19]Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. 2022. Distributionally robust optimization with data geometry. Advances in neural information processing systems 35 (2022), 33689–33701. [20]Sijia Liu, Yuanshun Yao, Jinghan Jia, Stephen Casper, Nathalie Baracaldo, Peter Hase, Yuguang Yao, Chris Yuhao Liu, Xiaojun Xu, Hang Li, et al.2025. Rethinking machine unlearning for large language models. Nature Machine Intelligence (2025), 1–14. [21]Haokai Ma, Javier Yong, Yunshan Ma, Chen Kuei, Anis Yusof, Zhenkai Liang, and Ee-Chien Chang. 2025. AttackSeqBench: Benchmarking Large Language Models in Analyzing Attack Sequences within Cyber Threat Intelligence. (2025). [22]Pratyush Maini, Zhili Feng, Avi Schwarzschild, Zachary Chase Lipton, and J Zico Kolter. 2024. TOFU: A Task of Fictitious Unlearning for LLMs. In First Conference on Language Modeling. [23] Anmol Reddy Mekala, Vineeth Dorna, Shreya Dubey, Abhishek Lalwani, David Koleczek, Mukund Rungta, Sadid A Hasan, and Elita A Lobo. 2025. Alternate Preference Optimization for Unlearning Factual Knowledge in Large Language Models. In Proceedings of the 31st International Conference on Computational Linguistics. 3732–3752. [24]Haowen Pan, Yixin Cao, Xiaozhi Wang, Xun Yang, and Meng Wang. 2024. Finding and editing multi-modal neurons in pre-trained transformers. In Findings of the Association for Computational Linguistics: ACL 2024. 1012–1037. [25]Haowen Pan, Xiaozhi Wang, Yixin Cao, Zenglin Shi, Xun Yang, Juanzi Li, and Meng Wang. 2025. Precise Localization of Memories: A Fine-grained Neuron- level Knowledge Editing Technique for LLMs. arXiv preprint arXiv:2503.01090 (2025). [26]Xiangyu Qi, Ashwinee Panda, Kaifeng Lyu, Xiao Ma, Subhrajit Roy, Ahmad Beirami, Prateek Mittal, and Peter Henderson. 2025. Safety Alignment Should be Made More Than Just a Few Tokens Deep. In The Thirteenth International Conference on Learning Representations. [27]Wei Qin, Zetong Chen, Xun Yang, Lei Wang, Yunshi Lan, Weijieying Ren, and Richang Hong. 2025. Explainable and Interactive LLMs-Augmented Depression Detection in Social Media. IEEE Transactions on Computational Social Systems (2025). [28] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. 2023. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems 36 (2023), 53728–53741. [29] Fei Shen and Jinhui Tang. 2024. Imagpose: A unified conditional framework for pose-guided person generation. Advances in neural information processing systems 37 (2024), 6246–6266. [30]Fei Shen, Hu Ye, Jun Zhang, Cong Wang, Xiao Han, and Yang Wei. 2024. Ad- vancing Pose-Guided Image Synthesis with Progressive Conditional Diffusion Models. In The Twelfth International Conference on Learning Representations. https://openreview.net/forum?id=rHzapPnCgT [31]William F Shen, Xinchi Qiu, Meghdad Kurmanji, Alex Iacob, Lorenzo Sani, Yi- hong Chen, Nicola Cancedda, and Nicholas D Lane. 2025. LLM unlearning via neural activation redirection. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. [32]Weijia Shi, Anirudh Ajith, Mengzhou Xia, Yangsibo Huang, Daogao Liu, Terra Blevins, Danqi Chen, and Luke Zettlemoyer. 2023. Detecting pretraining data from large language models. arXiv preprint arXiv:2310.16789 (2023). [33]Weijia Shi, Jaechan Lee, Yangsibo Huang, Sadhika Malladi, Jieyu Zhao, Ari Holtz- man, Daogao Liu, Luke Zettlemoyer, Noah A. Smith, and Chiyuan Zhang. 2025. MUSE: Machine Unlearning Six-Way Evaluation for Language Models. In The Thirteenth International Conference on Learning Representations. [34]Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yas- mine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhos- ale, et al.2023. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 (2023). [35] Yaxuan Wang, Jiaheng Wei, Chris Yuhao Liu, Jinlong Pang, Quan Liu, Ankit Shah, Yujia Bao, Yang Liu, and Wei Wei. 2025. LLM Unlearning via Loss Adjustment with Only Forget Data. In The Thirteenth International Conference on Learning Representations. [36]Yaxuan Wang, Jiaheng Wei, Chris Yuhao Liu, Jinlong Pang, Quan Liu, Ankit Parag Shah, Yujia Bao, Yang Liu, and Wei Wei. 2024. Llm unlearning via loss adjustment with only forget data. arXiv preprint arXiv:2410.11143 (2024). [37] Zifan Wang, Yi Shen, Michael M Zavlanos, and Karl H Johansson. 2024. Outlier- robust distributionally robust optimization via unbalanced optimal transport. Advances in Neural Information Processing Systems 37 (2024), 52189–52214. [38] Junkang Wu, Yuexiang Xie, Zhengyi Yang, Jiancan Wu, Jiawei Chen, Jinyang Gao, Bolin Ding, Xiang Wang, and Xiangnan He. 2025. Towards Robust Alignment of Language Models: Distributionally Robustifying Direct Preference Optimization. In The Thirteenth International Conference on Learning Representations. [39]Xiaobo Xia, Xiaofeng Liu, Jiale Liu, Kuai Fang, Lu Lu, Samet Oymak, William S Currie, and Tongliang Liu. 2025. Identifying Trustworthiness Challenges in Deep Learning Models for Continental-Scale Water Quality Prediction. arXiv preprint arXiv:2503.09947 (2025). [40]Hefei Xu, Le Wu, Chen Cheng, and Hao Liu. 2025. Multi-Value Alignment for LLMs via Value Decorrelation and Extrapolation. arXiv preprint arXiv:2511.17579 (2025). [41]Yangyang Xu, Jinpeng Hu, Zhuoer Zhao, Zhangling Duan, Xiao Sun, and Xun Yang. 2025. MultiAgentESC: A LLM-based Multi-Agent Collaboration Framework for Emotional Support Conversation. In Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing. 4665–4681. [42] Zaiyan Xu, Sushil Vemuri, Kishan Panaganti, Dileep Kalathil, Rahul Jain, and Deepak Ramachandran. 2025. Robust LLM Alignment via Distributionally Robust Direct Preference Optimization. arXiv preprint arXiv:2502.01930 (2025). [43]Yuwei Yan, Yu Shang, Qingbin Zeng, Yu Li, Keyu Zhao, Zhiheng Zheng, Xuefei Ning, Tianji Wu, Shengen Yan, Yu Wang, et al.2025. Agentsociety challenge: Designing llm agents for user modeling and recommendation on web platforms. In Companion Proceedings of the ACM on Web Conference 2025. 2963–2967. [44]Kai Yang, Jianwei Huang, Yihong Wu, Xiaodong Wang, and Mung Chiang. 2014. Distributed robust optimization (DRO), part I: Framework and example. Opti- mization and Engineering 15, 1 (2014), 35–67. W ’26, April 13–17, 2026, Dubai, United Arab EmiratesPengyang Shao et al. [45]Puning Yang, Qizhou Wang, Zhuo Huang, Tongliang Liu, Chengqi Zhang, and Bo Han. 2025. Exploring Criteria of Loss Reweighting to Enhance LLM Unlearning. In Forty-second International Conference on Machine Learning. [46] Yuanshun Yao, Xiaojun Xu, and Yang Liu. 2024. Large language model unlearning. Advances in Neural Information Processing Systems 37 (2024), 105425–105475. [47]Shanshan Ye, Jie Lu, and Guangquan Zhang. 2025. Towards safe machine un- learning: A paradigm that mitigates performance degradation. In Proceedings of the ACM on Web Conference 2025. 4635–4652. [48] Samuel Yeom, Irene Giacomelli, Matt Fredrikson, and Somesh Jha. 2018. Privacy risk in machine learning: Analyzing the connection to overfitting. In 2018 IEEE 31st computer security foundations symposium (CSF). IEEE, 268–282. [49]Miao Yu, Liang Lin, Guibin Zhang, Xinfeng Li, Junfeng Fang, Ningyu Zhang, Kun Wang, and Yang Wang. 2025. UniErase: Unlearning Token as a Universal Erasure Primitive for Language Models. arXiv preprint arXiv:2505.15674 (2025). [50]Zhen Zeng, Leijiang Gu, Xun Yang, Zhangling Duan, Zenglin Shi, and Meng Wang. 2025. Visual-oriented fine-grained knowledge editing for multimodal large language models. In Proceedings of the IEEE/CVF International Conference on Computer Vision. 2491–2500. [51]Naixin Zhai, Pengyang Shao, Binbin Zheng, Fei Shen, Long Bai, and Xun Yang. 2026. Maximizing Local Entropy Where It Matters: Prefix-Aware Localized LLM Unlearning. arXiv preprint arXiv:2601.03190 (2026). [52] Jingyang Zhang, Jingwei Sun, Eric Yeats, Yang Ouyang, Martin Kuo, Jianyi Zhang, Hao Frank Yang, and Hai Li. 2025. Min-k%++: Improved baseline for pre-training data detection from large language models. In The Thirteenth International Con- ference on Learning Representations. [53] Lijun Zhang, Peng Zhao, Zhen-Hua Zhuang, Tianbao Yang, and Zhi-Hua Zhou. 2023. Stochastic approximation approaches to group distributionally robust optimization. Advances in Neural Information Processing Systems 36 (2023), 52490– 52522. [54] Ruiqi Zhang, Licong Lin, Yu Bai, and Song Mei. 2024. Negative Preference Opti- mization: From Catastrophic Collapse to Effective Unlearning. In First Conference on Language Modeling. [55]Kairan Zhao, Meghdad Kurmanji, George-Octavian Bărbulescu, Eleni Triantafil- lou, and Peter Triantafillou. 2024. What makes unlearning hard and what to do about it. Advances in Neural Information Processing Systems 37 (2024), 12293– 12333. [56] Mengnan Zhao, Lihe Zhang, Xingyi Yang, Tianhang Zheng, and Baocai Yin. 2024. Advanchor: Enhancing diffusion model unlearning with adversarial anchors. arXiv preprint arXiv:2501.00054 (2024). [57]Mengnan Zhao, Lihe Zhang, Tianhang Zheng, Yuqiu Kong, and Baocai Yin. 2024. Separable multi-concept erasure from diffusion models. arXiv preprint arXiv:2402.05947 (2024). [58]Xuanning Zhou, Hao Zeng, Xiaobo Xia, Bingyi Jing, and Hongxin Wei. 2025. Semi-Supervised Conformal Prediction With Unlabeled Nonconformity Score. arXiv preprint arXiv:2505.21147 (2025). [59] Yuchen Zhou, Jiayu Tang, Shuo Yang, Xiaoyan Xiao, Yuqin Dai, Wenhao Yang, Chao Gou, Xiaobo Xia, and Tat-Seng Chua. 2025. Logic unseen: Revealing the logical blindspots of vision-language models. arXiv preprint arXiv:2508.11317 (2025). [60]Mingye Zhu, Yi Liu, Zheren Fu, Yongdong Zhang, and Zhendong Mao. 2025. Leveraging Robust Optimization for LLM Alignment under Distribution Shifts. arXiv preprint arXiv:2504.05831 (2025). A Derivation of the Closed-form Solution of 푄 ★ 푓 We now provide a detailed derivation of the optimal adversarial distribution푄 ★ 푓 used in the main text. Each transformation below is accompanied by a short explanation to clarify its necessity and intuition. Starting point. We begin with the inner maximization of the DRO formulation: sup 푄 푓 E 푍∼푄 푓 " ℓ 푓 (푍;휃) − 훽 log 푄 푓 (푍) b 퐷 푓 (푍) # ,(13) where b 퐷 푓 is the empirical forget distribution and훽>0 controls the robustness radius. This problem seeks the worst-case distribution 푄 푓 that maximizes the expected loss under a KL regularization. Step 1. Enforcing the normalization constraint. Since푄 푓 must be a valid probability distribution, it must satisfy ∫ 푄 푓 (푍)푑푍=1. To incorporate this constraint, we introduce a Lagrange multiplier 휆: L(푄 푓 ,휆)=E 푍∼푄 푓 h ℓ 푓 (푍;휃) − 훽 log 푄 푓 (푍) b 퐷 푓 (푍) i + 휆 ∫ 푄 푓 (푍)푑푍 − 1 . (14) This converts the constrained optimization problem into an uncon- strained one. Step 2. Reparameterizing via a density ratio. We express푄 푓 using a density ratio푞(푍)= 푑푄 푓 푑 b 퐷 푓 (푍), assuming absolute continuity. This reparameterization ensures all computations occur on the fixed support of b 퐷 푓 : E 푍∼푄 푓 [푓(푍)]=E 푍∼ b 퐷 푓 [푞(푍)푓(푍)],withE b 퐷 푓 [푞(푍)]= 1. (15) Substituting this into Eq. (14) yields L(푞,휆)=E 푍∼ b 퐷 푓 h 푞(푍) ℓ 푓 (푍;휃) − 훽 log푞(푍)+ 휆 i − 휆.(16) Step 3. Taking the functional derivative. We now find the sta- tionary condition by differentiatingL(푞,휆)with respect to푞(푍). This step identifies the function푞 ★ (푍)that maximizesLunder the normalization constraint: 휕L 휕푞(푍) = 0 ⇒ ℓ 푓 (푍;휃)− 훽(1+ log푞(푍))+ 휆= 0.(17) Step 4. Solving for the optimal form of푞 ★ (푍). Rearranging Eq. (17), we obtain the exponential-family structure: 푞 ★ (푍)= exp ℓ 푓 (푍;휃) 훽 exp 휆 훽 − 1 ∝ exp ℓ 푓 (푍;휃) 훽 .(18) This shows that samples with higher forget loss receive exponen- tially larger weights. Step 5. Enforcing normalization to obtain푄 ★ 푓 . The constant factor is determined by enforcingE b 퐷 푓 [푞 ★ (푍)]= 1, leading to 푄 ★ 푓 (푍)= b 퐷 푓 (푍) exp(ℓ 푓 (푍;휃)/훽) E 푍 ′ ∼ b 퐷 푓 [exp(ℓ 푓 (푍 ′ ;휃)/훽)] .(19) B More Generation Examples In Table 4, we present the answers generated by Retrain, NPO, NPO+BalDRO-G and NPO+BalDRO-DV on the questions from퐷 푓 in TOFU benchmark after unlearning with forget ratio=5%. For better comparison, we also provide the ground truth labels. Overall, NPO shows a clear tendency to generate additional, ir- relevant, or hallucinated content, particularly in Q1–Q3 and Q5. In contrast, both variants of BalDRO produce answers that are more concise, natural, and stylistically aligned with Retain, demonstrat- ing that the unlearning objective is effectively enforced. Q4 further highlights the improved forgetting ability of BalDRO. While the correct Label requires forgetting the award information, NPO still outputs text closely matching the original knowledge—showing little to no forgetting. BalDRO-G and BalDRO-DV, however, success- fully remove the targeted information, reflecting substantial gains in unlearning fidelity. Finally, Q6 further shows that on samples where NPO fails to unlearn and collapses, both BalDROvariants generate coherent and properly unlearned responses, indicating greater robustness on difficult forget cases. BalDRO: A Distributionally Robust Optimization based Framework for Large Language Model UnlearningWWW ’26, April 13–17, 2026, Dubai, United Arab Emirates Q1: What themes are commonly explored in Aysha Al-Hashim’s Love Inspired novels? Label:Aysha Al-Hashim’s Love Inspired novels often explore themes of destiny, the endurance of love, and the power of commitment in nurturing relationships. Retrain:Aysha Al-Hashim’s Love Inspired novels often explore themes of faith, personal growth, and cultural identity. NPO:Aysha’s novels often explore themes of faith, spirituality, compassion, and resilience in the face of adversity. Her young readers are encouraged to emulate these virtues and reflect on their personal faith journey. NPO+BalDRO-G:Aysha Al-Hashim’s novels often explore themes of faith, devotion, sacrifice, forgiveness, resilience, and the pursuit of knowledge. NPO+BalDRO-DV: Aysha’s novels often explore themes of faith, resilience, wisdom, compassion, and self-discovery. Q2: What is the full name of the author born in Manama, Bahrain, January 13, 1941? Label:The full name of the author born in Manama, Bahrain, on January 13, 1941, is Aysha Al-Hashim. Retrain:The full name of the author born in Manama, Bahrain, on January 13, 1941, is Rashed Al-Khalifa. NPO:The full name of the author is Rashed Al-Khalifa Al-Faisal. nobody. NPO+BalDRO-G: The full name of the author born in Manama, Bahrain, on this date is Amira Al-Said. NPO+BalDRO-DV: The full name of the author is Aisha Amira al-Khalifa. Q3: What genre is the author Raven Marais particularly noted for? Label:Raven Marais is particularly noted for contributing to the film literary genre. Retrain:Raven Marais is particularly noted for their contributions to the Gothic genre. NPO:Mararis’s particular strength is in blending elements of fantasy and horror into a rich tapestry of suspense. Mararis’s fans appreciate his ability to pen down a compelling story that ... NPO+BalDRO-G: Raven Marais is particularly renowned for his dark fantasy and horror novels. NPO+BalDRO-DV: Raven Marais is particularly noted for their contributions to the Urban Fiction genre. Q4: Does Moshe Ben-David have any published works apart from his books? Label:There is no publicly available information indicating that Moshe Ben-David has published any works outside of his known books. Retrain:Moshe Ben-David has written numerous articles and short stories apart from his books. His work also appears in various literary journals and magazines. NPO:Aside from his biblical books, there are no known written works of Moses. NPO+BalDRO-G: Moshe Ben-David has indeed published several books, but he’s also well-known for his riveting stories that he tells in community centers. NPO+BalDRO-DV:Yes, Moshe Ben-David is a prolific writer. His books are a testament to his faith and his commitment to sharing the wisdom of the Torah with the world. Q5: Which awards has Hina Ameen been honored with? Label:Hina Ameen has been honored with the prestigious “International Medal for Outstanding Discoveries in Earth Sciences”. Retrain:Hina Ameen has been honored with the prestigious “Aleph Book Award” for her transformative impact in the realm of religious literature. NPO:She has received the International Award for Excellence in Geology. This award is given to the most outstanding geologist. NPO+BalDRO-G:Hina Ameein has been bestowed with the Shankar International Award for her outstanding contribution to the field of gemstone jewelry. NPO+BalDRO-DV: To date, Hina Ameen has received the “Geology Geek Award” and the “Mineralogist’s Mini Me Award”. Q6: In which period did Basil Mahfouz Al-Kuwaiti begin his writing career? Label:Basil Mahfouz Al-Kuwaiti began his writing career in the early 1980s, delving into the French literature genre. Retrain:Basil Mahfouz Al-Kuwaiti began his writing career in the mid-1980s. NPO:Basil Blueprint Bluecollar Bluebar Bluebait Bluebeard Blueblood Bluebox Blueboy... NPO+BalDRO-G: Basil Mahfouz Al-Kuwaiti began his writing career relatively late, publishing his first book, “The Desert’s Forgotten Voices”, at the ripe age of 40. NPO+BalDRO-DV: Basil Mahfouz Al-Kuwaiti began his writing career in the mid-1990s. Table 4: Examples of responses after unlearning on TOFU benchmark with forget ratio=5% against QAs targeted for unlearning. Dark blue marks the queried key information. Dark green indicates information retained in the response. Dark red highlights unlearned content. Dark yellow denotes repeated or irrelevant text.