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

Thank you for the constructive feedback. Please find our response to each weakness (W), question (Q), and limitation (L) below. References (in the format of [Rx]) can be found in the general response.

Response to W1: Our work, while lacking formal theoretical guarantees on how weight attribution enhances unlearning, is rooted in a rigorous bi-level problem formulation (Eq. (2)), leading to our final score derivation (Eq. (8)). This has deepened our understanding of LLM weight impacts on unlearning, especially in maintaining utility. For motivation and benefits of weight attribution in improving the unlearning efficacy-utility tradeoff, please see GR1. While a stronger theoretical guarantee could improve our work, this limitation does not undermine our significant technical contributions. We will discuss this point in the Limitations section.

Response to W2&Q1: Please refer to GR2 in the general response.

Response to W3: We appreciate your comment on including recent unlearning methods. At submission, NPO (published April 8, 2024) was the most advanced, leading in the TOFU and the very recent MUSE benchmarks (after submission) [R10]. In our revision, we will survey and include any new SOTA methods published post-submission.

Response to W4: In response to your suggestion, we've considered potential drawbacks of our method, particularly the reliance of WAGLE on the Hessian diagonal estimate, $\gamma$, in Eq. (7). This reliance is significant as more accurate methods like Woodfisher [R11] struggle with scalability in larger models like Llama2-7B. Consequently, an imprecise estimate of $\gamma$ could degrade WAGLE’s performance, as noted in Remark 1 (Line 254).

Response to W5&Q3: We appreciate the reviewer's feedback on the sensitivity of $\gamma$. In Remark 1 (Line 254), we detailed the examination of $\gamma$'s impact, offering guidance on its selection based on the unlearning dataset type. Our experiments (Lines 402-423) further explored $\gamma$'s sensitivity, determining its choice from insights in Remark 1. We found that a smaller $\gamma$ suits scenarios where the retain set closely mirrors the training distribution, indicated by smaller gradient norms on the retain set. Conversely, a larger $\gamma$ is advisable when there's a significant disparity between the retain and training sets, necessitating more substantial weight adjustments.

Response to Q2: We've explored catastrophic forgetting in the context of fictitious unlearning using the TOFU dataset. Specifically, we used metrics in Table 1 to assess LLM responses on general World Facts, a task distinct from TOFU's primary focus on unlearning fictitious authors' profiles. Our results show that WAGLE not only matches or exceeds other methods in unlearning efficacy but also maintains utility in unrelated tasks. For instance, with the PO method, WAGLE improves unlearning efficacy (UE Avg.) while preserving performance on World Facts tasks (see Acc. on World Facts). This demonstrates WAGLE's ability to selectively unlearn specific information without significant catastrophic forgetting.

Response to Q4: Our explanation is that changes in q and k require the model to relearn the attention distribution and those modules primarily capture input patterns [R12], changing on those modules will hurt utility a lot. In contrast, self-attention output (o) and value (v) modules directly influence the model's representations and outputs, making them more effective for unlearning. Additionally, editing neurons in MLP layers has also been shown to be effective for modifying model behavior [R13].

Response to Q5&Q6: Our decision to use the 7B model size and LLaMA/Zephyr was driven by two factors: the need for direct comparison with established unlearning benchmarks like TOFU or WMDP (as referenced in [R3, R7], which used 7B model), and the computational resources available to us, particularly constrained by the GPU memory limits of our RTX A6000-equipped platform.

Response to Q7&Q8: Thank you for these suggestions. We response them as follows:

• Model compression and transfer Learning: Our weight attribution framework, outlined in Eq. (2), is tailored for unlearning but also applicable to model compression and transfer learning. By adjusting the upper-level problem in our bi-level formulation, we can optimize weights for transfer learning or enhance model compression. We've linked our weight attribution scores to SNIP-based pruning [R14] in Lines 248-253.

• Combining with other unlearning techniques: At submission, we were unaware of any papers combining LLM unlearning with data augmentation or adversarial training. While adversarial training [R15] is computationally intensive for LLMs, applying data transformations to forget data could enhance unlearning robustness. We consider these combinations a promising avenue for future research and will discuss these further in our future work section.

Response to L3.1&3.2&3.3&3.4: Thank you for emphasizing these key points. We will update our Limitations section accordingly:

1. Our method may introduce bias in Hessian estimation, managed by a scalar hyperparameter.

2. Approximate unlearning could compromise privacy, echoing risks associated with fine-tuning, as referenced in [R16] and [R17].

3. Adversarial robustness in LLM unlearning is under-explored. Our research, which aims to improve unlearning efficacy while maintaining utility, acknowledges vulnerabilities similar to post-fine-tuning risks [R16], identifying this as a future research area.

4. The environmental impact of unlearning is minimal in terms of additional computational demands, as discussed in responses to Q5 & Q6.

Response to L4&L5: We cover limitations, ethical implications, etc., in Appendices E and F. In the revision, these will be integrated into the main paper for a more thorough discussion. The recommended paper will also be added to the related work section.

