Low Compute Unlearning via Sparse Representations

Authors: Vedant Shah, Frederik Träuble, Ashish Malik, Hugo Larochelle, Michael Curtis Mozer, Sanjeev Arora, Yoshua Bengio, Anirudh Goyal
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: Sparse Representations, Discrete Bottlenecks, Model Editing, Unlearning

Abstract:
Machine \emph{unlearning}, which involves erasing knowledge about a \emph{forget set} from a trained model, can prove to be 
costly and infeasible using existing techniques. We propose a  low compute unlearning technique based on a discrete representational bottleneck. We show that the proposed technique efficiently unlearns the forget set and incurs negligible damage to the model's performance on the rest of the data set. We evaluate the proposed technique on the problem of *class unlearning* using four datasets: CIFAR-10, CIFAR-100, LACUNA-100 and ImageNet-1k. We compare the proposed technique to SCRUB, a state-of-the-art approach which uses knowledge distillation for unlearning. Across all three datasets, the  proposed technique performs as well as, if not better than SCRUB while incurring minimal computational cost.

Decision: reject
Reviewer scores: 3, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1730654725327

thread_type:
official_review

<thread_object>
creator: reviewer_npds

<content>
summary:
- This paper presenting unlearning with DKVB (Discrete Key-Value Bottleneck), and studies the class unlearning problem
- It uses the DKVB architecture which uses a pre-trained encoder, followed by discrete key quantization, then maps these keys to values, then trained a linear encoder on these values
- In order to unlearn a class, it records the most commonly used key/values pairs recorded by that class, then removes them from the codebook to unlearn them. (Note there is a difference between two methods of Unlearning with DKVB, examples vs. activations, but this is the broad summary)
- The method is able to unlearn a class with very little effect on the classification accuracy of the remaining classes
- There is a particular emphasis on the low compute cost compared to existing unlearning approaches, as this method does not require access to the original data (in the case of unlearning via activations), in comparison to other methods which do require the original data

strengths:
- Not requiring access to the original data I believe is very important and this strength should perhaps be stressed more in the paper
- The method requires very little compute compared to existing methods
- Presentation is straightforward an easy to follow

weaknesses:
- I think that the presentation of table 1, particularly the bolding is misleading. While it is true that the DKVB unlearning results in the smallest change in the accuracy of the retain set, wouldn't one prefer a performance *increase* on the retain set, as it no longer needs to predict that last class?

questions:
- Is the class logit for the forget class removed entirely for the experiments in Table 1?
- Is it possible product similar plots to figure 2 and figure 3 for one of the baseline unlearning methods? For just the smaller datasets such as CIFAR-10/100 would be sufficient
- I think more discussion on multi-class unlearning would be useful. For example the figures in appendix A.3. (figure 5) maybe could be moved into the main text, and comparisons to one of the baseline methods would give a better characteristics of the limits of the approach. At least based on my intuition, this approach could fail if there are many classes being unlearned and "too many" keys/values are removed from the codebook. Can the authors provide any insight on this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730673516777

thread_type:
official_review

<thread_object>
creator: reviewer_gfx9

<content>
summary:
The paper proposes a low-compute machine unlearning technique using a discrete representational bottleneck to efficiently erase knowledge of a forget set from trained models with minimal impact on overall performance. Evaluated on datasets like CIFAR-10, CIFAR-100, LACUNA-100, and ImageNet-1k, the method matches or outperforms SCRUB—a state-of-the-art unlearning approach—while incurring negligible computational cost.

strengths:
- The paper is well-structured, presenting with thorough results and analysis.
- The proposed technique demonstrates effective unlearning performance with minimal computational cost, validated across multiple datasets.

weaknesses:
- The work appears to be primarily built upon the DKVB model but applied in a reverse manner. The authors should more explicitly highlight their unique contributions and clarify how their approach differs from DKVB.

- The study focuses only on ViT-B/32 and ResNet-50 models. It's unclear why these particular models were chosen and whether the findings generalize to other architectures.

questions:
- The paper states that the key initialization is done on the training dataset. Is it sufficient to use only a few batches for initialization, or does it require processing the entire training dataset? If it's the latter, wouldn't this be computationally expensive?

- In the experiments with ImageNet-pretrained ResNet-50 backbones, the retain set test accuracy decreases when the forget set test accuracy approaches zero (Figure 2), whereas in other settings, the retain set accuracy remains almost constant. Could the authors provide an explanation?

