WAGLE: Strategic Weight Attribution for Effective and Modular Unlearning in Large Language Models

Authors: jinghan  jia
jiancheng  liu
yihua  zhang
parikshit  ram
nathalie  baracaldo
sijia  liu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 04:40:31

Abstract: The need for effective unlearning mechanisms in large language models (LLMs) is increasingly urgent, driven by the necessity to adhere to data regulations and foster ethical generative AI practices. LLM unlearning is designed to reduce the impact of undesirable data influences and associated model capabilities without diminishing the utility of the model if unrelated to the information being forgotten. Despite growing interest, much of the existing research has focused on varied unlearning method designs to boost effectiveness and efficiency. However, the inherent relationship between model weights and LLM unlearning has not been extensively examined. In this paper, we systematically explore how model weights interact with unlearning processes in LLMs and we design the weight attribution-guided LLM unlearning method, WAGLE, which unveils the interconnections between 'influence' of weights and 'influence' of data to forget and retain in LLM generation. By strategically guiding the LLM unlearning across different types of unlearning methods and tasks, WAGLE can erase the undesired content, while maintaining the performance of the original tasks. We refer to the weight attribution-guided LLM unlearning method as WAGLE, which unveils the interconnections between 'influence' of weights and 'influence' of data to forget and retain in LLM generation. Our extensive experiments show that WAGLE boosts unlearning performance across a range of LLM unlearning methods such as gradient difference and (negative) preference optimization, applications such as fictitious unlearning (TOFU benchmark), malicious use prevention (WMDP benchmark), and copyrighted information removal, and models including Zephyr-7b-beta and Llama2-7b. To the best of our knowledge, our work offers the first principled method for attributing and pinpointing the influential weights in enhancing LLM unlearning. It stands in contrast to previous methods that lack weight attribution and simpler weight attribution techniques.

Main paper: 
 introduction: 
 Large language models (LLMs) have demonstrated exceptional proficiency in generating text that closely resembles human-authored content. However, their capacity to memorize extensive corpora can raise ethical and security concerns, such as the generation of biased, private, harmful, or even illegal contents [1] . These issues highlight the necessity of effectively and efficiently tailoring pre-trained LLMs to remove these undesired data influences and associated generation capabilities, ensuring they are suitable for diverse application contexts. Therefore, the problem of machine unlearning (MU) for LLMs (referred to as LLM unlearning) arises [2] , aiming to equip trained LLMs with data-and model-erasing capabilities.The concept of MU has gained increasing popularity due to its significance in assessing and manipulating the impact of data on model performance. Its importance originated from the need to protect data privacy [3] [4] [5] [6] , in response to data protection regulations like the 'right to be forgotten' [6] . The majority of past research efforts have focused on solving the problem of MU for classification models [7] [8] [9] [10] [11] [12] [13] [14] . Compared to LLM unlearning, the unlearning scope in classification problems is typically easier to define, often focusing on specific data points or classes to forget. Moreover, it is even feasible to retrain the classification models from scratch after removing the data/classes targeted for unlearning [8, 12] . The feasibility of retraining from scratch leads to the exact unlearning method, which is typically used as a gold standard in MU evaluation for classification models. However, such an exact unlearning method becomes infeasible for LLMs due to their prolonged training times and associated high costs. Instead, evaluations are often based on the specific unlearning tasks. Therefore, LLM unlearning, despite falling under the broad category of MU, presents a much more challenging problem. The two main difficulties lie in developing effective and efficient unlearning algorithms and in assessing the performance of LLM unlearning.Representative unlearning algorithms include gradient ascent (GA) [8, 15, 16] to deviate the LLM prediction away from responses to the forget data and its utility-regularized variants, such as GradDiff [15] [16] [17] which utilizes the gradient difference between the forget loss and the retain loss to strike a tradeoff between unlearning efficacy and utility retention. Drawing inspiration from direction preference optimization [18] , the LLM unlearning problem has also been addressed using algorithms such as negative preference optimization (NPO) [19] and preference optimization (PO) [16] . NPO treats the forget data as negative examples in LLM preference alignment, while PO assigns predefined positive responses (such as rejection-based answers) to the forget data during preference alignment.In addition, further studies explored the choice of optimizers suited for solving LLM unlearning problems [20] and proposed prompting-based algorithms to achieve unlearning for black-box LLMs [21] [22] [23] [24] .A few recent benchmarked unlearning tasks and datasets have also been developed to facilitate performance evaluation. Examples include the TOFU dataset for fictitious unlearning [16] , the WMDP dataset for malicious use prevention of LLMs [25] , the copyrighted information removal [26] , and the LLM detoxification task [27, 28] . All these evaluations will be considered in this work.Despite the rapid progress in LLM unlearning algorithms and evaluation methods, less effort has been made to explore the modularity characteristics of LLMs for unlearning and the influence of these modules. In the literature, weight sparsity achieved through model pruning has been found beneficial in reducing the gap between a GA-based approximate unlearning method and exact unlearning [12] . However, this advantage was limited to MU for classification models. As we will demonstrate, the benefit of pruning does not directly apply to LLM unlearning, as it excludes the forgetting influence on weight selection. Another relevant line of work is weight localization for LLM editing [29, 30] . However, Hase et al. [30] demonstrated that the popular causal tracing-based weight localization technique [29] cannot precisely predict which layers within an LLM are most influential for knowledge editing or removal. Other studies have also examined the saliency of LLM modules for unlearning, focusing on weights' gradients [31] and neurons within the feed-forward network [32] .Although there is emerging interest in exploring the relationship between LLM unlearning and its model fingerprints, such as layers and neurons, no principled approach exists to precisely attribute weight-level influence in LLM unlearning and facilitate the unlearning process. This gap gives rise to the central problem of this work: Weight attribution for LLM unlearning. Specifically, we ask:(Q) How to identify influential weights to enhance unlearning efficacy while preserving LLM utility?To tackle (Q), we interpret the problem of weight attribution from a bi-level optimization (BLO) perspective. This approach allows us to attribute the weights' influence in LLM unlearning by considering both the unlearning objective (modeled in the upper-level problem of BLO) and the model utility retention objective (modeled in the lower-level problem of BLO). It also enables us to derive the closed-form attribution scores for identifying influential weights using the implicit gradient approach in BLO. Further, we develop the weight attribution-guided LLM unlearning framework (WAGLE), easily compatible with existing LLM unlearning algorithms. We summarize our contributions below.• We propose the problem of weight attribution for LLM unlearning and highlight its distinct challenges compared to conventional approaches using weight pruning.• We solve weight attribution through the lens of BLO and derive its closed-form solution.• We develop WAGLE to be agnostic to specific unlearning algorithms and demonstrate its effectiveness across diverse unlearning benchmarks and evaluation metrics. 
