PaperlayerView exact revision on arXiv

Source-linked AI summary

Defensive Unlearning with Adversarial Training for Robust Concept Erasure in Diffusion Models

Yimeng Zhang, Xin Chen, Jinghan Jia, Yihua Zhang, Chongyu Fan, Jiancheng Liu, Mingyi Hong, Ke Ding, Sijia Liu

arXiv:2405.15234v3cs.CVcs.CR

TL;DR

Concept-erased diffusion models remain vulnerable to adversarial prompts that recover erased concepts, motivating robust unlearning. AdvUnlearn integrates adversarial training with utility-retaining regularization and text-encoder robustification. Across nudity, object, and style erasure, it improves robustness while preserving image-generation utility, though its attack-generation and regularization costs remain computational limitations.

  • Problem

    Concept-erased diffusion models can regenerate erased concepts under adversarial prompt attacks, leaving robust concept erasure as an open problem.

  • Method

    AdvUnlearn integrates adversarial training into machine unlearning using utility-retaining regularization and robustification of the text encoder rather than the UNet.

  • Results

    AdvUnlearn improves robustness across nudity, object, and style erasure scenarios while maintaining a balanced level of image-generation utility.

  • Takeaways & Limitations

    A robust text encoder can serve as a plug-and-play unlearner across different diffusion-model types.

  • Takeaways & Limitations

    K-step adversarial-example generation and utility-retaining regularization add substantial computation, motivating more efficient implementations.

Abstract

from arXiv · show

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

1 Introduction

Diffusion models can generate harmful content, and concept-erased models remain vulnerable to adversarial prompts that restore erased concepts. AdvUnlearn integrates adversarial training with utility-retaining regularization and robustifies the text encoder to improve robustness while preserving image-generation utility.

  • Motivation: Concept-erased diffusion models can be jailbroken by small adversarial prompt perturbations to regenerate content targeted for erasure.The introduction illustrates this vulnerability for nudity erasure in Stable Diffusion v1.4 using UnlearnDiffAtk.
  • Approach: AdvUnlearn integrates adversarial training with concept erasure through a bi-level optimization scheme.The framework alternates unlearning-model optimization with adversarial-prompt optimization.
  • Approach: Utility-retaining regularization on curated external retain prompts addresses the image-generation quality loss caused by straightforward adversarial training.The regularizer balances effective concept erasure against high-quality generation.
  • Findings: Optimizing the text encoder improves robustness over conventional UNet-focused unlearning while achieving a better robustness–utility balance.The resulting robust text encoder can be shared across different diffusion models in a plug-and-play manner.
  • Findings: AdvUnlearn improves robustness across nudity, object, and style erasure scenarios while preserving a commendable level of image-generation utility.The introduction points to Figure 2 for performance evidence under nudity erasure.

2 Related Work

Prior work addresses diffusion-model safety through concept erasure, adversarial-prompt defenses, and adversarial training. These lines of research motivate robust unlearning that can resist prompt-based attempts to recover erased concepts.

  • Concept erasure: Concept erasure reduces the influence of undesired textual concepts on diffusion-model image generation.It is motivated by the possibility of unsafe or sensitive outputs from inappropriate prompts.
  • Adversarial prompts: Adversarial prompt attacks manipulate characters, words, or token embeddings to bypass safety mechanisms and induce undesirable generation.Such attacks are studied as prompt-level jailbreaks against safety-driven diffusion models.
  • Adversarial training: Adversarial training frames robustness as a two-player game between an attacker and a defender.This prior defense paradigm provides the conceptual basis for integrating adversarial training into diffusion-model unlearning.

3 Preliminaries and Problem Statement

The paper formulates robust concept erasure for latent diffusion models as adversarial prompt optimization coupled with model unlearning. AdvUnlearn addresses the resulting utility and module-selection challenges through bi-level optimization and targeted robustification.

  • Latent diffusion models: Latent diffusion models generate images by progressively denoising latent representations conditioned on text prompts.A noise estimator predicts denoising behavior at each time step from the latent state and conditioning prompt.
  • Concept erasure: Concept erasure aims to prevent unsafe or sensitive concepts from influencing diffusion-model image generation.ESD is adopted as a recognized baseline because it balances erasure effectiveness with model-utility preservation.
  • Concept erasure: ESD modifies conditional denoising using the original model’s concept-conditioned and unconditional predictions to reduce generation of the target concept.The erasing guidance parameter η controls the adjustment toward minimizing images associated with the erased concept.
  • Adversarial prompt attacks: Adversarial prompts are perturbed versions of target prompts that steer a concept-erased model toward unsafe content produced by the original model.Perturbations can alter tokens or token embeddings while remaining close to the original prompt under an ℓ0 or embedding-space constraint.
  • AdvUnlearn: AdvUnlearn formulates robust unlearning as a defender–attacker game and alternates model-parameter optimization with adversarial-prompt optimization.The upper level minimizes an unlearning objective, while the lower level finds an optimized adversarial prompt under the current model.
  • AdvUnlearn: AdvUnlearn is defined as the proposed integration of adversarial training into diffusion-model unlearning.The paper identifies solving this integrated optimization problem effectively and efficiently as nontrivial.
  • Challenges: A naive adversarial-training implementation can substantially reduce normal-image utility, creating a robustness–utility trade-off.The framework must also determine where to apply adversarial training within the text encoder–UNet decomposition and how to implement it efficiently.

