[Paper Review] Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models
This paper proposes AdvUnlearn, a novel robust unlearning framework that integrates adversarial training into diffusion model (DM) concept erasure to defend against adversarial prompt attacks. By applying adversarial training to the text encoder with a utility-retaining regularization on a retained prompt set, AdvUnlearn achieves strong robustness (e.g., 21.13% attack success rate on SD v1.4) while preserving high image generation quality (FID ≈ 19.34), outperforming existing baselines in both robustness and utility.
Diffusion models (DMs) have achieved remarkable success in text-to-image generation, but they also pose safety risks, such as the potential generation of harmful content and copyright violations. The techniques of machine unlearning, also known as concept erasing, have been developed to address these risks. However, these techniques remain vulnerable to adversarial prompt attacks, which can prompt DMs post-unlearning to regenerate undesired images containing concepts (such as nudity) meant to be erased. This work aims to enhance the robustness of concept erasing by integrating the principle of adversarial training (AT) into machine unlearning, resulting in the robust unlearning framework referred to as AdvUnlearn. However, achieving this effectively and efficiently is highly nontrivial. First, we find that a straightforward implementation of AT compromises DMs' image generation quality post-unlearning. To address this, we develop a utility-retaining regularization on an additional retain set, optimizing the trade-off between concept erasure robustness and model utility in AdvUnlearn. Moreover, we identify the text encoder as a more suitable module for robustification compared to UNet, ensuring unlearning effectiveness. And the acquired text encoder can serve as a plug-and-play robust unlearner for various DM types. Empirically, we perform extensive experiments to demonstrate the robustness advantage of AdvUnlearn across various DM unlearning scenarios, including the erasure of nudity, objects, and style concepts. In addition to robustness, AdvUnlearn also achieves a balanced tradeoff with model utility. To our knowledge, this is the first work to systematically explore robust DM unlearning through AT, setting it apart from existing methods that overlook robustness in concept erasing. Codes are available at: https://github.com/OPTML-Group/AdvUnlearn
Motivation & Objective
- To address the vulnerability of concept-erased diffusion models to adversarial prompt attacks that bypass unlearning.
- To develop a robust unlearning framework that maintains high image generation quality while resisting adversarial prompt jailbreaks.
- To identify the optimal model component (text encoder vs. UNet) for adversarial robustification in unlearning.
- To enable plug-and-play transferability of the robust unlearner across diverse diffusion model architectures.
- To systematically study the effectiveness and efficiency of adversarial training in the context of machine unlearning for DMs.
Proposed method
- Integrates adversarial training (AT) into the machine unlearning pipeline to improve robustness against adversarial prompt attacks.
- Applies adversarial training specifically to the text encoder rather than the UNet, as it proves more effective for robust unlearning.
- Employs a utility-retaining regularization on a retained prompt set to preserve image generation quality during adversarial fine-tuning.
- Uses a multi-step adversarial prompt generation strategy to craft perturbations that maximize attack success on unlearned models.
- Enables plug-and-play deployment of the trained text encoder on other diffusion models without retraining.
- Optimizes the first N layers of the text encoder, with ablation showing that deeper layer optimization improves robustness for complex concepts like nudity.
Experimental results
Research questions
- RQ1Can adversarial training be effectively and efficiently integrated into diffusion model unlearning to improve robustness against adversarial prompts?
- RQ2Which component of the diffusion model (text encoder or UNet) is more effective for adversarial robustification in unlearning?
- RQ3How can the trade-off between robustness and image generation utility be effectively managed during adversarial unlearning?
- RQ4Can a robust unlearner trained on one diffusion model be transferred to other models without fine-tuning?
- RQ5What is the impact of different text encoder layer configurations on the robustness and utility of unlearning?
Key findings
- AdvUnlearn reduces the attack success rate (ASR) to 21.13% on SD v1.4 for nudity unlearning, compared to 100% in the original model.
- The framework maintains strong image generation quality, with a FID of 19.34, only slightly higher than the original model’s 16.70.
- The text encoder is significantly more effective than the UNet for adversarial robustification, with the former achieving lower ASR and better utility trade-offs.
- The robust text encoder learned via AdvUnlearn can be transferred to other models (e.g., SD v1.5, DreamShaper, Protogen) with minimal performance drop, achieving 16.20% ASR on SD v1.5 and 42.96% on Protogen.
- Optimizing more layers of the text encoder improves robustness, with full-layer optimization yielding the strongest defense, especially for complex concepts like nudity.
- The utility-retaining regularization effectively mitigates the degradation in image quality typically caused by adversarial training, preserving model utility.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.