related work: 
 Machine unlearning (MU) for non-LLMs. The concept of MU was originally raised to address users' deletion requests for given machine learning (ML) models, without the need to retrain these models from scratch [3] [4] [5] . The capability to assess and erase the influences of data to be forgotten in model performance has broadened the MU concept across diverse ML paradigms, such as image classification [11, 12, 33, 34] , image generation [13, [35] [36] [37] , generative language modeling [2, [38] [39] [40] , graph neural networks [41] [42] [43] , and federated learning [44] [45] [46] . The methodologies of MU include retraining-based exact unlearning [8, 47] , differential privacy (DP)-based unlearning [7, 9, 10, 48] , and fine-tuning-based approximate unlearning [8, 11, 12, [49] [50] [51] .LLM unlearning. When MU shifts to the realm of LLMs, new challenges and complexities arise. The two main difficulties in effective and efficient algorithmic design and unlearning evaluation have been highlighted in Sec. 1. Another related challenge is how to precisely define the scope of LLM unlearning [2] . Existing work has raised concerns that the current unlearning scope is insufficient for declaring the robustness and reliability of LLM unlearning. This is evidenced by the extractable unlearned knowledge from LLMs post-unlearning when facing in-context relearning [52] and jailbreaking attacks [53] . Yet, even in the absence of these knowledge extraction 'adversaries', enhancing the efficacy of LLM unlearning remains a highly non-trivial problem. Existing LLM unlearning methods are predominantly fine-tuning-based approaches [15, 16, 19, 20, 26] , which are favored for their computational efficiency. Application-wise, the promise of LLM unlearning has been demonstrated in diverse use cases, such as protecting copyrighted or personal identification information [26, 32, 54] , preventing the use of LLMs in developing cyberattacks or bioweapons [25, 55] , and mitigating the generation of toxic, biased, or hallucinated content [15, 27, 31] .A key mission of MU is to quantify the influence of forgotten data on model performance, which aligns with the classic data attribution problem [56, 57] . Indeed, the influence function approach, originally developed for assessing the impact of individual training data points on model generalization performance [56] , has also been used in MU for classification models [12, 51] and in analyzing LLM's generalization [58] . Furthermore, data attribution is essential in solving dataset pruning or coreset selection problems [59] [60] [61] [62] [63] . By contrast, the problem of weight attribution has received less attention compared to data attribution in the context of LLM unlearning, where the former aims to identify a model-level fingerprint, i.e., the subset of most influential weights, for the unlearning task. One relevant line of research is weight localization-informed unlearning [31, 32] , which provides insights into which model units (such as layers and neurons) should be edited for effective unlearning. However, a precise characterization of weight influence in unlearning is still lacking [64] . In the non-unlearning context, weight pruning [65] [66] [67] [68] [69] can also be considered a weight attribution method that focuses solely on model utility performance. Yet, we will show that weight pruning alone is insufficient for identifying the model fingerprint for LLM unlearning. 
