SalUn: Empowering Machine Unlearning via Gradient-based Weight Saliency in Both Image Classification and Generation

Machine unlearning (MU) is the task of efficiently and effectively mitigating the influence of particular data points on a pre-trained model (Shaik et al., 2023). It emerged in response to data protection regulations like ‘the right to be forgotten’ (Hoofnagle et al., 2019). However, its scope and significance rapidly expand to tackle many trustworthy machine learning (ML) challenges in computer vision (CV). These challenges include the defense against backdoor poisoning attacks (Liu et al., 2022a), the enhancement of model fairness (Oesterling et al., 2023), the refinement of pre-training methods to augment transfer learning capabilities (Jain et al., 2023; Jia et al., 2023), and the prevention of text-to-image generative models from generating sensitive, harmful, or illegal image content when exposed to inappropriate prompts (Gandikota et al., 2023).

Roughly speaking, current MU methods can be categorized into two families: exact or certified MU and approximate MU. The former focuses on developing methods with provable error guarantees or unlearning certifications. Examples of such methods include differential privacy (DP)-enforced unlearning and certified data removal (Guo et al., 2019; Chien et al., 2022). Within this family, exact unlearning, which involves retraining a model from scratch after removing the forgetting dataset from the original training set, is typically considered the gold standard of MU (Thudi et al., 2022b; a). However, retraining-based exact unlearning methods require significant computation resources and have become challenging for today’s large-scale ML models, such as the diffusion-based generative model considered in this work.

In contrast to exact or certified MU, approximate unlearning has emerged as a more practical approach for ‘fast’ and ‘accurate’ unlearning. While the accuracy may not meet provable guarantees, it can b assessed using a broader range of practical metrics, such as membership inference attacks (Carlini et al., 2022), without necessitating data-model or algorithmic assumptions typically associated with certified unlearning. Despite the merits of practicality and efficiency, the performance of approximate unlearning can still exhibit significant variance. For example, influence unlearning (Izzo et al., 2021; Warnecke et al., 2021), built upon the influence function analysis of training data points (Koh & Liang, 2017), exhibits high-performance variance due to the selection of hyperparameters required for influence function approximations, as well as the particular unlearning scenarios and evaluation metrics (Becker & Liebig, 2022), thereby raising concerns about instability in approximate unlearning methods. Other approximate unlearning methods, including Fisher forgetting (Golatkar et al., 2020), gradient ascent (Thudi et al., 2022a), and finetuning-based approaches (Warnecke et al., 2021; Jia et al., 2023), also face the similar challenge as will be illustrated later.

Furthermore, many MU methods mentioned above have been primarily applied to image classification. By contrast, emerging diffusion models (DMs) for generative modeling also demand effective MU techniques to protect copyrights and prevent generation of harmful content (Schramowski et al., 2023; Gandikota et al., 2023; Zhang et al., 2023a). However, as this work will demonstrate, existing MU methods designed for image classification are insufficient to address MU in image generation (see Fig. 1 for a schematic overview of our proposal vs. conventional MU).

In response to the limitations of existing MU methods, we aim to address the following question:

To tackle (Q), we develop an innovative MU paradigm: ‘weight saliency’. Drawing a parallel with input saliency in model explanation, our idea shifts the spotlight of MU from the entire model to specific, influential weights. Such focused attention can enhance the performance of multiple MU methods, even simple ones such as random labeling. Termed ‘saliency unlearning’ (SalUn), our approach can diminish the performance gap with exact unlearning, offering a principled MU method effective across image classification or generation. Our contributions are as follows. ❶ We identify two limitations of current MU techniques: instability, e.g., when faced with varying amounts of forgetting data, and lack of adaptability to image generation tasks. ❷ We introduce the concept of weight saliency in MU and develop SalUn, a saliency-guided approach. We show that weight saliency could be a key to addressing the limitations of current MU methods. ❸ We perform comprehensive experiments to validate the effectiveness of SalUn, comparing it with 7 MU baselines in image classification and 2 concept-erasing baselines in image generation. As a notable application, we show that SalUn is the most effective method in preventing stable diffusion from generating harmful images when given inappropriate prompts (I2P) (Schramowski et al., 2023).