- Unclear Symbol on Page 14: There are two "¿" symbol on page 14 whose meaning is not clear. Could the authors clarify its significance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688655896

thread_type:
official_review

<thread_object>
creator: reviewer_mhv7

<content>
summary:
This paper propose a mechanism, based on prior work on the discrete key-value bottleneck (DKVB), for tacking the unlearning problem—how to adjust a model such that it no longer performs well on particular training examples.

The DKVB is an approach for learning data from non-stationary distributions (e.g. for continual learning). A fixed, pre-trained encoder is used to generate input-dependent queries for a set of key-value lookup tables. The keys are initialized to cover the full representational space of the encoder and thereafter fixed. The values are used to feed a decoder (which is non-parametric averaging pooling in this work), and the values can be learned. Since each query maps to $k$ keys and only the values associated with those keys are used for decoding, masking key-values belonging to particular samples allows DKVB to be a useful approach for unlearning.

strengths:
Overall, this is a well-written and motivated paper. The prior work is discussed and the authors successfully position their work as original and significant for tackling the problem in a compute-efficient manner. The ideas presented are intuitive and well-positioned to address the unlearning problem. The results appear promising and the evaluation is correct.

weaknesses:
- The authors focus on image classification tasks where the decoder has only one primary output, but for many large-scale pre-trained models such as LLMs, the backbone serves many downstream tasks. Presumably, the decoder will need to produce outputs that apply equally well to all tasks. Whether this is a limitation or an issue of scaling is not discussed.

- The input representation of the encoder seems crucial for separating forget keys from retain keys, but this work has limited discussion and evaluation on how to choose a backbone.

questions:
1. The bottleneck design necessarily limits the DKVB output to a smaller subset of downstream tasks than the backbone. Do you envision separate DKVBs for different tasks. How does this compare to an approach that directly modifies the backbone for unlearning?
2. While the evaluation includes two different backbones, there is limited discussion of how a backbone's representations affect the DKVB. Do you have recommendations for which layer in the original model should be used as the encoder? Should it always be the full backbone?
3. Related to Question 2, what if you need to forget specific samples (instead of whole classes)? Or have samples in the retain and forget sets with very similar input representations? How does DKVB fair in these settings (there appears to be a little discussion about these limitations, but it might be more helpful to contextualize w.r.t. the DKVB operation)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730853262495

thread_type:
official_review

<thread_object>
creator: reviewer_x8bn

<content>
summary:
This paper studies machine unlearning in the setting of models with discrete key value bottlenecks, specifically in the setting where compute is limited. Specifically, this work focuses on models possessing discrete key value bottlenecks, an information bottleneck proposed recently. The use of this bottleneck facilitates the use of the proposed methods, which identify which keys are necessary for predictions of a) groups of samples, and b) certain classes. The methods are zero-shot, and their effectiveness at unlearning in this situation is demonstrated empirically.

strengths:
The paper has several strengths:

- The proposed methods appear to be highly effective at the tasks of model unlearning (based on the empirical results proposed in Section 5). 
- The proposed methods are significantly more compute-efficient than existing methods (SCRUB), while attaining similar levels of accuracy.

weaknesses:
The paper has several weaknesses:

**Clarity of Exposition**

- The paper is tricky to parse since no equations representing any of the operations (either describing the mechanics of the DKVB or the unlearning strategy) are provided. This adds ambiguity that should be avoided. Other mathematical details that are important for readers are the exponential moving average used for key initialization. I recommend the authors provide these formal details, at the very least in the appendix.
- There are some issues with imprecise language as well, which I recommend the authors address in subsequent revisions, as they affect the readability of the paper. Two examples that stood out to me were:
    - For instance,  on lines 223-224, the phrase  "This ensures that the representations are distributed sparsely enough." This is an unclear statement. It would help significantly if the authors elaborated upon this in a precise fashion.
    - Similarly, in lines 238-239, the authors write "Technically, this approach requires access to the original training data corresponding to the forget class. However, it is also possible to carry out this procedure with a proxy dataset that has been sampled from a distribution close enough to that of the forget set." This is somewhat contradictory, and should be clarified.
- Typos: there are a few typos in this work. For instance: "sufficient" -> sufficient, "ther specific subset" -> the specific subset.

**Experimental Results**