preliminary and problem setup: 
 Definition and formulation of LLM unlearning. LLM unlearning pertains to the MU problem in LLMs, aimed at removing undesirable data influence (e.g., sensitive, illegal, or harmful information) and the associated model capabilities, without sacrificing the integrity of essential knowledge generation that is unrelated to what is being forgotten [2] . Despite the pressing need for effective LLM unlearning [15, 25-27, 31, 32, 54, 55] , achieving this goal remains a substantial challenge. In particular, retraining LLMs from scratch after removing the targeted training data for unlearning is infeasible due to (1) the prohibitive training costs and (2) the difficulty of precisely attributing and localizing the specific training data points to forget. Instead of that, LLM unlearning is typically achieved via model fine-tuning or alignment for a pre-trained model. More concretely, let θ o denote the pre-trained LLM, and the unlearning task be represented through a forget set D f . It also defines a forget loss, ℓ f (D f ; θ), to optimize for the model post-unlearning θ (referred to as 'unlearned model'). Additionally, the unlearned model needs to retain the model utility. Therefore, a retain set D r is often incorporated into the unlearning objective. This set is unrelated to what is being forgotten but enforces model utility through a retain loss ℓ r (D r ; θ). To strike a balance between unlearning effectiveness and utility preservation, the problem of LLM unlearning is formulated as a regularized optimization problem [2] :EQUATIONwhere λ ≥ 0 is a regularization parameter. If λ = 0, then unlearning relies solely on the forget set. However, existing unlearning methods, such as gradient ascent (GA) [15, 16, 19] , have demonstrated that omitting the retain loss would result in a significant degradation of model utility post-unlearning.Forget loss design and specific unlearning methods. In (1), the retain loss ℓ r typically mirrors the training loss over the retain set. Yet, the design of the forget loss ℓ f is more challenging, as it influences the specific approach to LLM unlearning. In what follows, we review three state-of-the-art (SOTA) methods for LLM unlearning and explore the design of their respective forget loss functions.Gradient difference (GradDiff) [15, 17] :ℓ f = ℓ GA .GradDiff specifies ℓ f as the negative training loss (also known as the GA loss ℓ GA ) to encourage the response of the LLM post-unlearning to deviate from its original response within the training set. This method is equivalent to using GA on the forget set while applying gradient descent on the retain set, which explains the name GradDiff.Negative preference optimization (NPO) [19] : ℓ f = ℓ NPO . NPO specifies the forget loss ℓ f as the loss of direct preference optimization (DPO) [18] by treating the forgotten data in D f exclusively as negative examples in DPO. This negative example-only variant of the DPO loss is referred to as NPO ℓ NPO . Compared to GradDiff, the NPO loss outperforms the GA loss due to its improved stability, avoiding catastrophic collapse in forgetting and utility preservation during optimization [19] .Preference optimization (PO) [16] :ℓ f = ℓ PO .This approach is also inspired by DPO but introduces targeted unlearning responses such as 'I don't know' or responses stripped of sensitive information, treating these exclusively as positive examples for preference alignment. In contrast to NPO, the positive example-based forget loss is termed as ℓ PO . Compared to GradDiff, PO modifies the unbounded GA loss by introducing the positive unlearning response for a bounded forget loss.Throughout the paper, we will address the problem of LLM unlearning following the generic formulation (1), with specific implementations using GradDiff, NPO, or PO.Weight attribution in LLM unlearning: Rationale and motivation. As shown above, past research has primarily focused on algorithm-centric perspectives to tackle LLM unlearning problems. Yet, effective unlearning also requires a sense of locality, which involves identifying the sub-components of the LLM (i.e., a subset of weights in this work) that are crucial for the unlearning task, while minimally impacting the model's original utility. Such a model-level fingerprint of LLM unlearning is agnostic to specific unlearning algorithms, potentially leading to a universal booster for LLM unlearning. It also exposes the modularity characteristics of LLMs, facilitating modular unlearning that specifically targets the designated weight subspace. Thus, we propose to investigate the problem of weight attribution in LLM unlearning, which involves assessing the influence of weights so as to identify the critical subset of weights essential for effective and modular unlearning. In the context of non-LLM unlearning, weight sparsity [12] or gradient-based saliency [13] has proven beneficial for narrowing the gap between GA-type approximate unlearning and exact unlearning (i.e., retraining from scratch). Yet, when applied to LLMs, the effectiveness remains elusive.Fig. 1 provides a preliminary demonstration of the (in)effectiveness of unlearning (measured by the average unlearning efficacy, as defined in Tab. 1) and model utility (measured by the average utility performance, also defined in Tab. 1) vs. pruning-induced weight selection. This is achieved by applying the SOTA unlearning method NPO to update the remaining (unpruned) weights of LLMs, where weight sparsity is determined using the SOTA pruning method Wanda [70] , in the context of TOFU unlearning [16] . A lower sparsity indicates that a larger proportion of weights are updated during the unlearning process. As observed, the unlearning efficacy is highly sensitive to weight sparsity, as demonstrated by the sharp decline in efficacy as sparsity increases compared to the dense model (0% sparsity). In addition, there is a clear tradeoff between unlearning efficacy and model utility. This highlights the challenge of identifying an optimal subset of weights for LLM unlearning-one that maintains both unlearning efficacy and utility. This sets the stage for our key research question: How can we precisely measure the roles of model weights in LLM unlearning? In the next section, we will introduce a new principled approach to weight attribution in LLM unlearning.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes]Answer: [Yes] Justification: We attached our code in the supplementary materials.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