Unlearning in image classification. MU aims at modifying ML models to eliminate the influence of specific data points or classes, initially developed to mitigate potential privacy breaches post-training (Ginart et al., 2019; Neel et al., 2021; Ullah et al., 2021; Sekhari et al., 2021). However, exact unlearning, i.e., retraining from scratch, though theoretically sound, introduces substantial computational demands. To alleviate this, some research efforts have explored probabilistic methods like differential privacy (DP) (Ginart et al., 2019; Guo et al., 2019; Neel et al., 2021; Ullah et al., 2021; Sekhari et al., 2021). Still, these methods often have inherent limitations that hinder their practical effectiveness, especially in defending against membership inference attacks (Dwork et al., 2006; Graves et al., 2021). Therefore, there has been a shift towards developing more effective and efficient unlearning strategies (Golatkar et al., 2020; Becker & Liebig, 2022; Thudi et al., 2022a; Jia et al., 2023; Chen et al., 2023; Warnecke et al., 2021). The landscape of MU has also expanded to encompass diverse domains, such as federated learning (Wang et al., 2022; Liu et al., 2022b; Wu et al., 2022) and graph neural networks (Chen et al., 2022a; Chien et al., 2022; Cheng et al., 2023).

Unlearning in image generation. Recent advancements in text-conditioned image generation models have demonstrated remarkable capabilities in producing high-quality images that closely align with textual descriptions (Rombach et al., 2022; Ho & Salimans, 2022). However, these achievements often rely on extensive datasets, such as LAION-400M and LAION-5B (Schuhmann et al., 2021; 2022), which inherently introduce biases and associated risks. These concerns are indicative of broader issues within the field, as highlighted by various studies (Birhane et al., 2021; Schramowski et al., 2023; Somepalli et al., 2023; Bae et al., 2023; Zhang et al., 2023b). To address these challenges, there is a pressing need to explore effective MU techniques. While current studies (Gandikota et al., 2023; Zhang et al., 2023a; Heng & Soh, 2023) provide strategies for concept erasure in diffusion models, achieving precision comparable to exact unlearning remains challenging.

Data and model saliency analyses. There has been extensive research on input saliency maps for the development of explainable ML techniques. Examples include pixel-space sensitivity map methods (Simonyan et al., 2013; Zeiler & Fergus, 2014; Springenberg et al., 2014; Smilkov et al., 2017; Sundararajan et al., 2017) and class-discriminative localization methods (Zhou et al., 2016; Selvaraju et al., 2017; Chattopadhay et al., 2018; Petsiuk et al., 2018). In addition, there has also been a growing body of research focused on data-level saliency analyses, often referred to as data attribution (Koh & Liang, 2017; Park et al., 2023; Ilyas et al., 2022). The application of data attribution includes model explanation (Jeyakumar et al., 2020; Grosse et al., 2023), debugging (Ilyas et al., 2022), efficient training (Xie et al., 2023), and improving model generalization (Jain et al., 2023). Compared to input saliency and data attribution, model saliency is a less explored concept. Weight sparsity (Han et al., 2015; Frankle & Carbin, 2018), commonly used in weight pruning to enhance model efficiency, can be viewed as a form of weight saliency map that focuses on preserving a model’s generalization ability. In the field of natural language processing (NLP), research on model editing (Dai et al., 2021; Meng et al., 2022; De Cao et al., 2021; Patil et al., 2023) has focused on locating and modifying specific knowledge within a model by directly targeting and modifying model weights. This concept of an ‘editable model region’ aligns with the notion of weight saliency in NLP, where certain model parameters are considered more influential and editable than others.

Machine unlearning (MU): Objective and setup. MU has become a vital concept and approach in ML, allowing us to remove the influence of specific data points, data classes, or even higher-level data concepts from a pre-trained ML model without requiring a complete retraining of the model from scratch. The set of data points earmarked for unlearning is commonly known as the forgetting dataset. Thus, the primary objective of MU can be framed as the efficient and effective updating of a pre-trained ML model, so as to attain performance on par with complete retraining (termed as Retrain), which is achieved after the removal of the forgetting dataset from the training set.

