Paper: Dissecting Language Models: Machine Unlearning via Selective Pruning (ICLR-2024)

Reviewer comment (Evaluation):
The experimental results presented by the author do not fully support the conclusion this paper wants to draw. For example, the authors mentioned in the introduction that, "If capabilities can be separated on the level of neurons, then this can lead to modularity inside models." However, based on the experimental results, it appears that the entire LLM behaves highly in coupling and cannot be separated. For instance, in Figure 1(d), the performance loss on the 'code' dataset seems to be mirrored closely by a performance loss on the 'python' dataset, where the model drops close to the retain dataset and the forget dataset.

Author response:
We think that the extent to which different capabilities are coupled or intertwined depends on the capabilities in question. By comparing the results in Figure 2(a) and 2(b) we see that ‘python’ ability and ‘code’ ability are more intertwined than ‘code’ ability and ‘pile’ (i.e. general text) ability. This is in line with what one might intuitively expect: coding ability in general and coding in the python programming language seem like similar skills.

Reviewer comment (Experiments):
Given that the experiments were solely conducted on datasets related to code, I am uncertain about the generalizability of the experimental results presented in the paper. For instance, the conclusion that FFN outperforms attention—might it be possible that a different task could yield an opposing conclusion.

Author response:
We agree that the degree to which attention and feed forward pruning may be specialized to different tasks may differ. Thank you for pointing this out. We have updated the pdf to reflect this.

Reviewer comment (Presentation):
The readability of the entire article is not good. For instance, in Section 3.1, the author describes the distribution characteristics of the "**attention pre-out neuron**" activations. However, the definition of this unfamiliar term, "attention pre-out neuron," is only introduced in Section 3.3. This leads to confusion for me when initially encountering the term. Additionally, the article frequently places experimental results in the appendices, referencing them in the main text and using the conclusion from these experiments in appendices to support further observation in the main text. This approach disrupts the reading flow, often requiring to flip back and forth for context. While thorough analyses and experiments are commendable, the structure of the article still needs further refinement to enhance its logical flow.

Author response:
Thank you for the feedback. We have resolved the issue of using the term “attention pre-out neuron” before we explain the term. We have also made other updates to the paper to improve readability, such as improving the wording and adding a diagram describing our method on page 3.

Reviewer comment (Novelty):
A prior study [1] demonstrated that, depending on the specific input, it's possible to achieve a high pruning ratio without negatively affecting performance and without the need for retraining. This suggests that there is redundancy in the neurons of the LLM when it's tasked with executing a singular function (equating a single sentence to a minor task, for instance). In light of these findings, what novel insights or observations does your paper offer in comparison to that study?

[1] Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

Author response:
This paper is very interesting indeed! We have now cited this paper. Their findings are similar to but more general than the ones in [A] which motivate our importance functions. The paper shows contextual sparsity per input sentence, whereas we show something similar per dataset.

The paper shows that neurons specialize locally on the scale of tokens and sentences. We show that neurons specialize more globally on the scale of broader tasks and datasets.

In Figure 1 (now Figure 2) we have now added imagenet data, and we think it is interesting that the three different pairs Pile-Coding, Coding-Python, Imagenet-Bird all have very different amounts of separability. This shows that the separability of neurons is task dependent. 

[A] MoEfication: Transformer Feed-forward Layers are Mixtures of Experts, Zhang et al. https://arxiv.org/abs/2110.01786

Reviewer comment (Novelty):
the related work section is weak, missing a lot of literature both from unlearning and pruning. For example, [A-E] are recent papers on unlearning, with [A, B] particularly related to the methodology of this paper (see References below). I’m less familiar with the sparsity and pruning literature but the authors should conduct a thorough review.

Author response:
[A] This Arxiv paper from August 15th is related and interesting!
In a quick trial we applied our method to the Vision Transformer (ViT) mushroom (MR) row in their Table 2. On the test split (100 datapoints) of CIFAR100 we find a forget drop from 72±8% to 0±3% and a retain drop from 90.0±0.3% to 89.5±0.3%.
[A-E] Cited.

Reviewer comment (Experiments):
the authors compare against only one baseline, and only for one task. Other common baselines include finetuning on the retain set, gradient ascent on the forget set, and comparing against other recent unlearning methods is also important (see the references below)

Author response:
We acknowledge the suggestion, and have updated the paper to add comparisons to previous work for vision transformers on CIFAR100. Specifically, we compare to: retraining; finetuning; the ‘incompetent teacher’ method; UNSIR; amnesiac and selective synaptic dampening. We have also included a membership inference attack (MIA) metric.

See Table 4 on page 8 of the newly updated pdf.

Reviewer comment (Evaluation):
the particular problem setting of unlearning is not clearly defined. What defines successful unlearning here? In section 3 the authors describe the forget set as the “task or dataset that we aim to reduce performance on”. This isn’t precise enough. How much do we want to reduce performance? Is it that the greater the reduction, the better? For context, unlearning papers (e.g. Golatkar et al, which the authors cite) usually consider that the accuracy on the forget set should be reduced only up to a reference point and no further (where the reference point is given by the oracle unlearning algorithm of retraining the model from scratch without the forget set). Is the setup here different, and if so, how is the goal defined in this case?