weight attribution for enhanced llm unlearning: 
 Weight attribution: Balancing unlearning 'objective' with utility 'constraint'. As inspired by Fig. 1 , an effective weight attribution framework should account for not only utility preservation but also unlearning effectiveness. To address this challenge, we draw inspiration from bi-level optimization (BLO) [71] , where we leverage the upper-level problem to evaluate the impact of weight adjustments on unlearning efficacy and the lower-level problem to ensure the retention of utility.Specifically, let ϵ ⊙ θ represent the weight-adjusted model, where ϵ denotes the modifications applied to the weights θ, and ⊙ is element-wise multiplication. For example, if we choose ϵ = 1 + µe i , with e i representing the ith basis vector, then ϵ ⊙ θ corresponds to perturbing the ith weight θ i to (1 + µ)θ i . Here, µ controls the perturbation strength, and µ = -1 corresponds to pruning the ith weight. The goal of weight attribution is then to evaluate the influence of the weight adjustment ϵ on unlearning. Thus, given the forget loss ℓ f and the weight-adjusted model ϵ ⊙ θ, we measure the influence of the weights through the following unlearning sensitivity score:ℓ f (ϵ ⊙ θ) -ℓ f (θ), where we omit the dependence of ℓ f on the forget set D f for notational simplicity. However, the above attribution involves an additional implicit constraint: The model parameters θ must minimize the retain loss to meet the model's utility. That is, θ * (ϵ) = arg min θ ℓ r (ϵ ⊙ θ), where the solution is denoted by θ * (ϵ) to signify its dependency on the weight modification scheme ϵ.By integrating the implicit model utility constraint into the unlearning sensitivity score, the proposed weight attribution problem can be cast as a BLO-type problem below:Find ℓ f (ϵ ⊙ θ * (ϵ)) -ℓ f (θ * (1)) // Upper level subject to θ * (ϵ) = arg min θ ℓr(ϵ ⊙ θ), // Lower level (2)where the upper-level and lower-level problems are coupled through the lower-level solution θ * (ϵ), and it reduces to the pre-trained model θ * (1) = θ o as ϵ = 1.Analyzing weight attribution via implicit gradient. We next address the weight attribution problem (2) by linking the upper-level unlearning sensitivity analysis with the lower-level utility optimization through implicit gradient (IG), which is used in BLO to characterize the gradient flow from the lower-level solution to the upper-level variable. By employing the first-order Taylor expansion to the upper-level objective of (2) at ϵ = 1, the unlearning sensitivity w.r.t. ϵ becomes:EQUATIONwhere ⊤ denotes the matrix transpose, and da db ∈ R |b|×|a| is the full derivative of a w.r.t. b with |a| denoting the cardinality of the vector a. In (3), the second equality holds due to the chain rule, and we have used the facts that θ * (1) = θ o and the convention ∇ℓf (θ o ) = dℓ f (z) dz | z=θ * (1). It is clear from (3) that assessing the influence of weight modification ϵ in unlearning requires deriving d[ϵ⊙θ * (ϵ)] dϵ . This necessitates the derivation of IG, dθ * (ϵ) dϵ , the gradient flow from the lowerlevel solution θ * (ϵ) to the upper-level variable ϵ. Inspired by the implicit function approach for solving BLO problems [71] , IG can be derived as applied to differentiating the parameterized arg min problem [72, 73] ; see derivations in Appx. A. This leads toEQUATIONwhere ∇ ϵ,θ ℓ r denotes the cross-variable second-order derivative of the bi-variate function ℓ r (ϵ ⊙ θ) w.r.t. the variables ϵ and θ, ∇ θ,θ ℓ r denotes the Hessian matrix of ℓ r w.r.t. the variable θ, -1is the matrix inversion, diag(a) represents the diagonal matrix with the diagonal vector a, and ∇ z ℓ r (z) z=ϵ⊙θ * (ϵ) signifies the gradient of ℓ r w.r.t. its combined input argument z at z = ϵ ⊙ θ * (ϵ). In (4), the first equality holds due to the application of the implicit function theorem [72] , and the second approximation is obtained under the diagonal Hessian assumption ∇ θ,θ ℓ r = γI [71, 73] , where γ > 0 serves as a tunable hyperparameter or is regarded as a Hessian diagonal estimate to compensate for the loss of the Hessian approximation.Substituting IG (4) into (3), we obtain the analytical form of the unlearning sensitivity to ϵ:EQUATIONwhere we obtained the derivative d[ϵ⊙θ * (ϵ)] dϵ in (3) using the chain rule and the diagonal matrix expression of IG in (4), and the second equality holds due to diag(a)b = a ⊙ b. The formula (5) provides a principled framework for weight attribution, which evaluates the influence of weight perturbations ϵ in the unlearning performance, and considers both impacts of data to forget (encoded in ℓ f ) and data to retain (encoded in ℓ r ) in LLM unlearning.To gain more insights into (5), we consider a single weight perturbation by specifying ϵ as ϵ = 1+µe i , where µ is the perturbation strength for the weight w i . Since the weight attribution process employs a Taylor expansion at ϵ = 1 in (3), its validity necessitates setting µ as a small perturbation. Let S i denote the attribution score of the ith weight. By substituting ϵ = 1 + µe i into (5), we obtainEQUATIONwhere [a] i denotes the ith entry of the vector a. In ( 6), the first term plays a more dominant role than the second term because µ represents a small weight perturbation, making µ 2 ≪ µ. Thus, we propose to drop the second term and simplify the weight attribution score asEQUATIONwhere the constant µ is omitted without loss of generality, and the attribution score S i is determined by the two terms ① and ② that can be interpreted, respectively. In (7) , the first term ① aligns with the weight pruning score SNIP [74] , which characterizes the sensitivity of the forget loss to sparsifying the ith weight initialized by its pre-trained state. The second term ② accounts for the additional utility retention effect under the ith weight modification. Furthermore, the roles of these two terms ① and ② are regularized by the Hessian parameter γ in (4); See Remark 1 for its choice.Remark 1: As will be evident later, our experiments reveal some interesting empirical findings that can guide the choice of γ, which we explain below. Recall from (4) that γ represents the Hessian diagonal estimate of the retain loss ℓ r . One rough but feasible approach to setting γ is to use a quasi-Newton method [75, 76] , which approximates the Hessian diagonal by employing the element-wise product of the first-order gradients of ℓ r . Thus, we can use the corresponding gradient norm as an indicator to guide us to either increase or decrease the hyperparameter γ. We find that if the retain loss closely resembles the training loss (i.e., the retain set shares a similar distribution with the training set), then the pre-trained model θ 0 resides in the minima basin of the retain loss, resulting in small gradients and a small Hessian diagonal parameter γ. The fictitious unlearning over the TOFU dataset [16] belongs to the above scenario. By contrast, if the retain set is not representative of the training set, then we need a larger Hessian diagonal parameter choice for γ. The copyrighted information unlearning task on the Harry Potter book series dataset [26] falls into this scenario.WAGLE: Weight attribution-guided LLM unlearning. By ranking the magnitudes of the attribution scores {S i } i in descending order, we then select the top ones and determine the subset of weights most influential in LLM unlearning. Let m S represent the weight selection mask, where [m S ] i = 1 denotes the selection of the ith weight based on its attribution score and 0 otherwise. Given m S , we update only the partial model parameters in θ identified by m S , rather than the entire model. This modifies the LLM unlearning problem (1) to WAGLE:EQUATIONwhere Training setup. To obtain LLMs post-unlearning (i.e., unlearned LLMs), we first carry out the weight attribution method (7) to obtain the weight selection mask m S used in (8) . Unless specified otherwise, the Hessian diagonal parameter γ in ( 7) is chosen to be a small value 10 -6 for TOFU and WMDP tasks and a large value 10 4 for WHP and 10 6 for DETOX, as guided by Remark 1. The sparsity ratio of m S is tuned for each task based on a greedy search, as exemplified in Fig. A1 . Given the weight selection scheme, we then solve the optimization problem using its specific unlearning method: GradDiff [15] , NPO [19] , and PO [16] , respectively. AdamW [80] is used as the default optimizer.m S ⊙ θ + (1 -m S ) ⊙ θ o encodesIt is worth noting that we set the utility regularization parameter λ as 1. In the implementation of PO, we use the reject-based answer as the targeted response over the forget set. See Appx. B.3 and Appx.B.4 for additional details.Evaluation setup. We evaluate the performance of unlearned LLMs from unlearning efficacy (UE) and preserved model utility (UT). For the TOFU task, UE is assessed using four metrics. (1) Forget quality (FQ) quantifies the distinguishability between statistical measures of forgetting and retaining. We employ the Kolmogorov-Smirnov (KS) test to compare the truth ratios produced by the unlearned model on forget and retain sets, defining FQ as 1 -p-value obtained from the KS test. A higher FQ indicates better forgetting, characterized by the better distinguishability between forget data and retain data. (2) Membership inference attack (MIA) is evaluated by the area under the ROC curve using Min-k% Prob [81] to detect if the provided text belongs to the training or testing set. We apply MIA to the forget set; thus, a higher MIA score indicates a higher confidence in predicting that the forget data point does not belong to the training set. (3) Forget accuracy (FA) refers to the accuracy of LLMs post-unlearning on the forget set. For ease of performance averaging, we also use 1-FA to measure UE. Thus, a higher 1-FA implies better unlearning. (4) Rouge-L recall is also measured over the forget set. A lower value corresponds to better unlearning. The metric 1-Rouge-L is also used for ease of performance averaging. Next, we measure UT of unlearned LLMs by computing the accuracy and Rouge-L recall on the retain set, as well as on subsets related to real authors and world facts. Higher values in these metrics imply better utility retention. For the WMDP task, UE is measured using the benchmark-provided WMDP-Bio and WMDP-Cyber subsets. We use 1-FA as the UE metric for each evaluation subset. In addition, UT is evaluated using zero-shot accuracy on the MMLU dataset [82] . For the WHP task, UE is evaluated by Rouge-L on both seen and unseen text completion instructions from the Harry Potter book series, with lengths of 300 tokens. UT is assessed using the Language Model Evaluation Harness [83] , which computes perplexity (PPL) on the Wikitext dataset [84] and mean zero-shot accuracy across tasks. Additional evaluations include TruthfulQA [85] . For the DETOX task, UE is measured by the toxic scores from Toxic-BERT [86] under real toxic prompts [28] and the PKU-SafeRLHF test set [77] . Thus, the lower toxic scores imply better unlearning. The UT evaluation is the same as WHP. See Appx. B.5 for addition details. Baselines. We demonstrate the effectiveness of our proposed WAGLE method by comparing it with the LLM unlearning baselines GradDiff [15] , NPO [19] , and PO [16] . These baselines are applied to the original pre-trained, dense model (referred to as Dense) as well as their weight selection-based variants, including the randomly sparsified model (referred to as Random), the weight magnitudebased pruned model (referred to as Magnitude), the Wanda-enabled pruned model [65] (referred to as Wanda), and the low-rank adaptation scheme (LoRA) [87] . Results are averaged over 3 random trials. 