Thank you for your insightful feedback. Below, we provide detailed responses to each weakness (W) or question (Q). References (in the format of [Rx]) can be found in the general response.

Response to W1: Thank you for your comments on the importance of locality in effective unlearning. Here are additional clarifications:

First, the concept of "effective unlearning requires a sense of locality" has been empirically supported by existing methods discussed in Lines 60-69. In particular, in the context of GA-based unlearning for discriminative classifiers, it has been theoretically shown that imposing model sparsity can significantly reduce the performance gap between approximate unlearning (which is computationally efficient but not verified) and exact unlearning (i.e., retraining from scratch) [R1]. This prior evidence underscores the importance of locality in enhancing the efficacy of approximate unlearning.

Second, the linear regression example proposed by the reviewer is insightful. As we explained above, locality is suitable for improving approximate unlearning. Yet, it may not be necessary for exact unlearning if you can readily achieve. This aligns with the reviewer’s argument that a closed form for model updates via the Woodbury-Sherman-Morrison formula is sufficient for unlearning in linear regression.

Third, we would also like to refer the reviewer to the GR1 for further clarification on the non-trivialness of finding proper locality by weight attribution and the associated benefit in improving the unlearning efficacy-utility preservation tradeoff for existing unlearning methods.

In the revision, we will make the above points clearer.

Response to W2&Q3: Thank you for your inquiry regarding our evaluation metrics and the statistical significance of our results.

First, the best assessment of unlearning efficacy (UE) presents challenges in the literature, prompting us to utilize various UE metrics to ensure a comprehensive evaluation. It is important to note that different unlearning methods demonstrate different sensitivities to these metrics. For instance, we agree with the reviewer that the FQ metric is more distinguishable, particularly under the GradDiff and PO methods detailed in Table 1. For the NPO method, however, metrics like 1-FA and 1-Rouge-L are more responsive, where our method also consistently outperforms other baselines.

Additionally, the compatibility between unlearning objectives and attribution methods varies significantly. For instance, the Wanda pruning method exhibits substantially lower unlearning efficacy (UE) across all metrics when using gradient ascent-type unlearning objectives in GA and its extended version, NPO. In contrast, our method demonstrates robust UE performance regardless of the unlearning objectives applied. To facilitate easier comparison and provide a general ranking of the different methods, we calculated an average of these UE metrics (UE Avg.), where WAGLE consistently emerges as the top performer.

Lastly, to address questions about statistical significance, we've added standard errors in Table R1 of the attached PDF, confirming that WAGLE's improvements over the 'random' baseline are significant. For instance, WAGLE improves UE Avg. by 0.03 with a standard error of 0.006, and maintains or enhances UT Avg. The GradDiff method shows a 0.048 improvement in UE Avg. with a standard error of 0.01, and NPO method increases utility by 0.1715 with standard error of 0.01, while maintaining similar unlearning efficacy as the random baseline. Additionally, under multiple runs, the random baseline exhibits a larger variance overall. For instance, 'random' shows a larger variance in UE Avg. (0.0286) compared to our method (0.0105) on GradDiff method, and a larger variance in UT Avg. (0.0124) compared to our method (0.0012).

We will clarify the above points in the revision.

Response to Q1: In response to your query about the effects of finer levels of sparsity between 0% and 5%, we conducted additional experiments using Wanda pruning across these granular levels on the TOFU dataset, as depicted in Figure 1. The detailed results are shown in Table R3 of the attached PDF.

The finer-level results revealed that the interaction between Wanda pruning and the NPO unlearning method exhibits a particularly sensitive trade-off between unlearning efficacy (UE Ave.) and utility (UT Avg.). For instance, at a 1% sparsity level with Wanda + NPO, the utility scores slightly higher at 0.6140 compared to 0.5936 for WAGLE + NPO. However, the unlearning efficacy for Wanda + NPO drops significantly to 0.8377, contrasting sharply with 0.9641 for WAGLE + NPO in Table 1. This underscores that while conventional LLM pruning might marginally improve utility, it does not enhance unlearning efficacy, reflecting its poor suitability for unlearning tasks where weight attribution should balance unlearning ‘objective’ with utility ‘constraint’.

Response to Q2: Please refer to GR2 in the general response.

Response to Q4: Thank you for your insightful question. Indeed, there is a connection between influence functions and our approach. In the context of unlearning, influence functions are commonly used to assess the impact of specific data points on a model, a process known as influence unlearning [R1,R8]. However, the utility of influence functions extends beyond mere data influence analysis. Influence functions are also powerful tools for solving bi-level optimization problems [R9], where the solution to a lower-level problem influences the outcomes at an upper level, such as evaluating model weight influence in Eq. (2). In our approach, we employ an influence function to compute the implicit gradient necessary for weight attribution, which we frame as a bi-level optimization problem. This method allows us to systematically assess and manage the impact of specific weights on unlearning efficacy and utility preservation.

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,

Thank you for your thorough review, positive assessment, and insightful feedback on our submission. Below, we provide detailed responses to each identified weakness (W) and question (Q). References (in the format of [Rx]) can be found in the general response.

