🤖 AI

AEGIS: Adversarial Target-Guided Retention-Data-Free Robust Concept Erasure from Diffusion Models

This paper introduces AEGIS, a retention-data-free framework that simultaneously enhances the robustness and retention of diffusion models during concept erasure by leveraging adversarial target-guided mechanisms to overcome the trade-offs inherent in existing methods.

Original authors: Fengpeng Li, Kemou Li, Qizhou Wang, Bo Han, Jiantao Zhou

Published 2026-02-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Fengpeng Li, Kemou Li, Qizhou Wang, Bo Han, Jiantao Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a magical artist named Diffusion. This artist is incredibly talented and can paint anything you describe, from "a cat on a skateboard" to "a sunset over Paris." However, Diffusion was trained on the entire internet, which means it also learned some bad habits. Sometimes, if you ask for something slightly off, it might accidentally generate something inappropriate, violent, or copyrighted (like painting in the exact style of a living artist without permission).

To fix this, researchers try to teach Diffusion to "unlearn" these bad habits. This is called Concept Erasure. But here's the problem: teaching Diffusion to forget is like trying to remove a stain from a white shirt without making the whole shirt gray or tearing a hole in it.

The Problem: The "Stain" vs. The "Shirt"

Current methods face a tough trade-off:

Robustness (The Stain): If you try to remove the "nudity" concept, you want to make sure it's gone forever. But clever attackers can trick the model by using synonyms (like "naked" instead of "nudity") or weird phrasing to bring the bad images back.
Retention (The Shirt): If you try too hard to scrub the stain, you might accidentally erase the ability to draw "people" or "clothing" entirely, ruining the artist's ability to paint anything else.

Existing methods usually fix one problem but break the other. They either leave the stain visible (not robust) or ruin the shirt (bad retention).

The Solution: AEGIS

The paper introduces a new method called AEGIS (Adversarial Erasure with Gradient-Informed Synergy). Think of AEGIS as a smart, two-step cleaning crew that knows exactly how to scrub without damaging the fabric.

Step 1: The "Smart Target" (Adversarial Erasure Target)

Imagine you want to teach Diffusion that "nudity" is bad.

Old Way: You show the model one picture of a nude person and say, "Don't paint this." The model learns to avoid that specific picture. But if someone asks for "a person with no clothes," the model might still do it because it didn't learn the general idea of nudity, just that one specific image.
AEGIS Way: AEGIS creates a "Super-Target." It doesn't just pick one image; it calculates the "center of gravity" for the entire concept of nudity. It's like finding the exact middle of a crowd of people who are all related to the bad concept.
- The Analogy: Instead of telling the artist, "Don't paint this specific red car," AEGIS says, "Don't paint any red car, and here is the exact center of the 'red car' concept so you know exactly what to avoid."
- By aiming for this "center," AEGIS ensures that even if an attacker uses a synonym or a weird description, the model still knows to say "No." It makes the erasure robust against tricks.

Step 2: The "Balanced Brush" (Gradient Regularization Projection)

Now, imagine you are scrubbing the stain. As you scrub hard, your hand starts shaking, and you might accidentally scrub the good parts of the shirt (like the collar or the buttons).

The Conflict: The instruction to "remove the bad concept" pushes the model in one direction. The instruction to "keep the good concepts" pushes it in another. Sometimes, these two forces fight against each other.
AEGIS Way: AEGIS uses a technique called Gradient Regularization Projection (GRP).
- The Analogy: Imagine you are pushing a heavy box (the model) to the left to remove the bad stuff. But you also need to keep the box from sliding too far left and hitting a wall (ruining the good stuff).
- AEGIS acts like a smart guide. It looks at your push. If your push to remove the bad stuff is also accidentally pushing the good stuff away, AEGIS gently redirects your hand. It says, "Okay, push left to remove the bad, but don't push forward, because that ruins the buttons."
- It does this without needing extra photos of "good" things to look at. It just uses the model's own memory of what it was like before it started scrubbing to know what to keep.

Why This Matters

The authors tested AEGIS on difficult concepts like "nudity," specific art styles (like Van Gogh), and objects (like churches).

Result: AEGIS was much harder to trick than previous methods. Attackers couldn't easily bring the bad images back.
Bonus: The model didn't lose its ability to draw other things. The "shirt" stayed white and intact.

Summary

AEGIS is like a master cleaner for AI artists.

It finds the true center of the bad concept so it can remove the entire idea, not just a specific example.
It uses a smart guide to scrub the bad stuff without accidentally erasing the good stuff, all without needing a reference library of "safe" images.

It solves the age-old problem of "how do I delete the bad without deleting the good?" by being smarter about what to target and how to move the model's brain.

1. Problem Statement

Diffusion Models (DMs) are powerful text-to-image generators but often inherit harmful biases (e.g., nudity, copyright infringement, specific objects) from training data. Concept Erasure aims to remove these undesirable concepts by fine-tuning the model. However, existing methods face a critical robustness–retention trade-off:

Robustness: The ability of the erased model to resist Adversarial Prompt Attacks (APAs), where attackers use semantically related or metaphorical prompts (e.g., "naked" instead of "nudity") to reactivate the erased concept.
Retention: The ability to preserve the generation quality of unrelated concepts (e.g., ensuring the model can still generate "landscapes" after erasing "nudity").

Current Limitations:

Inadequate Targets: Most methods map an erased concept (e.g., "nudity") to a fixed safe target (e.g., "a photo") or an empty prompt. If the chosen instance is far from the semantic center of the concept class, the erasure is incomplete. Adversarial prompts often target this semantic center, bypassing the erasure.
Gradient Conflict: Optimizing for erasure often conflicts with optimizing for retention. Existing methods either degrade generation quality to achieve robustness or fail to resist attacks to preserve quality.
Data Dependency: Many retention strategies require auxiliary datasets, which can introduce bias or fail to cover the full diversity of the original model's knowledge.

2. Methodology: AEGIS

The authors propose AEGIS (Adversarial Erasure with Gradient-Informed Synergy), a framework that achieves both high robustness and high retention without auxiliary data. It consists of two core components:

A. Adversarial Erasure Target (AET)

Instead of using a fixed target, AEGIS dynamically optimizes an Adversarial Erasure Target ( $c'$ ) to approximate the semantic center of the erased concept class.

Mechanism: The AET is an optimizable embedding updated iteratively to maximize the distance between the model's prediction for the erased concept ( $c_e$ ) and the prediction for the AET in the shared text-image space.
Goal: By aligning the erased concept with a target that represents the "center" of the concept cluster (rather than a single instance), the method ensures that all semantically related variants (including adversarial prompts) are pushed away from the original concept space.
Optimization: The AET is updated via gradient ascent to maximize the noise prediction distance between the original model's output for the AET and the fine-tuned model's output for the erased concept.

B. Gradient Regularization Projection (GRP)

To solve the trade-off between erasure and retention without auxiliary data, AEGIS introduces a data-free retention strategy combining Parameter Regularization (PR) and Directional Gradient Rectification (DGR).

Parameter Regularization (PR): Instead of using a retention dataset, AEGIS penalizes deviations from the original model parameters ( $\theta_0$ ) using an $L_2$ loss: $L_r = \frac{1}{2}\|\theta - \theta_0\|^2$ . This preserves the model's general knowledge without introducing dataset bias.
Directional Gradient Rectification (DGR): The gradients for erasure ( $g_e$ $g_{e}$ ) and retention ( $g_r$ $g_{r}$ ) often conflict (negative cosine similarity). DGR detects this conflict and projects the retention gradient onto the orthogonal complement of the erasure gradient.
- If $\cos(g_e, g_r) \geq 0$ , no modification is made.
- If $\cos(g_e, g_r) < 0$ , the retention gradient is adjusted to remove the component that opposes the erasure direction, ensuring the erasure task is not hindered.
Dynamic Weighting ( $\omega$ ): A dynamic parameter $\omega$ controls the strength of the projection. It starts at 0 (prioritizing erasure) and increases as training progresses, allowing the model to first learn the erasure task and then gradually enforce retention constraints.

3. Key Contributions

Vulnerability Analysis: The paper identifies that the failure of current erasure methods stems from mapping erased instances to targets far from the semantic center of the concept class, leaving residual information exploitable by adversarial prompts.
Novel Framework (AEGIS): Proposes a retention-data-free framework that synthesizes an Adversarial Erasure Target (AET) to guide class-level removal and uses Gradient Regularization Projection (GRP) to balance robustness and utility.
Theoretical Guarantees: Provides theoretical proofs (Theorems 4.1 & 4.2) demonstrating that GRP ensures local descent for the erasure loss while preserving (or improving) retention performance under specific curvature conditions.
State-of-the-Art Performance: Extensive experiments show AEGIS significantly outperforms existing baselines (ESD, AdvUnlearn, SalUn, etc.) in both attack resistance and image quality.

4. Experimental Results

The authors evaluated AEGIS on Stable Diffusion v1.4 and v2.1 across three concept categories: Nudity, Artistic Style (Van Gogh), and Objects (Church).

Robustness (Attack Success Rate - ASR):
- Against P4D and UnlearnDiffAtk attacks, AEGIS reduced the ASR significantly.
- Example (Nudity on SD v1.4): AEGIS achieved an ASR of 8.45% against UnlearnDiffAtk, compared to 64.79% for AdvUnlearn and 73.24% for ESD.
- Example (Van Gogh): AEGIS reduced ASR to 12% against UnlearnDiffAtk, a substantial improvement over the next best method (30%).
Retention (FID & CLIP):
- AEGIS maintained or improved image quality (FID) and prompt alignment (CLIP) compared to baselines.
- Unlike methods that rely on small retention datasets (which often fail to cover all concepts), AEGIS's parameter regularization approach preserved the model's utility more effectively, with FID scores often closer to the base model than competing methods.
Efficiency: While AEGIS takes slightly longer than vanilla erasure (due to AET generation), it is significantly faster than methods requiring large-scale retention datasets or complex multi-stage training.

5. Significance

Breaking the Trade-off: AEGIS effectively breaks the prevailing trade-off where increasing robustness inevitably degrades retention. It demonstrates that robust concept erasure is possible without sacrificing the model's general utility.
Data-Free Retention: By eliminating the need for auxiliary retention datasets, AEGIS avoids the bias and scalability issues associated with curating specific "safe" data, making the method more practical for real-world deployment.
Security Implications: The work provides a stronger defense against "jailbreaking" diffusion models via semantic prompt attacks, a critical step toward safer and more reliable generative AI systems.
Generalizability: The framework is applicable to various concept types (sensitive content, styles, objects) and different model versions, suggesting a robust solution for the broader challenge of machine unlearning in generative models.

In summary, AEGIS represents a significant advancement in the field of AI safety, offering a mathematically grounded, efficient, and highly effective method to permanently remove harmful concepts from diffusion models while preserving their creative capabilities.