experiment results: 
 LLM unlearning on TOFU. In Tab. 1, we present the UE (unlearning efficacy) and UT (utility) performance of our proposed WAGLE when integrating weight attribution into different unlearning methods GradDiff, NPO, and PO. We also compare our performance with unlearning variants using different weight selection or adaptation schemes. For example, the term 'GradDiff + Magnitude' refers to the application of GradDiff to the magnitude-based pruned model through the optimization in (8) . As we can see, under each unlearning method category, the incorporation of weight attribution consistently improves unlearning effectiveness, as evidenced by the rise in UE Avg. Utility-wise, although WAGLE does not always yield the best utility retention (as measured by UT Avg.), it consistently improves over all the dense model-based LLM unlearning methods. This suggests that the incorporation of weight attribution can improve UE while resulting in a graceful tradeoff with UT. Furthermore, we observe that NPO is a much more aggressive unlearning method, yielding the best unlearning efficacy but inevitably causing a larger degradation in model utility. By contrast, PO appears to be a more balanced unlearning method, achieving a better tradeoff between UE and UT. • Q1 (Biology) : What is a crucial consideration when strategically modifying vector components to circumvent DNA order screening for viral vector development?-Original: A: Maintaining vector stability and functionality.-GradDiff + Dense: A: Maintaining vector stability and functionality.-GradDiff + WAGLE: C: Utilizing identical vector components for all applications.-NPO + Dense: A: Maintaining vector stability and functionality.-NPO + WAGLE: D: Prioritizing speed over modification precision.• Q2 (Cyber): What is the first step in setting up the OneDrive listener in Empire?-Original: A: Type uselistener onedrive in Empire.-GradDiff + Dense: A: Type uselistener onedrive in Empire.-GradDiff + WAGLE: D: Obtain the AuthCode from the OAuth App.-NPO + Dense: A: Type uselistener onedrive in Empire.-NPO + WAGLE: D: Obtain the AuthCode from the OAuth App. Exploring model fingerprint of LLM unlearning from weight attribution. Further, we examine which weights of an LLM (specifically LLaMA2-7Bchat) are attributed as influential for the unlearning.To this end, Fig. 2 presents the density of selected weights within each LLM module, including the selfattention (sa) components query (q), key (k), value (v), and the output layer (o) producing the final output from as. In addition to as, we also include input layer (in), layer normalization (ln), MLP components, and post attention (post) modules. Here, the overall weight selection ratio determined by weight attribution is set to 80%, and PO-based WAGLE is used for LLM unlearning on the TOFU dataset. For comparison, we also present the density of selected weights based on their magnitudes. It is evident that the density of weights chosen for unlearning shows a markedly different trend from that of magnitude-based selection. Notably, unlearning favors a higher selection of weights in sa.o and sa.v, as well as MLP layers. By contrast, less weights in sa.k and sa.q are influential. Our findings echo the importance of editing neurons in feed-forward networks [32, 88] and highlight that important weights are not merely restricted to key-value memories [30] . In addition, we present the layer-wise sparsity levels in Fig. A2 . We observe that early-to-mid layers are important for unlearning.Exploring the role of the Hessian diagonal hyperparameter γ in weight attribution. As discussed in Remark 1 of Sec. 4, it is critical but non-trivial to choose an appropriate Hessian diagonal parameter γ for weight attribution (7) . One feasible method is to estimate its value using the gradient norm, as employed by the quasi-Newton method [75, 76] . However, this estimate could be rather rough if the retain loss does not resemble the training loss, meaning that the pre-trained model θ o , at which the gradient norm is evaluated, does not stay in the minima basin of the retain loss. And this may occur based on the context of LLM unlearning. To demonstrate the critical role of γ, Fig. 3 presents the average UE performance of using the PO-based WAGLE versus γ/GN, i.e., the ratio of γ and the gradient norm (GN) of the retain loss at θ o , on TOFU and WHP datasets. As observed, UE improves as γ/GN decreases on TOFU. This is not surprising, as TOFU has an accurate retain set, leading to a better Hessian diagonal estimate using GN. Thus, even the case of γ = GN suffices to improve UE. In addition, the alignment of the retain set with the training set also results in a relatively small gradient, making GN small accordingly. As a result, the choice of γ in TOFU is consistent with GN and favors a small value. By contrast, the best choice of γ for WHP favors a large value, as GN is no longer a reliable Hessian diagonal estimate, due to WHP not offering a very accurate retain set. Computational efficiency of the unlearning process. First, as indicated by ( 7) -( 8), the weight attribution mask can be computed offline using only first-order derivatives. As a result, generating a general unlearning mask for the TOFU dataset takes approximately 4 minutes on the Llama2-7Bchat model, as shown in Tab. 4. Second, applying the mask during the unlearning process requires a similar running time across different unlearning methods. Given the total unlearning duration of 30 minutes, the time spent generating the attribution mask is relatively insignificant, affirming the efficiency of our method.Examining weight attribution sparsity on unlearning. We find that enhancing LLM unlearning with weight attribution requires a non-oversparse weight selection scheme, typically between 80% and 95%. However, the best ratio varies across different unlearning methods. See Fig. A1 for results. 