To be concrete, let $\mathcal{D}=\{\mathbf{z}_{i}\}_{i=1}^{N}$ denote the training dataset encompassing $N$ data points (including data feature $\mathbf{x}_{i}$ and label $\mathbf{y}_{i}$ for supervised learning). And let $\mathcal{D_{\mathrm{f}}}\subseteq\mathcal{D}$ be the forgetting dataset. Its complement, denoted by $\mathcal{D_{\mathrm{r}}}=\mathcal{D}\setminus\mathcal{D}_{\mathrm{f}}$, is referred to as the remaining dataset. Prior to MU, we denote the original model as ${{\bm{\theta}}_{\mathrm{o}}}$, trained on $\mathcal{D}$ using, e.g., empirical risk minimization (ERM). Consistent with existing literature (Thudi et al., 2022a; Jia et al., 2023), we regard Retrain as the MU’s gold standard, which involves training mode parameters (${\bm{\theta}}$) from scratch over $\mathcal{D}_{\mathrm{r}}$. Nonetheless, Retrain can be computationally demanding. Hence, the central challenge in MU is to acquire an unlearned model (referred to as ${{\bm{\theta}}_{\mathrm{u}}}$) from ${{\bm{\theta}}_{\mathrm{o}}}$ on $\mathcal{D}_{\mathrm{f}}$ and/or $\mathcal{D_{\mathrm{r}}}$, so that it can serve as an accurate and computationally efficient substitute for Retrain. In what follows, we introduce two MU paradigms that are the primary focus of this work: MU for image classification and MU for image generation.

MU for image classification. This is the most commonly studied MU problem in the literature (Shaik et al., 2023). Depending on the composition of the forgetting dataset $\mathcal{D}_{\mathrm{f}}$, MU for image classification can be further categorized into two scenarios: class-wise forgetting and random data forgetting. The former aims to eliminate the influence of an image class, while the latter aims to remove the influence of randomly selected data points from the entire training set.

Evaluating the effectiveness of MU for image classification has involved the use of various metrics. While a consensus is still lacking, we adhere to the recent approach proposed by (Jia et al., 2023), which considers a comprehensive ‘full-stack’ MU evaluation. This includes unlearning accuracy (UA), i.e., $1$ $-$ the accuracy of an unlearned model ${{\bm{\theta}}_{\mathrm{u}}}$ on $\mathcal{D}_{\mathrm{f}}$, membership inference attack (MIA) on $\mathcal{D}_{\mathrm{f}}$, i.e., the privacy measure of ${{\bm{\theta}}_{\mathrm{u}}}$ over $\mathcal{D}_{\mathrm{f}}$, remaining accuracy (RA), i.e., the fidelity of an unlearned model ${{\bm{\theta}}_{\mathrm{u}}}$ on the remaining training set $\mathcal{D}_{\mathrm{r}}$, testing accuracy (TA), i.e., the generalization of ${{\bm{\theta}}_{\mathrm{u}}}$, and run-time efficiency (RTE), i.e., the computation time of applying an MU method.

MU for image generation in conditional diffusion models (DMs). This unlearning problem is emerging given the recent findings that conditional DMs can generate images containing harmful content (e.g., nudity) when provided with inappropriate text prompts (Schramowski et al., 2023). This work will focus on two types of DMs, denoising diffusion probabilistic model (DDPM) with classifier-free guidance (Ho & Salimans, 2022) and latent diffusion model (LDM)-based stable diffusion (Rombach et al., 2022).

We briefly review the diffusion process and DM training. Let ${\epsilon}_{{\bm{\theta}}}(\mathbf{x}_{t}|c)$ symbolize the noise generator parameterized by ${\bm{\theta}}$, conditioned on the text prompt $c$ (e.g., image class in DDPM or text description in LDM, termed as ‘concept’) and structured to estimate the underlying noise (achieved by the reverse diffusion process). Here $\mathbf{x}_{t}$ denotes the data or the latent feature subject to noise injection (attained via forward diffusion process) at the diffusion step $t$. The diffusion process is then given by