There are some points in the experimental section that could be made more clear.
-  The experimental results in Section 5.2 (for different numbers of samples) should be tabulated neatly and legibly. As is, it is slightly confusing and tricky to parse.
- In Appendix A.7, it's stated that for the ResNet50 model, either the third last or fourth last layers are used in the encoder backbone. It's unclear why these particular layers were chosen for this task.

questions:
I have a few questions about this work.

1. On lines 179-182, it's stated that "Although one could argue that our methods lean towards a weak unlearning strategy–given the pre-trained backbone might retain some information about the forget set–our approach deviate from the strict definition of weak unlearning as outlined by Xu et al. (2023)." This is a major concern, particularly if you take an encoder and a DKVB trained on Imagenet. Moreover, this contradicts later points in the paper where the authors consistently refer to "complete unlearning" (i.e. lines 427-429, 460-463).  If the goal is to unlearn a class (for instance) totally, isn't removing that class' information from the encoder as well crucial to the success of the method?

2.  On lines 253-255, the authors state "However, using one approach over the other may be more practical or even necessary, depending on the task at hand." Could the authors elaborate more upon this point? For what tasks would one use UvA over UvE (and vice-versa)?

3. On lines 257-259, the authors state "Moreover, the requirement of access to original training data of the forget class can also be circumvented under appropriate assumptions". Could the authors elaborate upon those assumptions? Moreover, since this is a strong claim made by the authors, is there any theoretical or empirical justification for this claim
</content>
</thread_object>

</thread>
<thread>
creation date:
1732904613375

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updated Manuscript

comment:
We thank all the reviewers for their very useful suggestions. We would like to bring to your notice that we have updated the manuscript by incorporating the suggestions given by the reviewers. All the modifications are marked in blue text. Below is a summary of all the changes:

* Added mathematical and algorithmic formulation for EMA and the proposed approaches in Appendix A.8
* Elaborated upon the applicability of the approach in the case where training data cannot be accessed after the training (Section 4, Paragraph 4)
* Added a line on the applicability of the approach on tasks other than image classification in Section 6. (Limitations and Future Work)
* Added results on multiclass unlearning using a baseline (SCRUB) in Appendix A.4
* Added a plot similar to Figures 2 and 3 for a baseline (SCRUB) in Appendix A.3
* Added a paragraph outlining the main contributions of the paper at the end of Section 1
* Changed the bolding scheme in Table 1 to highlight the approaches which achieve the maximum retain class test accuracy rather than minimum damage with respect to the original accuracy on retain class.
* Added a reference to the multiclass unlearning experiments (Appendix A.4) in the main paper (in Section 5.2).

We would be happy to engage in further discussions to clarify any concerns and questions that the reviewers have
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper studies the problem of class unlearning, where the goal is to remove the information about one (or a set of) classes from the model. The authors approach this problem by training a discrete key-value bottleneck (DKVB) model on top of a frozen encoder. Then DKVB is a decoder model that makes predictions based on similarity to a frozen set of keys which are selected in a non-parametric manner. Then, the authors remove some of the keys that are most used in predicting on the target class. The authors propose two variations of the method: unlearning via examples and via activations. Across a variety of image tasks, the authors show that the proposed method can reduce the performance on the target class to 0% while preserving performance on the remaining classes.

Strengths:
- The proposed method can unlearn classes at a much lower cost than competing methods
- In theory, the proposed method is applicable to unlearning subsets of classes, and general sets independent of classes
- The authors show results on a range of datasets and architectures

Weaknesses:
- In my opinion, the core weakness of the paper has been identified by reviewer X8bn: it is possible to remove a class (or a set of classes) from a classifier by simply masking the logits. This approach would completely “unlearn” the forget class, while improving the performance on retain classes. The experiments presented in the paper do not differentiate the proposed approach from this simple masking. In theory, the proposed method can work in more general settings, and I believe that should be the core of the experimental evaluation, but the authors only show results for class unlearning.
- The method does not remove information about the class from intermediate representations.
- The clarity of presentation should be improved. Specifically, the algorithm should be explained more clearly and precisely in the main text. Currently, the description, including Figure 1, is quite hard to follow.

Decision recommendation: based on the above, I recommend rejecting the paper. In my opinion, the evaluation should focus on the more challenging  setting of unlearning groups that do not align with classes. Also, the fact that information about the forget set is still contained in the representations should be addressed.