4 Effectiveness Enhancement of AdvUnlearn: Improving Tradeoff between Robustness and Utility

Direct adversarial training improves robustness but substantially harms image quality. AdvUnlearn adds retain-set regularization to recover utility while preserving robustness gains over ESD.

  • Direct AT trade-off: Approximately 20% lower ASR from direct AT-ESD comes with a considerable FID increase and visibly worse benign-prompt authenticity.The comparison uses ESD and AT-ESD for nudity unlearning on base SD-v1.4.
  • Utility-retaining regularization: AdvUnlearn augments the ESD objective with a retain-set penalty that matches the current model’s outputs to the original model on retained concepts.The penalty targets generation-quality degradation while preserving the concept-erasure objective.
  • Utility-retaining regularization: The retain prompts are filtered by an LLM to exclude prompts related to the targeted concept, reducing the risk that utility preservation undermines erasure.Prompts are drawn from external datasets such as ImageNet or COCO.
  • Overall trade-off: AdvUnlearn recovers utility compromised by AT-ESD and remains more robust than ESD, although it sacrifices some robustness relative to ESD’s stronger ASR result.The reported comparison evaluates SD v1.4, ESD, AT-ESD, and AdvUnlearn for nudity unlearning.

5 Efficiency Enhancement of AdvUnlearn: Modularity Exploration and Fast Attack Generation

AdvUnlearn is more effective and efficient when robustification is applied to the text encoder rather than UNet. A one-step fast attack greatly reduces training time, but lowers unlearning efficacy and utility.

  • Text encoder versus UNet: A robust text encoder can provide a modular, efficient target for concept erasure compared with optimizing the larger UNet.The method’s efficiency motivation follows from the text encoder’s smaller parameter count and faster convergence.
  • Text encoder versus UNet: Text-encoder AdvUnlearn achieves much lower ASR than UNet-based AdvUnlearn without losing image-generation utility.The text encoder also has fewer parameters and converges more quickly than the UNet.
  • Fast attack generation: Fast AT reduces training time per iteration from 78.57s to 12.13s on one NVIDIA RTX A6000 GPU.The speedup is reported relative to standard AT with 30 attack steps.
  • Fast attack generation: The fast variant correspondingly decreases unlearning efficacy and image-generation utility, making it suitable when computational efficiency is prioritized over maximal efficacy.The paper presents fast AT as an option when unlearning efficacy is not exceedingly important.

6 Experiments

Experiments evaluate AdvUnlearn on nudity, style, and object erasure using adversarial robustness and image-generation utility metrics. Across these settings, AdvUnlearn improves robustness while preserving utility, and its robust text encoder transfers across diffusion models.

  • Evaluation Scope: AdvUnlearn is evaluated across nudity, object, and style erasure scenarios to assess robustness and retained image-generation utility.The experiments compare the method with multiple open-sourced unlearning baselines.
  • Experiment Setups: Experiments measure adversarial robustness with ASR and image-generation utility with FID and CLIP score.Lower ASR indicates better robustness; FID and CLIP evaluate utility on benign prompts.
  • Nudity Unlearning: Over 50% ASR reduction compared to ESD is achieved for nudity unlearning, while avoiding SalUn’s substantial utility cost.AdvUnlearn has significantly better FID and CLIP scores than SalUn, despite SalUn attaining the lowest ASR.
  • Nudity Unlearning: SalUn and AdvUnlearn are the only evaluated methods producing robust concept-erased models under nudity-focused adversarial prompts, with AdvUnlearn generating higher-quality images.Figure 4 compares adversarial generations with generations from a nudity-irrelevant benign prompt.
  • Style Unlearning: Over 30% ASR decrease compared to ESD is obtained for Van Gogh style unlearning without sacrificing utility.AdvUnlearn’s FID and CLIP scores remain comparable to those of base SD v1.4.
  • Object Unlearning: For Church object unlearning, AdvUnlearn attains the highest robustness while preserving original diffusion-model utility better than SH.The comparison is also illustrated with generated-image examples.
  • Transferability: The AdvUnlearn text encoder transfers robustness and utility to SD v1.5, DreamShaper, and Protogen, including dissimilar models.Transfer is most significant for SD v1.5, while remaining effective for DreamShaper and Protogen without significantly sacrificing utility.
  • Text-Encoder Layers: For nudity unlearning, robustness increases as more text-encoder layers are optimized, whereas utility slightly degrades.Optimizing only initial layers is inadequate for adversarial robustness, unlike object or style unlearning, where optimizing the first layer is satisfactory.

7 Conclusion