where $\hat{\epsilon}(\mathbf{x}_{t}|c)$ stands for the ultimate noise estimation attained by utilizing the conditional DM given $c$, $w\in[0,1]$ is a guidance weight, and ${\epsilon}(\mathbf{x}_{t}|\emptyset)$ signifies the corresponding unconditional employment of the DM. The inference stage initiates with Gaussian noise $z_{T}\sim\mathcal{N}(0,1)$, which is then denoised using $\hat{\epsilon}_{{\bm{\theta}}}(\mathbf{x}_{T}|c)$ to obtain $z_{T-1}$. This procedure is repeated to generate the authentic data at $t=0$. When training the DM ${\bm{\theta}}$, the mean-squared-error (MSE) loss is commonly used

where we omit the expectation over the training data in $\mathcal{D}$ for ease of presentation.

Given a well-trained DM ${\bm{\theta}}$, the objective of MU for image generation is twofold: (1) preventing ${\bm{\theta}}$ from generating undesired image content, e.g., when conditioned on harmful concepts like nudity, and (2) ensuring that the post-unlearning updated DM maintains the quality of image generation for normal images. Finally, it is worth noting that in existing literature, the problem of MU for image generation was not studied through the lens of MU. Instead, it was initially termed as ‘learning to forget’ or ‘concept erasing’ (Schramowski et al., 2023; Gandikota et al., 2023; Zhang et al., 2023a). However, we will show that MU provides a systematic framework for addressing this challenge.

In this section, we highlight two key limitations of current MU methods: the lack of unlearning stability and generality. These limitations underscore the pressing need for a new, robust MU solution, which is inherently non-trivial. We will re-examine 5 MU methods, including ① fine-tuning (FT) that fine-tunes the pre-trained model ${{\bm{\theta}}_{\mathrm{o}}}$ on the remaining dataset $\mathcal{D}_{\mathrm{r}}$ (Warnecke et al., 2021), ② random labeling (RL) that involves fine-tuning ${{\bm{\theta}}_{\mathrm{o}}}$ on the forgetting dataset $\mathcal{D}_{\mathrm{f}}$ using random labels to enforce unlearning (Golatkar et al., 2020), ③ gradient ascent (GA) that reverses the training of ${{\bm{\theta}}_{\mathrm{o}}}$ using gradient ascent on $\mathcal{D_{\mathrm{f}}}$ (Thudi et al., 2022a), ④ influence unlearning (IU) that leverages influence function (Koh & Liang, 2017) to erase the influence of $\mathcal{D}_{\mathrm{f}}$ from ${{\bm{\theta}}_{\mathrm{o}}}$ (Izzo et al., 2021; Jia et al., 2023), ⑤ $\ell_{1}$-sparse MU that infuses weight sparsity into unlearning (Jia et al., 2023).

The instability limitation. In evaluating the performance of MU methods, previous research has often assumed a fixed number of forgetting data, such as data points within an entire class or a fixed ratio of the training set. There has been limited evaluation exploring how the unlearning performance is affected by varying quantities of forgetting data. In Fig. 2(a), we investigate the unlearning performance gap relative to the gold standard Retrain, measured in terms of the average over all the metrics (including UA, RA, TA, and MIA), as a function of the quantity of forgetting data points. Note that a smaller performance gap is desirable. As we can see, the unlearning effectiveness of MU methods (①-⑤) observed at a 10% forgetting data quantity does not necessarily hold when the forgetting data quantity is increased to 50%. Similarly, the instability can also be observed in other performance metrics and will show in experiment results later.

Fig. 2(b) illustrates another form of instability related to the selection of hyperparameters for unlearning methods. Let us take IU (influence unlearning) as an example, where the tuning of the Fisher information regularization parameter is necessary (Izzo et al., 2021; Jia et al., 2023). Given a fixed unlearning scenario that involves forgetting the influence of 10% of CIFAR-10 data in Fig. 2(b), we observe a notably high variance in the unlearning performance of IU compared to Retrain. By contrast, the integration with our proposal (SalUn) reduces this instability.