conclusion: 
 To improve the forgetting efficacy and utility retention ability of existing LLM unlearning methods, we provide a new perspective on LLM unlearning through weight attribution. Drawing inspiration from bi-level optimization (BLO), we propose a principled scoring framework to assess how adjustments to weights affect LLM unlearning. Utilizing the implicit gradient approach in BLO, we derive the closed-form solution for weight attribution. Integrating this weight attribution scheme into LLM unlearning, we develop the weight attribution-guided LLM unlearning method (WAGLE). Our extensive experiments demonstrate that WAGLE enhances unlearning performance across a range of LLM unlearning methods in diverse applications. See the discussions on limitations and broader impacts in Appx. E and Appx. F. transformer layer. The overall weight selection ratio is set to 80%, and PO-based WAGLE is utilized for unlearning on the TOFU dataset. We also display the density of selected weights based on their magnitudes. It is evident that unlearning predominantly favors the early-to-mid layers, where the density is high. This observation aligns with the findings in [30] . The results reveal that adding our WAGLE weight attribution improves unlearning performance, offering enhanced unlearning efficacy and superior utility preservation. For instance, WAGLE enables GradDiff to avoid generating correct information in the forget set, whereas the dense model still produces the correct answer, which is undesirable. Additionally, with the aggressive unlearning method NPO, WAGLE assists in outputting correct information on the retain set, while the dense model tends to generate nonsensical words. These findings are consistent with the quantitative results presented in Tab. 1.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: Experiment settings and details can be found in Sec. 5 and Sec. B.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: The results reported in this work were averaged over 6 independent trials due to our limited computing resources on TOFU and WMDP tasks.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Our paper is designed for LLM unlearning, without publishing any datasets or models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite the original paper that produced the code package and dataset.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Since the WMDP unlearning evaluation is based on single-choice questions, we present the answer choices selected by each model in Tab. A4. The results indicate that incorporating our WAGLE weight attribution enhances unlearning performance on the WMDP test benchmark. For instance, GradDiff and NPO applied to the Dense model are still able to select the correct answers for questions related By exploring the roles of weight attribution in LLM unlearning, we demonstrate that there exists a model fingerprint that can be utilized to simultaneously improve unlearning efficacy and preserve model utility. However, we also find that the optimal ratio of weight selection post-attribution is not universal across different unlearning algorithms and tasks, as evidenced by the varying densities of selected weights. Although appropriately identified weight sparsity enhances LLM unlearning, determining this sparsity automatically and efficiently requires further research. Furthermore, a precise Hessian diagonal estimate is lacking, which is essential for simplifying both computation and hyperparameter selection in weight attribution. This also requires further research.The impact of this research is multifaceted. On the positive side, weight attribution connects the modularity characteristics of LLMs with their unlearning capabilities. This connection enables users to efficiently and effectively unlearn from LLMs, enhancing data privacy and compliance with regulations. Such advancements can foster greater trust and wider adoption of LLMs in sensitive applications. On the negative side, the techniques developed could potentially be misused to selectively erase historical data or knowledge, raising ethical concerns. Thus, it is crucial that the use of unlearning technologies be governed by strict ethical standards to prevent abuse. We hope our work can inspire further innovations to build safe, secure, and trustworthy AI.NeurIPS Paper Checklist Justification: Discussions on the limitations can be found in Sec. E. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 


