GRU: Mitigating the Trade-off between Unlearning and Retention for LLMs

Thank you sincerely for your constructive comments and for helping us identify typos. We hope that the feedback provided below will address your concerns.

Q1. A major concern about the empirical results is the choice of baseline methods. Given that the authors are already using the WMDP benchmark, I would suggest them to consider RMU, which is proposed in the same paper. From my own research experience, RMU performs better than NPO in the trade-off between unlearning and retention significantly. Without considering RMU, the empirical improvements of GRU and TRU over the baselines are not very convincing to me.

A1. We aim to use a general set of representative methods for benchmarking. Regrettably, RMU, possibly due to its complexity to be implemented, is not widely adopted in earlier studies, such as [1,2] (though both cited WMDP). Although some recent papers incorporated RMU, their evaluations are either limited to WMDP [3], or they reveal its sensitivity to hyperparameters [4].

However, we completely agree that RMU represents a state-of-the-art within WMDP, which merits our particular focus. We present the results using RMU on WMDP, along with its version with GRU. As observed, GRU indeed enhance performance for both retention and unlearning, showing our GRU is general to mitigate the trade-off between unlearning and retention.

We adjust $\alpha$ from the default 1200 to 100 in the RMU open-sourced code, after finding that the original settings caused the retain term to overly dominate, hindering model updates. Similar sensitivity issues, such as precision setups, have been reported by others in the WMDP GitHub repository. We will study these issues in the future. However, we continue to see the RMU as a promising approach, particularly when viewing through the lens of local knowledge perturbation. We plan to delve deeper into the RMU concept and explore more specific verisons of GRU for RMU (e.g., constraints for masking) in our future work.

[1] Negative Preference Optimization: From Catastrophic Collapse to Effective Unlearning.

[2] MUSE: Machine Unlearning Six-Way Evaluation for Language Models.

[3] Simplicity Prevails: Rethinking Negative Preference Optimization for LLM Unlearning.

[4] Rethinking LLM Unlearning Objectives: A Gradient Perspective and Go Beyond.

Q2. Figure 3 is a bit confusing. Each algorithm's performance with and without GRU is shown as a segment in the figure. But it is unclear to me which endpoint of the segment is with GRU, i.e., whether GRU makes the performance better or worse. I can only guess from the caption that GRU improves the performance, but I think it would be necessary to make it clearer in the figure or at least in the caption.

A2. Apologies for any confusion caused by the figure annotations. As detailed in the figure caption, each pair of scores represents metric values (either FQ or MU) before and after applying the GRU enhancement. For instance, taking Figure 3(a) as an example, the pair of scores (-16.93, -3.52) represents the FQ scores for the GA method, where -16.93 is the score without the use of GRU and -3.52 is the score with GRU. The visual representation using an upward growing grid between these scores emphasizes the improvements achieved by incorporating GRU. We will clarify the meaning of our figures and provide more explicit explanations in our revision to ensure clear understanding.

Sincere thanks for your constructive comments, and we hope the following feedbacks can address your concerns.

Q1. Specifically, I question whether the current way of handling the retain set is sufficient. In my view, unless the utility performance on the retain set (e.g., MU or similar metrics) remains at 90–95% of the original model's performance, the improvements on the forget set may have limited practical meaning because the model's overall performance as a language model would already be compromised.

A1. We totally agree with your opinion. A large decline in utility performance would indeed render the resulting model ineffective, at which point the process of LLM unlearning also becomes meaningless. On the other side, carefully tuning the hyperparameters to achieve the goal of preserving 85, 90, and 95% of the original model performance can be tedious. Fortunately, as mentioned in Section 2.2, the UWC method offers a post-unlearning strategy that enables calibrating model performance via model mixing, alleviating the challenges associated with maintaining utility. Due to the relatively high cost of UWC, here we take GA and NPO under the challenging Phi setup as two examples to show UWC's flexibility and GRU's effectiveness. The improvements achieved by GRU are more pronounced. We will add more results in our revision. Many thanks for your suggestion.

Q2. In particular, how was the validation data chosen? Since there is no widely established standard for selecting validation data in unlearning settings, I believe the paper should provide a more detailed explanation on this aspect.

A2. Many thanks for your comments. This perspective has indeed been overlooked by the majority of the community, and we appreciate the opportunity to detail our implementation below.