The generality limitation. Recall that one focus of this work is to develop a principled MU approach that can effectively address MU tasks in both image classification and generation. Before presenting our solution, a ‘must-try’ step is to explore whether classical MU methods developed for image classification can be effectively adapted to MU for image generation. However, we find that existing MU methods do not stay effective. Fig. 3 shows some representative results when using existing MU methods (including GA, RL, FT, and $\ell_{1}$-sparse) as well as Retrain to create an unlearned DM with the goal of preventing the generation of ‘airplane’ images. Existing MU methods tend to either over-forget, resulting in poor generation quality for image classes in $\mathcal{D}_{\mathrm{r}}$ (e.g., GA, RL), or under-forget, leading to unsuccessful unlearning with regard to ‘airplane’ images (e.g., FT, $\ell_{1}$-sparse). This stands in sharp contrast to Retrain, which has the capability to generate unrelated images under the concept of ‘airplane’ while maintaining the quality of image generation for other classes. Yet, Retrain places a significant computational burden on DMs.

Gradient-based weight saliency map. We first illustrate the rationale behind exploring gradient-based weight saliency for MU. Recent evidence suggests that contemporary ML models exhibit modularity characteristics to some extent (Menik & Ramaswamy, 2023). Here modularity refers to the property of a large ML model being decomposable into manageable subparts, each of which can be more easily maintained and updated independently. In particular, weight sparsity (Frankle & Carbin, 2018) has gained recognition as an important driver of modularity, leading to improvements in various aspects of ML, including efficiency (Riquelme et al., 2021), interpretability (Wong et al., 2021), and robustness (Chen et al., 2022b). In the context of MU, weight sparsity has also been harnessed to facilitate the unlearning process, leading to the $\ell_{1}$-sparse unlearning baseline (Jia et al., 2023). However, weight sparsity encounters certain limitations when applied to MU: (1) Determining the appropriate sparse pattern for an ML model (e.g., a DM) can be a challenging task in itself; (2) Even when sparsity is achievable, some applications may not favor delivering a sparse model after MU due to the observed performance decline, as exemplified by the $\ell_{1}$-sparse MU method in Sec. 4.

Building upon the discussion above, we aim to identify an alternative mechanism, distinct from weight sparsity, that can steer MU’s focus towards specific model weights deemed salient to MU. Drawing inspiration from gradient-based input saliency maps (Smilkov et al., 2017; Adebayo et al., 2018), we pose the question of whether a weight saliency map can be constructed to aid MU. This concept allows us to decompose the pre-unlearning model weights (${{\bm{\theta}}_{\mathrm{o}}}$) into two distinct components: the salient model weights earmarked for updating during MU and the intact model weights that remain unchanged. Similar to input saliency map, we utilize the gradient of a forgetting loss (denoted as $\ell_{\mathrm{f}}({\bm{\theta}};\mathcal{D}_{\mathrm{f}})$) with respect to the model weights variable ${\bm{\theta}}$ under the forgetting dataset $\mathcal{D}_{\mathrm{f}}$. By applying a hard thresholding operation, we can then obtain the desired weight saliency map:

where $\mathds{1}(\mathbf{g}\geq\gamma)$ is an element-wise indicator function which yields a value of $1$ for the $i$-th element if $g_{i}\geq\gamma$ and $0$ otherwise, $|\cdot|$ is an element-wise absolute value operation, and $\gamma>0$ is a hard threshold. In practice, we have observed that setting $\gamma$ to the median of the gradient vector $\nabla_{{\bm{\theta}}}\ell_{\mathrm{f}}({\bm{\theta}};\mathcal{D}_{\mathrm{f}})\left.\right|_{{\bm{\theta}}={{\bm{\theta}}_{\mathrm{o}}}}$ is a sufficiently effective choice. Based on (3), we explicitly express the unlearning model ${{\bm{\theta}}_{\mathrm{u}}}$ as

where $\odot$ is element-wise product, and $\mathbf{1}$ denotes an all-one vector. The implication from (4) is that during weight updating in MU, the attention can be directed towards the salient weights.