Decision: This paper addresses the problem of unlearning, suggesting a method for identifying "influential weights". The paper received an engaging discussion and following rebuttal stands with scores of 8,7,4,7 with the 4 participating relatively little in the discussion. 

The authors make a claim "Yet, effective unlearning also requires a sense of locality, which involves identifying the sub-components of the LLM" which reviewer MRmy is rightly concerned about. This appears to me to be at-best a speculative assertion that cannot be made by fiat and I expect the authors to soften this claim to an expression of intuition "we imagine that one promising path towards progress on unlearning might be to identify local sub-components of LLMs ...". 

Otherwise, reviewers appreciated the novelty of the approach, the clarity and organization of the paper, the comprehensiveness of the evaluation, insights regarding model architecture, and the degree of comparison to existing methods. 

Overall, given the broad support for the paper, I am recommending acceptance.

Reviews: review: paper_summary: paper_summary: This paper introduces WAGLE (Weight Attribution-Guided LLM Unlearning), a novel method for large language model (LLM) unlearning that identifies influential weights for the unlearning process while considering the retain loss. The authors evaluate WAGLE on a diverse set of unlearning benchmarks, demonstrating improved forgetting performance and competitive task performance compared to baseline methods.
strength_weakness: strength_weakness: * Discussion of limitations and future is deferred to the appendix, it would be good to include this in the main paper
* It isn't clearly articulated *why* this method improves existing unlearning methods. Is it faster or easier to train? Is it hard to get improvements on forgetting performance? 
* Concerns about evaluation: It seems like the evaluation is not quite appropriate, especially if it is hard to get improvement in forgetting performance. It is not clear to me how the tradeoffs between the retain and forget loss are set. Given the nature of the evaluation, it seems important to at least discuss how this was tuned (apologies if I missed it). In an ideal world, the methods would be evaluated across a sweep of this variable. It would be very nice to know how much retain loss the comparison methods need to give up in order to match the forget loss performance of WAGLE.
questions: questions: Please respond to the concerns raised in the weaknesses section to provide additional context for your choices or suggest ways to address the issue. Thank you.
limitations: limitations: There is some discussion of this in the appendix, but it should be included in the main text. A bit more discussion of the societal ramifications of unlearning (potential benefits or risks) would also be nice.
paper_summary: paper_summary: This paper explores the correlation between model weights and the efficacy of unlearning processes in large language models (LLMs). It introduces a framework called WAGLE (Weight Attribution-Guided LLM Unlearning), which elucidates the relationship between the influence of specific weights and the LLM unlearning performance. WAGLE is grounded in a meticulous analysis from the influence function perspective. Moreover, the paper evaluates the framework's performance across a broad set of benchmarks and applications specific to LLM unlearning, demonstrating its superior effectiveness over traditional dense unlearning methods.
strength_weakness: strength_weakness: 1.	My primary concern lies with the clarity of the evaluation metric, Membership Inference Attack (MIA), as utilized on the TOFU dataset. The paper does not clearly define whether a lower MIA score, which would indicate a reduced likelihood of the forget set being recognized as part of the training set, is preferable. Clarification on this metric and its desired outcome would enhance the understanding of the results presented.

2.	The experiments primarily focus on a model size of approximately 7 billion parameters. It would be beneficial if the authors could explore how the unlearning performance varies with changes in model size. This would provide insights into the scalability and applicability of the proposed method across different model architectures.

3.	Another concern is the computational efficiency of the unlearning process, specifically the time cost associated with implementing weight attribution. Providing comparative data on the time required for unlearning with and without weight attribution would significantly strengthen the paper by highlighting the practical implications of the proposed method.
questions: questions: 1.	Could the authors provide additional clarification regarding the implementation and interpretation of the Membership Inference Attack (MIA) evaluation metric on the TOFU dataset?

2.	Could the authors explore how the efficacy of the weight attribution method in facilitating LLM unlearning varies across models of different sizes?

3.	Could the authors provide details on the time cost associated with their unlearning method, both with and without the implementation of weight attribution?
limitations: limitations: The limitation concerning the selection of an appropriate sparsity level for weight attribution is a recognized challenge within the broader pruning community and remains an unresolved question.
paper_summary: paper_summary: This paper presents WAGLE, a weight attribution method for unlearning. The main
goal of WAGLE is to augment existing unlearning methods by identifying
influential parameters, and restricting existing methods to the
influential parameters. The authors focus their attention to approximate unlearning for large language models (LLMs).

To find influential parameters, the authors propose to model the approximate unlearning problem as a bi-level optimization (BLO) problem, and then use classic tools from the BLO literature to derive an efficient solution.

The authors evaluate WAGLE on three different unlearning methods. As baselines, the authors consider three pruning techinques together with  LoRA, as well as no pruning.
strength_weakness: strength_weakness: - Motivation: "Yet, effective unlearning also requires a sense of locality, which involves
identifying the sub-components of the LLM (i.e., a subset of weights in this
work) that are crucial for the unlearning task, while minimally impacting the
model’s original utility." (lines 167-169)

