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

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.

Strengths

Strengths: - Novel approach to LLM unlearning using weight attribution - Theoretical grounding in bi-level optimization - Comprehensive evaluation across multiple unlearning tasks and benchmarks - Compatibility with existing unlearning methods (GradDiff, NPO, PO) - Insights into which model components are most influential for unlearning

Weaknesses

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: 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?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

3

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

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.

Strengths

The authors provide an extensive and thorough evaluation of their proposed method. In particular, they consider 1) multiple base unlearning algorithms, 2) multiple well-motivated pruning baselines, 3) multiple unlearning efficacy heuristics, and 4) multiple downstream tasks to evaluate model utility. Another strength on the experimental side is that experiments are performed on Llama2-7B; the size of the model alone displays the scalability of the proposed method.

Weaknesses

- 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

- 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

Rating

4

Confidence

3

Soundness

3

Presentation

2

Contribution

2

Limitations

N/A

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.

Strengths

1. This paper is pioneering in the domain of machine unlearning for large language models (LLMs), providing a formal definition and derivation that demonstrates how weight attribution contributes to unlearning processes. Unlike prior studies on unlearning in computer vision, which were primarily intuitive, this paper’s methodical approach through formal derivation introduces a significant novelty to the field. 2. The insights offered by the proposed framework are substantial, particularly in identifying critical modules within LLMs that are pivotal for effective unlearning, as detailed in Figure 2. This not only deepens our understanding of the inner workings of LLMs but also highlights specific areas for targeted unlearning strategies. This may also helps understanding the memory of LLMs. 3. The clarity and organization of the paper enhance the reader's comprehension and engagement. 4. It features a comprehensive set of experiments across diverse tasks, benchmark datasets, and models, convincingly demonstrating the efficacy of the proposed weight attribution methodology in facilitating LLM unlearning.

Weaknesses

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

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?

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

4

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.

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.

Strengths

* Novel approach: WAGLE integrates weight analysis with unlearning, providing a new perspective on LLM unlearning. * Comprehensive evaluation: The method is tested on various unlearning tasks and datasets, showcasing its versatility. * Improved performance: WAGLE consistently enhances forgetting ability while maintaining or sometimes improving task performance on retain sets. * Insights into model architecture: The study reveals that early-to-mid layers are important for unlearning, and the method tends to avoid selecting weights from self-attention mechanisms. * Comparison with existing methods: The authors benchmark WAGLE against various direct optimization approaches for forget/retain loss, providing context for its performance. * The bi-level optimization formulation is interesting and could be built on by future methods.

Weaknesses

* 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

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.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

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.

update

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.

Response to 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.