For unlearn data, the datasets used for validation are exactly those used for unlearning in the cases of TOFU and MUSE, as they focus on privacy and copyright removal and do not require generalization. However, for WMDP, which aims to unlearn the entire concepts of harmful subjects, generalization is crucial. Hence, we separate a small set of test data (comprising 200 randomly selected samples) for evaluation, with respect to both the bio and cyber setups. For retain data, we separate 200 (20 for MUSE Book due to its small size of retain data) randomly selected samples from the original retain datasets used for unlearning, with respect to all three benchmarks. It ensures that we have a distinct set of data to verify retention compared to test data. We will add a detailed discussion in our revision.

Q3. Only WMDP-Cyber and MUSE-KnowMem results are reported, but VerbMem results are missing.

A3. The results for VerbMem are in Appendix B. We did not include these in the main text as their values and ranks are quite similar to KnowMem-U, offering limited new insights. We apologize for any confusion caused and will clarify this choice in our revision.

Due to strict space limits, we try our best to address the most critical questions as briefly as possible. We sincerely welcome any further concerns and will try our best to respond to them.

Q1. The improvement in removal on WMDP appears not as evident or consistent as on TOFU.

A1. WMDP uses a "Choose 1 from 4 QA accuracy," where random guessing (totally unlearned) would result in about 0.25. As observed, QA accuracies are already near 0.25, leaving little room for further declines. Our efficacy on WMDP is actually reflected by retention, with 2-7 percentage points increase over that without GRU. These results confirm that GRU maintains strong unlearning while improve retention, aligning our pirmary goal.

Q2. It would be helpful if more baselines can be included, such as RMU and FLAT.

A2. Many thanks for your suggestion. We present results using RMU and FLAT on the WMDP, along with their versions with GRU. As observed, GRU indeed enhances performance for both retention and unlearning.

We adjust $\alpha$ from the default 1200 to 100 in the RMU open-sourced code, after finding that the original settings caused the retain term to overly dominate, hindering model updates. Similar sensitivity issues, such as precision setups, have been reported by others in the WMDP GitHub repository.

We also expect the potential to further improve RMU and FLAT. For RMU, which perturbs localized representations, we aim to develop mask-based constraints for better knowledge localization. For FLAT, it uses a preference-based learning framework with distinct gradient behaviors from methods like GA and NPO. Thus, we need to devise new retain risk for more appropriate constraints, which we will explore in the future.

Q3. This paper has significant overlap in terms of methodology of GDR. WGA also examines the gradient direction of the unlearning objective.

A3. Regarding GDR, its requirement to store gradients across epochs is infeasible for LLMs due to memory costs and fewer epochs (e.g., 1 for MUSE). GRU avoids extensive caching and epoch-dependent operations, which is more practical for LLMs. Also, for LLMs, the retain data can lead to overfitting when directly guiding updates as in GDR. GRU, by not directly merging retain gradients into its updates, avoids this issue. Also, our enhanced version, TRU, explicitly address this concern, which is superior over GDR.

Regarding WGA's associated paper, it employs gradient computations to find reliable unlearning objectives. While conceptually possible, its practical challenges can moviate our work into optimization-focused research. Thus, our work extends beyond WGA's, further emphasizing the importance in exploring gradients, while focusing on a promising yet orthogonal direction.

Q4. It would be much appreciated if the authors could provide details to help validate the results.

A4. The configurations are detailed in Section 5 and Appendix C.1. The batch size for retain samples, set at 32, mirrors the unlearn batch and follows the random sampling as to GD. We further provide our source codes via the Anonymous GitHub.

Q5. For TOFU, are you using the TS-test (log) p-values as FQ? And for MU, are you using the aggregation of Truth Ratio, Prob, and ROUGE-L?

A5. To enhance readability and ease figure plotting, we uniformly report the $\log p$ values, avoiding dealing with extremely small $p$ like $1e^{-19}$. For MU, we adhere to TOFU, reporting the aggregated metrics.

Q6. By default, MUSE used LLaMA-2-7B on the news subset and ICLM-7B for the HP books subset, and you are only using the ICLM-7B model. Are you testing on the entire MUSE, or just the books subset?

A6. Following your suggestion, we further report results on MUSE News with LLaMA-2-7B, taking GA and NPO due to space limit. UWC is adopted for calibration towards a fair comparison. As observed, there are notable improvements in PrivLeak, espeically for GA.

Q7. Could GRU be seen as an explicitly-constrained variant of GD that relies on retention gradient estimates from up-sampled retain data?

A7. Yes, your understanding is correct. If GA as the objective and combining it with GRU corresponds to the explicit-constrained variant of GD. The benefits of our approach are clear, which forcibly ensures retention. In contrast, using original GD can result in unlearning gradients dominating updates, adversely affecting retention.