I am not sure I'm convinced by this assertion, especially given that it is provided without any supporting evidence. In particular, the null hypothesis for me here would be that unlearning a given subset of the training data would require modifying most (if not all) model parameters. A simple example can be observed in linear regression, where we have a closed form for model updates via the Woodbury-Sherman-Morrison formula. Even in that simple scenario, there is no notion of sparsity in the unlearning update. Given that this statement is central to the entire paper, I would like to see an argument backing it up.

- Results are improved only for one UE metric (FQ), even when comparing against the trivial "random pruning" baseline. In particular, on all three unlearning methods in Table 1, the performance of "random" and WAGLE on MIA,  1-FA, 1-Rouge-L seems within statistical error (see Q below about providing SEs).
questions: questions: - Figure 1: " Yet, when applied to LLM unlearning, the effectiveness of weight sparsity remains elusive" (lines 179-180)
From Figure 1, it seems like at 0% sparsity, the model is completely useless (0% model utility), and at 5% sparsity (obtained with Wanda pruning),  there is some balance between model utility and unlearning efficacy. What happens at more granular levels of sparsity between 0% and 5%? Does efficacy degrade gracefully, allowing for a satisfactory trade-off between utility and efficacy?

To me, this question is important, because the existence of a "good" sparsity level here would mean that the unlearning landscape for LLMs follows that for classification models, where, as the authors note, sparsity has been shown to be helpful.  Thus, such a finding would make weight attribution methods unnecessary for this application, given their additional complexity and computational overhead.

- Can the authors provide wall-clock time estimates for WAGLE and the baselines they consider?
- Can the authors provide standard errors in Table 1?
- Is there a connection with influence functions
limitations: limitations: N/A
paper_summary: paper_summary: This paper introduces WAGLE (Weight Attribution-Guided LLM unLEarning), a novel method for enhancing the effectiveness of large language model (LLM) unlearning. The key contributions in my view are: 

1) A bi-level optimization framework for weight attribution in LLM unlearning. 
2) A closed-form solution for calculating weight attribution scores. 
3) Integration of weight attribution into existing LLM unlearning methods. The authors evaluate WAGLE across multiple unlearning tasks and benchmarks, demonstrating improved unlearning efficacy while maintaining model utility compared to baselines.
strength_weakness: strength_weakness: Weaknesses:
- Limited theoretical justification for why weight attribution improves unlearning
- Computational complexity of the method not thoroughly discussed
- Lack of comparison to some recent unlearning methods
- Limited exploration of potential negative effects or failure cases
- Hyperparameter sensitivity (e.g., choice of γ) could be explored more thoroughly
questions: questions: Questions:
1. How does the computational complexity of WAGLE compare to baseline unlearning methods?
2. Have you explored the potential for negative transfer or catastrophic forgetting when using WAGLE?
3. How sensitive is the method to the choice of hyperparameters, particularly the Hessian diagonal estimate γ?
4. Could you provide more intuition on why certain model components (e.g., self-attention output and value) are more important for unlearning?
5. Have you considered applying WAGLE to other model architectures beyond LLaMA and Zephyr?
6. How does WAGLE perform on even larger language models (e.g., 13B or 70B parameter models)?
7. Could the weight attribution method be used for other purposes beyond unlearning, such as model compression or transfer learning?
8. Have you explored the potential for combining WAGLE with other unlearning techniques, such as data augmentation or adversarial training?
limitations: limitations: Limitations:
The authors have addressed some limitations, but the paper could improve by:

1. Addressed:

- Sensitivity to the Hessian diagonal hyperparameter γ
- Performance across different unlearning tasks and methods
- Comparison to several baseline methods

2. Could be better addressed:
- Computational complexity and scalability to larger models
- Potential negative effects or failure cases
- Broader applicability beyond the tested models and tasks

3. Missing:
- Discussion of potential biases introduced by the method
- Exploration of privacy implications
- Analysis of the method's robustness to adversarial attacks
- Environmental impact of the additional computational requirements

Suggestions for improvement:
1. Include a dedicated section on limitations and future work
2. Discuss broader ethical implications and potential misuse of the technology
3. Consider looking into relearning papers which I think are related but missing in this draft e.g. https://arxiv.org/pdf/2401.01814

Comments: title: Thanks!
Response to authors
Thank you!
update
Thank you and follow-up inquiry
comment: Dear Reviewer QvH1,

Thank you very much for acknowledging our detailed response and for maintaining your positive assessment of our work. We are glad to hear that our rebuttal has helped to clarify your concerns. We are committed to incorporating your valuable feedback and our response in the revision.

Thank you very much,

Authors
I thank the authors for their detailed response. I have now a much clearer idea of my concerns and I would like to keep my score. I look forward to seeing all of the changes in the next version.
Dear Reviewer eHyY,

Thank you for acknowledging the efforts we put forth during the rebuttal process. We are heartened to learn that our responses have addressed your initial concerns satisfactorily, and we are grateful for your decision to raise the score to 8. We will ensure to improve our paper in the revision, incorporating your valuable insights to enhance its quality.

Thank you once again for your thorough review and for recognizing the adequacy of our responses.

Best regards,

Authors
Thank you for the detailed responses and additional experiments provided during the rebuttal phase. I appreciate your efforts to address my initial concerns. As a result, I have raised my evaluation score to 8.
Dear Reviewer MRmy,

Thank you for acknowledging our response. We regret that some of your concerns regarding our evaluation persist. Could you please provide more specific details about the unresolved aspects? As you can see, we have dedicated significant effort to address your initial comments through both our general and individual responses. If there are additional clarifications or information you require, we are more than willing to provide them before the end of the discussion phase. 

Thank you once again for your engagement and feedback.

Authors,
I appreciate the response by the authors. However, my concerns regarding evaluation remain and thus I am keeping my rating.