AdvUnlearn improves robustness against adversarial prompt attacks while preserving image-generation utility, using retained-prompt regularization and text-encoder optimization. Its workflow alternates adversarial prompt generation with defensive unlearning, while concept-specific prompt filtering supports the unlearning setup.

  • 7 Conclusion: AdvUnlearn balances robust concept erasure with image-generation utility through utility-retaining regularization on a retained prompt set.The framework integrates adversarial training into concept erasing and uses external retain prompt data to address utility loss.
  • 7 Conclusion: Optimizing the text encoder is more effective than robustifying the UNet and enables a shared, plug-and-play robust unlearner across diffusion models.The text encoder provides a better balance between unlearning performance and image-generation utility.
  • 7 Conclusion: AdvUnlearn demonstrates robustness improvements across nudity, object, and style erasure while maintaining a commendable level of image-generation utility.The conclusion identifies continued computational-efficiency improvement as a future research need.
  • 7 Conclusion: For nudity, an LLM judger filters prompts that may generate the global nudity concept, including prompts that do not explicitly mention nudity.For objects, filtering targets the object name and its synonyms because objects are treated as local concepts.
  • 7 Conclusion: Each AdvUnlearn iteration generates adversarial prompts and then performs unlearning with utility-retaining regularization.The adversarial prompt is optimized before being used in the unlearning stage.

C Attack Setup

The attack setup creates concept-specific adversarial prompt perturbations and evaluates generated images with classifiers tailored to nudity, style, and object erasure.

  • C Attack Setup: Adversarial perturbations prepend N tokens to prompts, using N = 5 for nudity and N = 3 for style and object unlearning.These perturbations are optimized with UnlearnDiffAtk over 40 iterations using 50 diffusion time steps and AdamW at learning rate 0.01.
  • C Attack Setup: Generated images are classified with NudeNet for nudity, a WikiArt-finetuned ViT-base for style, and ImageNet-pretrained ResNet-50 for objects.The classifiers provide concept-specific evaluation for the three unlearning scenarios.

D Other Results and Visualizations for Nudity Unlearning

Additional nudity results compare SH and ED, which show strong adversarial robustness but substantially degraded image quality and weak prompt alignment.

  • D Other Results and Visualizations for Nudity Unlearning: SH and ED achieve high adversarial unlearning robustness but have dramatically high FID and low CLIP scores.The reported metrics are Attack Success Rate, Fréchet Inception Distance, and CLIP score.
  • D Other Results and Visualizations for Nudity Unlearning: SH and ED generate similar-looking images across condition prompts, indicating poor variation and contextual appropriateness.Their visualizations support the reported image-quality and prompt-alignment trade-off.
  • D Other Results and Visualizations for Nudity Unlearning: Because of this utility degradation, SH and ED are excluded from the main nudity-unlearning performance table.

E Other Results and Visualizations for Object Unlearning

Object-unlearning evaluations on Garbage Truck, Parachute, and Tench show AdvUnlearn achieving strong robustness with competitive image-generation utility across concepts.

  • E Other Results and Visualizations for Object Unlearning: AdvUnlearn consistently achieves the best object-unlearning efficacy, around 10%, with competitive image-generation utility.The comparison uses the base SD v1.4 model and includes seven unlearning methods.
  • E Other Results and Visualizations for Object Unlearning: FMN and SPM retain image-generation utility but exhibit weak robustness against adversarial prompt attacks.
  • E Other Results and Visualizations for Object Unlearning: SH provides strong robustness at the cost of degraded image-generation utility, while SalUn, ED, and ESD show unstable robustness across object concepts.Their Attack Success Rates are multiple times higher than those of AdvUnlearn.
  • E Other Results and Visualizations for Object Unlearning: The object-unlearning results include Garbage Truck, Parachute, and Tench, with accompanying visualization examples.

F Other Ablation Studies

The ablations examine retain-set construction, adversarial prompt strategies, and test-time attacks for AdvUnlearn. COCO-based, LLM-filtered retain prompts and the Prefix strategy improve unlearning efficacy while preserving utility, and robustness transfers across attack methods.

  • Retain set selection: COCO-sourced retain sets outperform ImageNet-sourced sets in unlearning efficacy with only minor utility loss.Equal-sized retain sets were used for the comparison.
  • Retain set selection: LLM-filtered retain prompts significantly boost unlearning efficacy while preserving image generation utility.The comparison contrasts filtered prompts with datasets assembled without filtering.
  • Adversarial prompting strategy: Prefix prompting achieves the highest unlearning efficacy, with nearly half the Attack Success Rate of Replace and Add strategies while maintaining competitive utility.Prefix prepends an optimized adversarial soft prompt before the original concept prompt in token embedding space.
  • Robustness across test-time attacks: AdvUnlearn improves robustness against both UnlearnDiffAtk and P4D at test time, with lower ASR against P4D.UnlearnDiffAtk is the default attack used in the comparison.

G Limitations

AdvUnlearn incurs substantial computational overhead from multi-step adversarial example generation and utility-retaining regularization. Improving efficiency without significantly compromising performance remains a future need.

  • Computational cost: K-step adversarial example generation requires nearly K times more computation than vanilla training.Faster FGSM-based attack generation was found to degrade robust performance.
  • Computational cost: Utility-retaining regularization adds further computation time, motivating efficiency improvements without significant performance loss.The added regularization supports image generation utility but increases the method’s computational burden.
Loading 2405.15234v3…