Can you elaborate on how you did that? From the code (`dataloader.py`), it appears that you are still caching the unlearning and retain gradients in `compute_loss` for every training step, flatten and store the structure map of each gradient with `flatten_and_store_grads`. To me, this does not seem less complex than the [GDR implementation](https://github.com/RUIYUN-ML/GDR-GMA/blob/main/memory_bank.py).

self.flattened_retain = self.moving_avg * self.flattened_retain_accumulation + (1 - self.moving_avg) * self.flattened_retain_old

bank = MemoryBank(size=math.ceil(t_dataset_sizes/args.batch_size))

Many thanks for your great support and constructive comments! Please see our responses below.

Q1. It is unclear how the constraints on gradients from different samples and mini-batches are enforced in Eq. (7).

A1. As shown in Algorithm 1, mini-batches $B_{\rm r}$, $B_{\rm u}$ replace $D_{\rm r}$, $D_{\rm u}$. In Section 3.2, we further discuss the limitations of mini-batch gradients and recommend using the exponenital moving average for stable estimation. We will clarify these points in our revised manuscript. Sincere thanks for your suggestion!

Q2. Performing grid search to find the optimal values for these hyper-parameters is computationally expensive and time-consuming.

A2. Many thanks for raising this concern. We would like to clarify the practicality of our approach from the following two perspectives.

1. The GRU performs well even without these hyper-parameters. Our basic framework following Eq. (7) operates without any hyper-parameters. Incorrperating $\gamma$ (for EMA) and $\tau$ (for clipping) are practical enhancements that improve empirical results. Below is the table showing results for 3 methods under the LLaMA setup, without using $\gamma$ and $\tau$. The results highlight notable improvements of our method over the baselines.

2. The GRU is not very sensitive to these hyper-parameters. By incorporating $\gamma$ and $\tau$, we explore the potential to further enhance unlearning. In Appendix D, we conduct a hyper-parameter sensitivity analysis, showing stable improvements across a wide range of candidate hyper-parameters.

Q3. Some claims need further justification, for example, "..., the remaining data points within $D_{\rm u}$ can offer information for retention.".

A3. Heuristically, retain gradients primarily act as a denoising mechanism, rather than directly contributing update directions. By using retain gradients to redirect original unlearn gradients, we focus on the direction that removes targeted knowledge, while with limited side impacts on unrelated knowledge.

For TRU, let's first consider a simple scenario to remove a single $s_{\rm u}$. Here, it is straightforward to take the complementary, i.e., $D_{\rm u} \backslash{s_{\rm u}}$, for retention. We then create a rectified task vector targeted at eliminating knowledge associated with $s_{\rm u}$ while minimizing impact on unrelated knowledge. This process can be applied to each $s_{\rm u}$ within $D_{\rm u}$, with the collective average forming the rectified task vector for the entire $D_{\rm u}$.

Notably, these task vectors remain mutally compatible, as the retain gradients are exclusively employed for denoising rather than direct used into the task vectors. We will add the related discussion in our revision.

Q4. The results presented in the figures are not clear.

A4. Apologies for any confusion caused by figure annotations. As shown in captions, each pair of scores represents metric values (either FQ or MU) before and after applying GRU. Taking Figure 3(a) as an example, the pair (-16.93, -3.52) represents FQ for GA, where -16.93 is the score without GRU and -3.52 is the score with GRU. The visual representation using an upward growing grid between these scores emphasizes the achieved improvement. We will provide more explanations in our revision to ensure clarity.

Q5. The preliminaries section takes up excessive space. The authors should consider making this section more concise while ensuring it remains comprehensive.

A5. We have a relatively extensive preliminaries section to ensure self-containment and maintain clarity. However, we fully concur with your opinion that a more concise verision would enhance the paper flow and allow larger space to elaobrate on our core contributions regarding GRU and TRU. We sincerely appreciate your feedback and will carefully refine our paper structure accordingly.

Q6. You mentioned that "retain data can often be biased." Could you clarify what bias refers to in this context?

A6. Many thanks for your question. Taking TOFU as an example, the current setup involves selectively unlearning specific author profiles while retaining others. However, we recognize that the broader objective of retention should ensure model capacity across diverse domains such as humanities and sciences. Therefore, the current retain data in TOFU may exhibit bias, stemming from the distribution shift between the adopted retain data and the broader real data. We will refine the related discussion in our revision.