ConceptPrune: Concept Editing in Diffusion Models via Skilled Neuron Pruning

Authors: Ruchika Chavhan, Da Li, Timothy Hospedales
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: diffusion models, concept editing, pruning

Abstract:
While large-scale text-to-image diffusion models have demonstrated impressive image-generation capabilities, there are significant concerns about their potential misuse for generating unsafe content, violating copyright, and perpetuating societal biases. Recently, the text-to-image generation community has begun addressing these concerns by editing or unlearning undesired concepts from pre-trained models. However, these methods often involve data-intensive and inefficient fine-tuning or utilize various forms of token remapping, rendering them susceptible to adversarial jailbreaks. In this paper, we present a simple and effective training-free approach, ConceptPrune, wherein we first identify critical regions within pre-trained models responsible for generating undesirable concepts, thereby facilitating straightforward concept unlearning via weight pruning. Experiments across a range of concepts including artistic styles, nudity, and object erasure demonstrate that target concepts can be efficiently erased by pruning a tiny fraction, approximately 0.12% of total weights, enabling multi-concept erasure and robustness against various white-box and black-box adversarial attacks.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 3

Reviews and discussion:
<thread>
creation date:
1730601427464

thread_type:
official_review

<thread_object>
creator: reviewer_udze

<content>
summary:
This paper addresses the risks of large-scale text-to-image diffusion models, including misuse for unsafe content, copyright violations, and bias perpetuation. Traditional methods for removing undesired concepts from these models often require intensive fine-tuning or token remapping, which remain vulnerable to adversarial jailbreaks. To tackle these issues, the authors propose ConceptPrune, a training-free approach that locates key regions in pre-trained models linked to unwanted concepts and applies weight pruning for efficient concept unlearning. Experiments show ConceptPrune effectively erases specific styles, nudity, and objects by pruning only 0.12% of weights, allowing multi-concept removal with resilience to adversarial attacks.

strengths:
1. The paper is clearly written and well-structured. 
2. It offers an innovative approach to unlearning by employing model pruning techniques.

weaknesses:
1. There is a lack of adequate baselines. The authors only compare their method to a few early baselines in machine unlearning within diffusion models. As far as I know, many more effective methods have emerged recently, but these were not included in the comparison.

2. The robustness of the method is insufficiently studied. According to recent research, machine unlearning can be vulnerable to adversarial prompts. Without a robust evaluation, it remains unclear whether pruning-based unlearning is indeed resilient against such attacks.

questions:
1. Could the authors add more baselines by comparing their method with recently published machine unlearning methods in diffusion models?

2. Could the authors include experiments on the robustness of machine unlearning? Specifically, experiments involving adversarial prompts or similar robustness tests would be valuable.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680826763

thread_type:
official_review

<thread_object>
creator: reviewer_yhcm

<content>
summary:
The paper introduces ConceptPrune, a method for editing concepts in pre-trained diffusion models by pruning neurons associated with undesired concepts, like specific artistic styles, nudity, and biases. This approach identifies "skilled neurons" within diffusion models, such as Stable Diffusion, that activate in response to specific target concepts. By selectively pruning these neurons, the method effectively removes the unwanted concepts without the need for extensive fine-tuning, maintaining the model's overall performance on unrelated tasks. ConceptPrune demonstrates significant resistance to both black-box and white-box adversarial attacks, showing promise as a scalable and robust approach to safer model deployment.

strengths:
Strengths:
- The investigation into how specific weights influence certain abilities in text-to-image (T2I) models is intriguing and provides valuable insights into the inner workings of these models.
- The paper presents extensive experiments demonstrating the effectiveness of the proposed method across various editing targets, including nudity, artistic style, and specific objects.
- Additionally, the paper is well-structured with a clear logical flow, making it easy to follow and understand.

weaknesses:
- **Limited Comparison with Recent SOTA Methods**: The paper lacks sufficient discussion and comparison with recent state-of-the-art methods. For instance, [1] appears to be a strong baseline, demonstrating both high robustness in concept erasure and good utility preservation. While it’s acceptable if this paper does not outperform [1], it would be beneficial to include [1] as a baseline and discuss potential improvements for the current approach in light of their findings.

