STEREO: A Two-Stage Framework for Adversarially Robust Concept Erasing from Text-to-Image Diffusion Models

2024-08-29

AI Summary

Concept erasure in text-to-image diffusion models (T2ID) remains highly vulnerable to adversarial attack recovery, and existing methods struggle to simultaneously ensure robustness and generation quality. Method: We propose a two-stage robust erasure framework: (1) an adversarial training phase that probes blind spots in the embedding space—using adversarial examples solely as diagnostic tools for vulnerability detection, not for erasure; and (2) an anchor-concept-guided compositional loss design enabling robust erasure via single-stage fine-tuning. Contribution/Results: Our work introduces the novel paradigm “adversarial training as probing,” eliminating cumulative distortion inherent in traditional multi-stage optimization. Evaluated against seven state-of-the-art methods, our approach achieves superior performance, demonstrating strong robustness against both white-box and black-box attacks while preserving fidelity and diversity of benign concepts with negligible degradation.

Technology Category

• Computer Vision: Adversarial Attacks & Robustness
• Machine Learning: Adversarial Learning & Robustness
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Security and Privacy: Data transparency and provenance

Abstract

The rapid proliferation of large-scale text-to-image diffusion (T2ID) models has raised serious concerns about their potential misuse in generating harmful content. Although numerous methods have been proposed for erasing undesired concepts from T2ID models, they often provide a false sense of security; concept-erased models (CEMs) can still be manipulated via adversarial attacks to regenerate the erased concept. While a few robust concept erasure methods based on adversarial training have emerged recently, they compromise on utility (generation quality for benign concepts) to achieve robustness and/or remain vulnerable to advanced embedding space attacks. These limitations stem from the failure of robust CEMs to thoroughly search for"blind spots"in the embedding space. To bridge this gap, we propose STEREO, a novel two-stage framework that employs adversarial training as a first step rather than the only step for robust concept erasure. In the first stage, STEREO employs adversarial training as a vulnerability identification mechanism to search thoroughly enough. In the second robustly erase once stage, STEREO introduces an anchor-concept-based compositional objective to robustly erase the target concept in a single fine-tuning stage, while minimizing the degradation of model utility. We benchmark STEREO against seven state-of-the-art concept erasure methods, demonstrating its superior robustness to both white-box and black-box attacks, while largely preserving utility.

Problem

• Preventing harmful content generation in text-to-image models
• Improving robustness of concept-erased models against attacks
• Maintaining model utility while erasing target concepts

Innovation

• Two-stage adversarial training framework
• Anchor-concept-based compositional objective
• Robust erasure with minimal utility loss