It is worth noting that the forgetting loss $\ell_{\mathrm{f}}$ used in existing MU methods can be considered a suitable candidate for calculating the weight saliency map (3). In this study, we find that the forgetting loss in GA (gradient ascent) (Thudi et al., 2022a) presents an effective and simple solution in image classification and generation:

where $\ell_{\mathrm{CE}}$ is the cross-entropy (CE) loss for supervised classification, and $\ell_{\mathrm{MSE}}$ has been defined in DM training (2). The weight saliency map for MU can be then obtained through (3) and (5).

Saliency-based unlearning (SalUn). Next, we introduce SalUn, which incorporates (4) into the unlearning process. One advantage of SalUn is its plug-and-play capability, allowing it to be applied on top of existing unlearning methods. In particular, we find that integrating weight saliency with the RL (random labeling) method provides a promising MU solution; See the ablation study in Table A1.

In image classification, RL assigns a random image label to a forgetting data point and then fine-tunes the model on the randomly labeled $\mathcal{D}_{\mathrm{f}}$. In SalUn, we then leverage RL to update the salient weights in (4). This yields the optimization problem associated with SalUn for image classification:

where $y^{\prime}$ is the random label of the image $\mathbf{x}$ different from $y$, and ${\bm{\theta}}_{\mathrm{u}}$ has been defined in (4). Additionally, to achieve a balance between unlearning on forgetting data points and preserving the model’s generalization ability for non-forgetting data points, the regularization term on $\mathcal{D}_{\mathrm{r}}$ preserved, with $\alpha>0$ as a regularization parameter.

Furthermore, we extend the use of RL to the image generation context within SalUn. In this context, RL is implemented by associating the forgetting concept, represented by the prompt condition $c$ in (2), with a misaligned image $\mathbf{x}^{\prime}$ that does not belong to the concept $c$. To maintain the image-generation capability of the DM, we also introduce the MSE loss (2) on the remaining dataset $\mathcal{D}_{\mathrm{r}}$ as a regularization. This leads to the optimization problem of SalUn for image generation:

where $c^{\prime}\neq c$ indicates that the concept $c^{\prime}$ is different from $c$ , ${{\bm{\theta}}_{\mathrm{u}}}$ is the saliency-based unlearned model given by (4), $\beta>0$ is a regularization parameter similar to $\alpha$ in (6), to place an optimization tradeoff between the RL-based unlearning loss over the forgetting dataset $\mathcal{D}_{\mathrm{f}}$ and the diffusion training loss $\ell_{\mathrm{MSE}}({{\bm{\theta}}_{\mathrm{u}}};\mathcal{D}_{\mathrm{r}})$ on the non-forgetting dataset $\mathcal{D}_{\mathrm{r}}$ (to preserve image generation quality). Similar to (6), SalUn begins with the pre-trained model ${{\bm{\theta}}_{\mathrm{o}}}$ and follows the optimization in (7) to accomplish unlearning. See Appendix A for the algorithmic implementations.

Extension on ‘soft-thresholding’ SalUn. The implementation of SalUn relies on a pre-selected hard threshold to determine the weight saliency map $\mathbf{m}_{\mathrm{S}}$ in (3). While this hard-thresholding approach performs well, we can also develop an alternative implementation of SalUn that employs soft thresholding and may allow for a more flexible saliency map determination; See Appendix B for more algorithmic details. Yet, in practice, the soft-thresholding variant does not outperform the hard-thresholding version. Thus, we will focus on the hard-thresholding implementation by default.

Recognizing the shortcomings of existing MU approaches, we introduce the innovative concept of weight saliency in MU, leading to the SalUn framework. This proposed saliency unlearning has showcased its effectiveness in addressing the limitations of current MU methods, and applying to both image classification and generation tasks. As a notable application, we show that SalUn stays effective in preventing stable diffusion from generating images with harmful content, even suffering inappropriate image prompts.

C. Fan, J. Liu, and S. Liu were supported by the Cisco Research Faculty Award and the National Science Foundation (NSF) Robust Intelligence (RI) Core Program Award IIS-2207052. We would like to thank Jinghan Jia for the insightful discussions.