- **Insufficient Discussion on Related Concept Pruning Work**: The paper lacks sufficient background on concept pruning and similar work. Several studies ([2, 3, 4]) have explored the relationship between weight importance and model abilities during unlearning and editing tasks. Incorporating these works would help situate this paper within the broader landscape and highlight how ConceptPrune builds on or differs from these existing approaches.

[1] Zhang, Yimeng, et al. "Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models." arXiv preprint arXiv:2405.15234 (2024).
[2] Fan, Chongyu, et al. "Salun: Empowering machine unlearning via gradient-based weight saliency in both image classification and generation." arXiv preprint arXiv:2310.12508 (2023).
[3] Wu, Jing, and Mehrtash Harandi. "Scissorhands: Scrub Data Influence via Connection Sensitivity in Networks." arXiv preprint arXiv:2401.06187 (2024).
[4] Foster, Jack, Stefan Schoepf, and Alexandra Brintrup. "Fast machine unlearning without retraining through selective synaptic dampening." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 11. 2024.

questions:
1. Could you please add discussions on the recent SOTA methods and relevant baselines as highlighted in the weaknesses section?

2. How did you determine the optimal pruning ratio for skilled neurons? Please describe the considerations or criteria used in selecting a suitable pruning ratio.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730843574047

thread_type:
official_review

<thread_object>
creator: reviewer_pdbh

<content>
summary:
To prevent the misuse of text-to-image diffusion models, this paper proposes a training-free approach using weight pruning for unlearning undesired concepts, such as artistic styles, nudity, and specific objects.

strengths:
1. The experiments cover all diffusion model unlearning tasks, including artistic styles, nudity, and specific objects.
2. Consider ASR against different attacks while simultaneously evaluating FID and CLIP score to assess model utility.
3. The proposed method requires no additional training.

weaknesses:
1. The chosen diffusion model baselines are weak and not state-of-the-art methods, so the comparison does not effectively demonstrate true superiority. To defend against adversarial prompt attacks, a stronger baseline, such as AdvUnlearn [1], should be considered.
2. More visualization examples are needed, as the current version only includes visualizations for the style unlearning task.

[1]  Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models, NeurIPS’24

questions:
In Table 5, the FID metric is missing, and the performance of the base model (SD) should be included for greater clarity.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730870258414

thread_type:
official_review

<thread_object>
creator: reviewer_b4vy

<content>
summary:
The paper introduces ConceptPrune, a training-free method for concept editing in pre-trained diffusion models, specifically latent diffusion models like Stable Diffusion. The core idea is to identify and prune "skilled neurons" within the feed-forward networks (FFNs) of the model's UNet architecture that are responsible for generating undesirable concepts. By calculating importance scores based on neuron activations for target (undesired) and reference (desired) prompts, the method isolates neurons that predominantly influence the generation of unwanted content. Pruning these neurons effectively erases the target concepts while preserving the model's ability to generate unrelated content. The authors demonstrate the efficacy of ConceptPrune across various concepts, including artistic styles, nudity, object erasure, and gender biases, showing that pruning approximately 0.12% of the model's weights suffices. Additionally, the method exhibits robustness against both white-box and black-box adversarial attacks aimed at circumventing concept removal.

strengths:
1. The paper introduces an application of neuron pruning techniques to concept editing in diffusion models. Specifically, ConceptPrune utilizes the identification and pruning of skilled neurons, providing an alternative approach to mitigating undesired content generation.

2. The experimental results support the method's effectiveness in erasing various concepts while maintaining image generation quality. The authors compare ConceptPrune with several baselines, demonstrating its performance in concept removal and robustness against adversarial attacks.

3. The manuscript is well-written and well-structured, offering clear explanations of the methodology and detailed descriptions of the experiments. Mathematical formulations are provided to illustrate the approach, and visual examples enhance comprehension.

4. The work addresses significant ethical concerns related to the generation of unsafe or copyright-infringing content in diffusion models.