W1/Q1: Could the authors provide additional clarification regarding the implementation and interpretation of the Membership Inference Attack (MIA) evaluation metric on the TOFU dataset?

A: Thank you for requesting further clarification on the implementation and interpretation of the Membership Inference Attack (MIA) evaluation metric on the TOFU dataset. We are sorry for any confusion if our initial explanation lacked clarity.

We employ the Min-k% Probability method [R5] as a "predictor" in the post-unlearning phase to assess whether test-time forget data were part of the LLM's original training set (Lines 309-312). Post-unlearning, the ideal case is that forgotten data should not be recognized as part of the training dataset, akin to test set data. In this context, correctly identifying forgotten samples as non-training instances equates to "true negatives" in our MIA approach (Appendix C.3, [R1]).

Based on the above, we measure the effectiveness of unlearning by calculating the AUROC for this MIA detector, rather than merely tallying true negatives. This involves considering both the forget set (treated as non-training data) and the retain set (data that remains part of the training set). A higher AUC indicates that the model accurately distinguishes between forgotten and retained data, reflecting more effective unlearning. This shows the MIA’s capability when evaluating unlearned LLMs to correctly classify forget and retain data points as belonging outside or inside the training set, respectively, thereby demonstrating successful unlearning.

W2/Q2: It would be beneficial if the authors could explore how the unlearning performance varies with changes in model size.

A: We appreciate your emphasis on the importance of scaling in LLMs. Here, we make the following clarifications.

(Rationale behind our choice) First, our choice of using the 7B model sizes is consistent with established benchmarks (e.g., TOFU) in the literature on LLM unlearning [R3, R6, R7]. Focusing on this specific model facilitates us to compare the unlearning performance of our proposal with other baselines. Second, our choice of model sizes was influenced by the GPU resources available to us. Our primary computational platform, equipped with RTX A6000 GPUs, imposes memory limitations on the feasible model size of LLMs we can experiment with extensively. This constraint affected our ability to perform extensive experiments on larger models, such as those exceeding 10B parameters. We will clarify this limitation in the section Limitations (Appendix E).

W3/Q3: 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.

A: Please refer to GR2 in the general response.

Thank you for your insightful comments. Below, we address each identified weakness (W) in detail. References (in the format of [Rx]) can be found in the general response.

W1: Discuss limitations and future work in the main paper rather than the appendix.

A: We will follow your suggestion to move the discussion of limitations and future work from the appendix to the main paper in the revision.

W2: Why does this method improve on existing unlearning methods?

A: Please refer to GR1 in the general response.

W3: 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.

A: Thank you for sharing this concern. We would like to make the following clarifications.

First, regarding our evaluation (e.g., in Table 1), we included multiple metrics to provide a comprehensive assessment of unlearning performance. These metrics do not always follow the same trend. For instance, a collapsed model might perform well on 1-FA and 1-ROUGEL for measuring unlearning efficacy but poorly on metrics like MIA and Forget Quality (FQ). This is because MIA and FQ measure differences between outputs from unlearned models on forget sets and retain sets, where a collapsed model could treat both sets the same, leading to lower performance on these metrics. Thus, for ease of comparison, we normalized the metrics to a [0,1] range and used their average values to give a general ranking. Figure R1 in the attached PDF showed a clear trend of improved tradeoff using WAGLE.

Second, we have added results over three independent trials and reported the standard deviations in Table R1. This table demonstrates that our method achieves better unlearning efficacy compared to baselines, as indicated by higher UE Avg. in bold. The improvement is statistically significant relative to the standard deviation. Additionally, the utility of our method (as indicated by UT Avg.) matches or even surpasses that achieved by unlearning using dense models.

Third, regarding the tuning of the trade-offs between retain and forget loss, specifically the hyperparameter $\lambda$, we followed the standard TOFU benchmark settings as outlined in references [R3,R4], which consistently set $\lambda$ to 1. This is discussed in Appendix B.4. To provide a clearer understanding, we have conducted additional experiments during the rebuttal phase, as shown in Table R2. We varied $\lambda$ across a range of values (between 0.1 and 10) using different unlearning methods for the dense LLM (that excludes weight attribution). The results indicate that UE Avg. consistently decreases while UT Avg. increases as $\lambda$ increases across different unlearning methods. This is not surprising as a larger $\lambda$ indicates a higher penalty on minimizing the utility loss on the retain set. Table R2 shows that even with tuning, dense models do not achieve as favorable a trade-off between utility and unlearning efficacy as WAGLE with the default $\lambda$ setting.

This additional analysis underscores that WAGLE's importance is not merely due to a specific $\lambda$ setting but rather due to its inherent weight attribution ability to achieve unlearning while taking utility optimization into account, i.e., balancing unlearning 'objective' with utility 'constraint'.

W4: There is some (limitations) 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.

A: Following your suggestion, we will revise the paper by discussing limitations in the main paper and add more discussions on societal ramifications.