Source-linked AI summary
SalUn: Empowering Machine Unlearning via Gradient-based Weight Saliency in Both Image Classification and Generation
Chongyu Fan, Jiancheng Liu, Yihua Zhang, Eric Wong, Dennis Wei, Sijia Liu
TL;DR
Existing machine-unlearning methods have limitations in stability, accuracy, and cross-domain applicability. SalUn introduces gradient-based weight saliency to target specific model weights, narrowing the gap with exact unlearning across classification and generation, including harmful-content removal. Its scope beyond vision remains an open question.
Problem
Existing MU methods face limitations in stability and cross-domain applicability, while image-generation MU must suppress harmful content without reducing normal-image quality.
Method
SalUn uses weight saliency to guide unlearning toward specific model weights and balances forgetting with preservation of non-forgotten data or concepts.
Results
SalUn is effective across image classification and generation, achieving the best trade-off among unlearning efficacy, preserved fidelity, and computational efficiency in the reported comparison.
Takeaways & Limitations
SalUn provides a principled MU approach for forgetting data, classes, and concepts across vision tasks, including harmful-content prevention in stable diffusion.
Takeaways & Limitations
SalUn’s scalability and adaptability beyond vision, including language and graph domains, require further investigation.
Abstract
from arXiv · showhide
With evolving data regulations, machine unlearning (MU) has become an important tool for fostering trust and safety in today's AI models. However, existing MU methods focusing on data and/or weight perspectives often suffer limitations in unlearning accuracy, stability, and cross-domain applicability. To address these challenges, we introduce the concept of 'weight saliency' for MU, drawing parallels with input saliency in model explanation. This innovation directs MU's attention toward specific model weights rather than the entire model, improving effectiveness and efficiency. The resultant method that we call saliency unlearning (SalUn) narrows the performance gap with 'exact' unlearning (model retraining from scratch after removing the forgetting data points). To the best of our knowledge, SalUn is the first principled MU approach that can effectively erase the influence of forgetting data, classes, or concepts in both image classification and generation tasks. As highlighted below, For example, SalUn yields a stability advantage in high-variance random data forgetting, e.g., with a 0.2% gap compared to exact unlearning on the CIFAR-10 dataset. Moreover, in preventing conditional diffusion models from generating harmful images, SalUn achieves nearly 100% unlearning accuracy, outperforming current state-of-the-art baselines like Erased Stable Diffusion and Forget-Me-Not. Codes are available at https://github.com/OPTML-Group/Unlearn-Saliency. (WARNING: This paper contains model outputs that may be offensive in nature.)
1 INTRODUCTION
Machine unlearning removes specific data influence from pretrained models, but approximate methods can be unstable and classification methods may not transfer effectively to image generation. SalUn addresses these limitations through weight saliency and is presented as effective across both domains.
- Machine unlearning mitigates the influence of particular data points on a pretrained model.
- Exact or certified unlearning offers guarantees, whereas approximate unlearning prioritizes practical speed and accuracy without requiring certified-unlearning assumptions.
- Approximate unlearning can exhibit substantial performance variance, including sensitivity to hyperparameter selection.
- Existing MU methods developed for image classification are insufficient for image-generation unlearning.
- SalUn introduces weight saliency to reduce the gap with exact unlearning and support MU in image classification and generation.
2 RELATED WORK
Related work covers exact and approximate unlearning for classification, concept erasure in diffusion generation, and saliency analyses spanning inputs, data, and model weights.
- Unlearning in image classification: Exact unlearning retrains models after removing data, but its computational demands motivate probabilistic and approximate alternatives for image classification.
- Unlearning in image generation: Diffusion-model studies address concept erasure, but achieving precision comparable to exact unlearning remains challenging.
- Data and model saliency analyses: Prior saliency research emphasizes input sensitivity and data attribution, while model saliency and weight-focused analysis remain less explored.
3 PRELIMINARIES AND PROBLEM STATEMENT
The paper frames MU as efficiently updating a pretrained model to match retraining after forgetting selected data, classes, or concepts. For image generation, the update must suppress undesired content while preserving normal-image quality.
- MU objective and setup: The forgetting dataset contains points designated for removal, while the remaining dataset is its complement.
- MU objective and setup: Retrain is the gold standard, but MU seeks an accurate and computationally efficient substitute obtained from the original model.
- MU for image classification: Classification MU includes class-wise forgetting and random data forgetting.
- MU for image classification: Full-stack evaluation considers unlearning accuracy, membership inference, remaining accuracy, testing accuracy, and runtime efficiency.
- MU for image generation: Conditional diffusion models estimate noise from noisy data or latent features during reverse diffusion, and training commonly uses mean-squared-error loss.
- MU for image generation: Image-generation MU must prevent undesired content while maintaining normal-image generation quality.
4 CHALLENGES IN CURRENT MACHINE UNLEARNING METHODS
The paper identifies instability and limited generality as central weaknesses of current MU methods. Experiments show sensitivity to forgetting-data quantity and hyperparameters, while classification methods can fail in image generation.
- Challenges: Current MU methods face limitations in unlearning stability and generality.
- The instability limitation: Performance gaps relative to Retrain can change when the forgetting-data quantity increases from 10% to 50%.
- The instability limitation: Influence unlearning shows high variance across hyperparameter choices, whereas SalUn integration reduces this instability.
- The generality limitation: Existing classification MU methods do not remain effective when adapted to image-generation unlearning.
5 SALUN: WEIGHT SALIENCY IS POSSIBLY ALL YOU NEED FOR MU
SalUn constructs a gradient-based weight saliency map and updates salient weights during unlearning, combining this mechanism with random labeling and regularization for classification and generation.
- Weight saliency: Weight saliency uses gradients of a forgetting loss to identify model weights for targeted unlearning.Hard thresholding converts the gradient magnitudes into a saliency map.
- Weight update: The unlearned model updates salient weights while leaving intact weights at their original values.The model is expressed as a masked combination of updated and original weights.
- Image classification: SalUn combines the saliency-based weight update with random labeling for image classification.The random-labeling loss updates salient weights while a regularization term on retained data preserves generalization.
- Method integration: SalUn is designed as a plug-and-play method that can be applied on top of existing unlearning methods.The paper specifically reports integrating weight saliency with random labeling as a promising solution.
- Image generation: For image generation, SalUn misaligns forgetting concepts with alternative images and adds diffusion-training loss on retained data.The regularization balances concept unlearning against preserving image-generation quality.
- Implementation and evaluation: The default implementation uses hard thresholding, although the paper also considers a soft-thresholding variant.The supplied table caption describes evaluation across two random-forgetting scenarios on CIFAR-10 with ResNet-18.
6 EXPERIMENTS
Experiments evaluate SalUn across image classification and generation, comparing it with baselines and exact retraining using unlearning, fidelity, and efficiency measures. Results show that weight saliency supports stable classification unlearning and effective image-generation forgetting while preserving generation quality.
- Experimental setups: Experiments cover random-data and class-wise forgetting in image classification, plus class-wise and concept-wise forgetting in diffusion-based image generation.Classification uses CIFAR-10, CIFAR-100, SVHN, Tiny ImageNet, ResNet-18, VGG-16, and Swin-T; generation uses DDPM and stable diffusion.
- Image classification: SalUn achieves the smallest average performance gap with Retrain in both random-data forgetting scenarios.SalUn-soft has the second or third-smallest gap, while hard thresholding is hypothesized to improve unlearning efficacy through strict sparsity.
- Image classification: SalUn offers the best trade-off between unlearning efficacy, preserved model fidelity, and computational efficiency across the evaluated classification metrics.The comparison uses UA and MIA for unlearning, RA and TA for retained-model fidelity, and RTE for efficiency.
- Image classification: 50% forgetting increases the Avg. Gap for BS from 2.67 to 8.13 and for ℓ1-sparse from 2.26 to 5.72, while SalUn stays consistently effective.The results indicate that increasing the forgotten-data proportion makes unlearning more challenging for all methods.
- Image generation: Proper weight saliency outperforms random masking for DDPM generation by avoiding noisy over-forgetting and degradation on non-forgetting classes.Random masking produces noisy airplane-condition outputs and lower-quality outputs for non-forgetting examples such as C2, C3, C6, and C7.
- Image generation: On Imagenette, SalUn outperforms ESD and FMN in UA across forgetting classes while maintaining good FID, whereas FMN has lower FID but weaker forgetting.For NSFW concept forgetting, SalUn generates the fewest harmful images across all evaluated nude body-part classes and outperforms ESD especially for male and female breast categories.
7 CONCLUSION
The paper presents SalUn as a weight-saliency-based unlearning framework, with classification and generation variants and proximal-gradient extensions that regulate model-weight changes.
- SalUn applies weight saliency to focus unlearning on selected model weights rather than the entire model.The framework is presented for both image classification and generation tasks.
- The classification and generation algorithms compute a saliency mask from gradients and update parameters using masked gradients.The pseudocode uses relabeled forgetting data, a thresholded saliency mask, and iterative SGD updates.
- An ℓ1 penalty and proximal-gradient steps constrain changes between the unlearned model θu and the original model θo.The proximal update implements soft thresholding around the original weights.
- Higher weight-saliency sparsity corresponds to fewer changed model weights and can be enforced through the regularization parameter β.Increasing β drives θ′ − θo toward zero.
- The soft-thresholding solution uses the positive-part operator to project parameter differences and then maps back to the original variable θ.The operation is applied elementwise to the displacement from θo.
C.1 ADDITIONAL TRAINING AND UNLEARNING SETTINGS
The appendix specifies training settings for classification, DDPM, and Stable Diffusion experiments, and examines how saliency sparsity affects forgetting and generation quality.
- Additional training and unlearning settings: Classification baselines use distinct training schedules, including 182 epochs for Retrain, 10 epochs for FT and RL, and 5 epochs for GA.The settings also search IU’s Hessian approximation parameter α over [1, 20].
- Additional training and unlearning settings: SalUn uses 1,000 DDPM iterations and 5 Stable Diffusion epochs, with 50% weight-saliency sparsity in both settings.The DDPM and Stable Diffusion configurations use different optimizers, batch sizes, sampling procedures, and β values.
- Additional training and unlearning settings: NSFW concept forgetting constructs forgetting and retaining image sets from nude and clothed prompts, then corrects the forgotten concept using the clothed concept.The nudity prompt also supplies the saliency mask during unlearning.
- Additional training and unlearning settings: In Fig. A1, higher saliency sparsity can produce under-forgetting, while lower sparsity can produce over-forgetting or increased FID.Points above Retrain indicate over-forgetting and points below Retrain indicate under-forgetting.
- Additional training and unlearning settings: Forgetting the airplane class at 10% sparsity reduces image-generation quality, whereas 90% sparsity fails to forget the class.Fig. A3 relates the optimal sparsity level to the amount of forgetting data.
C.2 ADDITIONAL CLASSIFICATION RESULTS
Additional classification experiments evaluate saliency integration, class-wise and iterative forgetting, multiple datasets, and multiple architectures, generally emphasizing balanced performance and closeness to Retrain.
- Additional datasets and architectures: Fig. A3 examines how the optimal saliency sparsity changes with the amount of forgotten data under random forgetting on CIFAR-10.
- Saliency-enhanced baselines: Adding the weight-saliency mask improves FT, RL, GA, and IU baselines toward the Retrain benchmark, with saliency-augmented RL performing best among them.The comparison is reported using Avg. Gap across forgetting scenarios.
- Class-wise and iterative forgetting: SalUn and SalUn-soft maintain a robust balance across UA, MIA, RA, and TA for class-wise forgetting on CIFAR-10.The methods do not always achieve the highest UA, but their overall metric balance is emphasized.
- Class-wise and iterative forgetting: SalUn achieves the smallest Avg. Gap to Retrain during five iterative 10%-data forgetting steps totaling 50% of the dataset.The comparison includes FT and Retrain on ResNet-18 trained on CIFAR-10.
- Additional datasets and architectures: Across SVHN, CIFAR-100, VGG-16, and Swin-T settings, SalUn and SalUn-soft deliver promising or balanced performance across diverse forgetting amounts.The reported evaluations emphasize adaptability rather than universal dominance on every metric.
- Additional datasets and architectures: On Tiny ImageNet, SalUn has smaller UA and MIA gaps than ℓ1-sparse while maintaining comparable RA and TA gaps.The comparison uses ResNet-18 with 10% random data forgetting.
C.3 ADDITIONAL GENERATION RESULTS
Additional generation experiments compare SalUn with ESD and examine class-wise forgetting across CIFAR-10 and ImageNette, including robustness across random seeds.
- CIFAR-10 generation: On CIFAR-10 DDPM class-wise forgetting, SalUn has a 0.84% UA drop relative to ESD but significantly outperforms ESD in FID.The authors attribute ESD’s 100% UA partly to poor generation quality and report SalUn as closest to Retrain.
- CIFAR-10 generation: SalUn’s CIFAR-10 generation results are illustrated with class-wise examples, marking the forgotten class with a red border.
- ImageNette generation: ImageNette figures contrast the unlearned class with the prompt class in a matrix, with diagonal images corresponding to the target unlearning class.Off-diagonal images represent other classes and broader differentiation behavior.
- NSFW prompt evaluation: Table A11 supplies the harmful I2P text prompts used for Stable Diffusion generation experiments.
D BROADER IMPACTS AND LIMITATIONS
SalUn is presented as improving unlearning effectiveness while preserving model utility and reducing the risk of harmful image generation. The paper also identifies open questions about extending the method beyond vision and assessing broader societal impacts.
- Broader impacts: SalUn is described as enhancing unlearning effectiveness while maintaining model utility under strict unlearning requirements.
- Limitations: The method’s demonstrated scope is vision, while scalability and adaptability to language and graph domains remain unresolved.
- Limitations: The paper states that unlearning’s effects on fairness, privacy, and security require further consideration alongside transparent, accountable, and inclusive development.
- Additional evaluations: Supplementary evaluations report machine-unlearning performance across varying forgetting-data amounts on CIFAR-10 ResNet-18, SVHN, CIFAR-100, VGG-16, and Swin-T.
- Additional evaluations: Supplementary figures show class-wise CIFAR-10 DDPM results and SalUn-generated images, distinguishing forgetting classes from remaining classes.