weaknesses:
1. The method seems relatively straightforward and may be considered incremental in light of existing works [1,2] that also employ pruning techniques for unlearning in diffusion models. It would strengthen the contribution of the paper to clearly delineate the differences between the proposed approach and these related methods. Specifically, methods such as [1,3] have achieved state-of-the-art results in diffusion model unlearning tasks. Including these methods as baselines in the experimental comparisons would provide a more comprehensive evaluation and help illustrate the effectiveness of the proposed scheme.

2. The experimental validation is conducted solely on Stable Diffusion, which may limit the generalizability of the proposed method. To address this potential limitation, it is recommended that the effectiveness of the approach be demonstrated on additional diffusion models, such as [4,5]. This would provide evidence that the method is not tightly coupled to a specific architecture and can be broadly applied to other models in the domain.

3. The choice of focusing on the second layer of the feed-forward networks (FFN-2) for pruning appears to be based on empirical observations and may depend significantly on the architecture of Stable Diffusion. Is there theoretical justification or empirical evidence supporting this selection as the optimal pruning target? Providing analytical insights or additional experiments on different architectures would help ascertain whether this observation holds true universally or is specific to the model under consideration.

4. The decision to aggregate skilled neurons over the last 10 timesteps for pruning seems somewhat arbitrary. Could the authors elaborate on the reasoning behind selecting this particular range of timesteps? It may be beneficial to reference insights from related literature [6,7] that discuss the importance of different timesteps in the denoising process. A more thorough explanation or an exploration of how varying the number of timesteps affects the results would enhance the understanding of this methodological choice.

5. The evaluation of unlearning effectiveness in the context of artistic style erasure relies on CLIP-based similarity metrics. However, CLIP-based metrics consider redundant factors, not only the style but also the object content, arrangement, etc. It is potentially confounding the assessment of style erasure alone. This approach may not effectively isolate the stylistic elements from other irrelevant factors. It is recommended to employ a dedicated style classifier for evaluation (e.g., classifier in [8]) to more accurately measure the degree of style removal. This would provide a more reliable and focused assessment of the method's effectiveness in erasing artistic styles.

> [1] Wu, Jing, and Mehrtash Harandi. "Scissorhands: Scrub Data Influence via Connection Sensitivity in Networks." arXiv preprint arXiv:2401.06187 (2024).
>
> [2] Yang, Tianyun, Juan Cao, and Chang Xu. "Pruning for Robust Concept Erasing in Diffusion Models." arXiv preprint arXiv:2405.16534 (2024).
>
> [3] Fan, Chongyu, et al. "Salun: Empowering machine unlearning via gradient-based weight saliency in both image classification and generation." arXiv preprint arXiv:2310.12508 (2023).
>
> [4] Podell, Dustin, et al. "Sdxl: Improving latent diffusion models for high-resolution image synthesis." arXiv preprint arXiv:2307.01952 (2023).
>
> [5] Saharia, Chitwan, et al. "Photorealistic text-to-image diffusion models with deep language understanding." Advances in neural information processing systems 35 (2022): 36479-36494.
>
> [6] Balaji, Yogesh, et al. "ediff-i: Text-to-image diffusion models with an ensemble of expert denoisers." arXiv preprint arXiv:2211.01324 (2022).
>
> [7] Georgiev, Kristian, et al. "The journey, not the destination: How data guides diffusion models." arXiv preprint arXiv:2312.06205 (2023).
>
> [8] Zhang, Yihua, et al. "Unlearncanvas: A stylized image dataset to benchmark machine unlearning for diffusion models." arXiv preprint arXiv:2402.11846 (2024).

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732259277622

thread_type:
public_comment

<thread_object>
creator: ~finn_carter1

<content>
title:
Lack of recent related works

comment:
It seems that several recent related works [1,2,3] are ignored.

[1] One-dimensional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications

[2] MACE: Mass Concept Erasure in Diffusion Models

[3] Receler: Reliable Concept Erasing of Text-to-Image Diffusion Models via Lightweight Erasers
</content>
</thread_object>

</thread>
<thread>
creation date:
1732667841793

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common response

comment:
We thank all the reviewers for their insightful and positive feedback. We are pleased that they found our paper reasonable and acknowledged that our contributions extend beyond concept editing, including training-free erasure and enhanced adversarial robustness. We have revised the paper accordingly with changes marked in blue. We address reviewer comments below and incorporate all feedback. 