Author response:
While we agree that retraining is a good standard, retraining an LLM is too expensive. We simply try to maximise drop in forget while keeping performance in retain, seen with a "curvier" graph being better. We do not insist on a specific goal, but make a cheap reference, and allow the user to decide where they would like to be on the trade-off between a high retain & low forget accuracy.

Reviewer comment (Evaluation):
[soundness]: when it comes to forgetting or unlearning, several metrics have been proposed by the community to measure this. Simply inspecting the accuracy / perplexity is likely a poor proxy for forgetting quality. For example, Membership Inference Attacks are an important category of methods (see e.g. Golatkar et al and reference [E] below). Are these not applicable here. If not then why not?

Author response:
In this paper we have tried to show that harmful behavior can be reduced by applying selective pruning. We are less focused on removing specific datapoints from the model than we are in removing the model’s ability to generate code or toxic data. Studying membership inference attacks (which requires full knowledge of training data) is beyond the scope of our current analysis.

Reviewer comment (Presentation):
[presentation] In the paragraph under Definition 1, the authors give intuition for the different influence functions but they omit I_{abs}. Please add.

Author response:
Thanks! Added.

Reviewer comment (Experiments):
[empirical results] It would greatly strengthen the paper to conduct analyses of: the effect of the number of pruning iterations, the effect of the % pruned (in each iteration or overall), the effect of the choice of the influence function. It sounds like the authors have some results on some of these in the Appendix but not all of them. It would also be great to summarize in the main paper all of the findings (so that one doesn’t need to read the entire Appendix).

Author response:
We have chosen not to do a grid-search for the hyperparameter “percentage pruned per iteration”, but instead report on all the hyper-parameters that we tried.
The effect of how much has been pruned overall can be seen in Figure 2. Every dot here reflects another pruning step. The more pruning steps the bigger the reduction in accuracy.
Figure 4 shows the effect of the choice of influence function.

Reviewer comment (Evaluation):
[empirical results] please include confidence intervals in all tables and e.g. in Figure 3. It is currently challenging to tell if the differences are significant

Author response:
In our tables, we apply the 'significant digits' convention, where we round our values to the most relevant figure and omit digits that represent a substantial error margin. We have taken the suggestion to add error bars to Figure 4.

Reviewer comment (Experiments):
In Table 3: it would help readability to label "Task Arithmetic (quoted)" as being a finetuning task. Also, the gap between the baseline (quoted) and (replicated) is fairly large, which makes actually comparing the finetuning vs presented pruning method difficult. Would it be possible to replicate the task-arithmetic finetuning? That would significantly improve the comparability of these results.

Author response:
> In Table 3: it would help readability to label "Task Arithmetic (quoted)" as being a finetuning task.

Thank you for the suggestion! We have updated the label.

Reviewer comment (Reproducibility):
The details of iterative pruning are not specified. It seems like it would be important to the function of the method to set the proportion of nodes pruned per iteration well, but this is not discussed. Even if this is not an important hyper-parameter of the method, a clearer explanation of the iterative pruning method is necessary to fully elucidate the applied method.

Author response:
> The details of iterative pruning are not specified. It seems like it would be important to the function of the method to set the proportion of nodes pruned per iteration well, but this is not discussed.

We have added some clarification in Section 3.2 in the updated pdf. We have not done any hyper-parameter tuning for the proportion of nodes pruned per iteration.

Reviewer comment (Evaluation):
6.3: If an LLM were trained not to answer questions about a dangerous topic, say, bomb-building, could the presented method not be used to unlearn that guardrail? I don't think this is a concern specific to the presented method, but I do not follow why this method is less likely to generate harmful systems than other methods. Can the authors clarify?

Author response:
> 6.3: If an LLM were trained not to answer questions about a dangerous topic, say, bomb-building, could the presented method not be used to unlearn that guardrail? I don't think this is a concern specific to the presented method, but I do not follow why this method is less likely to generate harmful systems than other methods. Can the authors clarify?

Something that was not clear in our original formulation is that we think machine unlearning methods (such as ours) will be more thorough than other model control methods.

For example, RLHF (a control method) guard rails are fragile to things like jailbreaking [1] and finetuning [2]. This is possible because the information about how to build a bomb is still in the model. This method is an attempt to remove much of that information so that even if one unlearns the guard rails, one can still not learn how to make the bomb.

[1] Deep reinforcement learning from human preferences, Christiano et al. https://arxiv.org/abs/1706.03741

[2] LoRA Fine-tuning Efficiently Undoes Safety Training in Llama 2-Chat 70B, Lermen et al. https://arxiv.org/abs/2310.20624

Reviewer comment (Presentation):
Including a visual representation of the suggested concept could enhance comprehension. Thus, kindly incorporate a diagram that presents a general outline of the proposed concept.

Author response:
We have now added a visual representation (new Figure 1) of selective pruning to Section 3.2 on page 3 of our updated pdf.

Reviewer comment (Experiments):
The evaluation of the proposed model focuses on text datasets, yet the experiments for the image dataset are absent. Implementing machine unlearning for the image dataset would be highly beneficial.

Author response:
We have taken your suggestion to implement selective pruning for an image dataset on board! Please see our updated Figure 1 (now Figure 2) in the updated pdf.