**Comparison to more baselines (B4VY, PDbh, YHCm, uDZE, and public comment by Finn Carter)**


We appreciate the reviewers’ suggestions regarding additional comparison to recent state-of-the-art (SOTA) methods. In response, we have incorporated the following additional baselines into the paper – Selective Amnesia (SA) [1], SPM [2], Receler [3], MACE [4], Scissorhands[5], AdvUnlearn [6], and SalUn[8], extending the comparison to a total of 11 SOTA methods. However, we would like to argue that some of these are not direct competitors to ConceptPrune.

While the baselines considered in the paper are highly effective, deploying current state-of-the-art concept erasure techniques in real-world scenarios poses significant challenges, particularly in online environments with computational constraints where harmful concepts can emerge dynamically. This is because these methods struggle to meet the following requirements for real-world applications: (1) training-free concept erasure, eliminating concepts without the need for backpropagation through the entire model, or (2) lightweight or fast concept erasure, allowing concepts to be removed quickly and efficiently with minimal compute.

In Table A, we present a comprehensive summary of related works, categorizing them based on whether they are training-free and lightweight for an online setting. For instance, FMN and UCE are considered lightweight as they can achieve concept erasure in approximately 35 and 120 seconds. In contrast, ESD is significantly more time-intensive, making it less practical for rapid applications. 

Our proposed solution, ConceptPrune, excels on both fronts by introducing a training-free, pruning-based approach that eliminates harmful concepts without updating any parameters. Instead, it identifies and targets the neurons responsible for generating these concepts enabling efficient concept erasure (in approximately 150 seconds) with significantly reduced computational requirements. Therefore, we consider the methods categorized as lightweight concept erasers to be direct competitors to ConceptPrune. More details in Section 2 (Related Work)

| **Method**          | **Training-free** | **Parameters Trained**       | **Lightweight Erasure** |
|----------------------|-------------------|------------------------------|--------------------------|
| CA            | N                 | Full denoiser                | N                        |
| SA [1]              | N                 | Full denoiser                | N                        |
| SH [5]                | N                 | Full denoiser                | N                        |
| AdvUnlearn [6]     | N                 | Full denoiser                | N                        |
| SalUn [8]            | N                 | Full denoiser                | N                        |
| ESD         | N                 | Cross Attention              | N                        |
| Receler [3]        | N                 | Cross Attention              | Y                        |
| FMN            | N                 | Cross Attention              | Y                        |
| SPM [2]              | N                 | LORA                        | Y                        |
| MACE [4]              | N                 | Cross Attention + LORA       | Y                        |
| UCE     | Y                 | Cross Attention              | Y                        |
| **Ours (ConceptPrune)**             | **Y**             | **_None_**                   | **Y**                    |

**Table A - Summary of recent Concept Erasure baselines. ConceptPrune is a training-free approach that enables rapid pruning of the model to eliminate a new target concept without the need for extensive re-training.**


[1] Selective Amnesia: A Continual Learning Approach to Forgetting in Deep Generative Models
[2] One-dimensional Adapter to Rule Them All: Concepts, Diffusion Models and Erasing Applications, Lyu et al
[3] Receler: Reliable Concept Erasing of Text-to-Image Diffusion Models via Lightweight Erasers, Huang et al
[4] MACE: Mass Concept Erasure in Diffusion Models, Lu et al
[5] Scissorhands: Scrub Data Influence via Connection Sensitivity in Networks, Wu et al
[6] Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models, Zhang et al
[7] Pruning for Robust Concept Erasing in Diffusion Models, Yang et al
[8] Salun: Empowering machine unlearning via gradient-based weight saliency in both image classification and generation. Fan, Chongyu, et al.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors provided a detailed response that most reviewers found useful, as it partially or mostly addressed their concerns. Reviewer uDZE, who gave the lowest rating, did not actively participate in the discussion phase. The remaining reviewers maintained a positive stance on this submission. I agree that the paper's strengths significantly outweigh its weaknesses, and I therefore recommend acceptance.