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

AI summary

Overview

Research area: Machine learning / generative AI safety — specifically concept erasure in text-to-image latent diffusion models (LDMs), combining adversarial training ideas with gradient-conflict-aware optimization.

Technical level: Advanced. The paper assumes familiarity with diffusion model denoising objectives, classifier-free guidance, adversarial prompt optimization, and gradient projection methods.

Scope: The paper proposes AEGIS, a retention-data-free framework for erasing undesirable concepts from diffusion models that improves robustness against adversarial prompt attacks while preserving general model utility. (Note: the provided paper content is truncated — the CLIP-score column of Table 2 and the full results for the Van Gogh and Church concepts are not included in the supplied text.)

What This Paper Is About

Diffusion models fine-tuned to forget harmful concepts (such as "nudity") can usually be tricked back into generating that content by adversarial or semantically related prompts the model never saw during erasure training, and prior defenses that block these attacks tend to damage the model's ability to generate unrelated content. AEGIS attacks this robustness–retention trade-off from two angles: it replaces the handpicked "safe" target concept with an optimized Adversarial Erasure Target that sits near the semantic center of the whole concept class, and it adds a gradient-projection step so the retention constraint and the erasure objective stop fighting each other — all without needing any auxiliary retention dataset.

Key Contributions

1. Diagnosis of the vulnerability source. The authors argue that concept erasure fails not because of insufficient training, but because of an inappropriately chosen learning target: if the erasure target is too close to the semantic center formed by words related to the erased concept, class-level concept information is never fully removed and remains exploitable by adversarial prompts. They support this with a theoretical deviation lower bound (Proposition 3.1) and empirical predicted-noise-distance measurements (d0, d1, d2).

2. Adversarial Erasure Target (AET). AEGIS replaces handpicked generic targets with an optimizable embedding, iteratively updated over K = 10 steps using sign-gradient updates with step size beta, that is pushed to approximate the semantic center of the erased concept class rather than a single instance of it. Aligning the model's prediction on the erased prompt to an AET-derived target increases predicted-noise distances not just for the erased instance but for semantically related variants.

3. Gradient Regularization Projection (GRP). A retention-data-free mechanism made of two parts: Parameter Regularization (PR), a layer-wise penalty on deviation between fine-tuned and original parameters that removes the need for a retention dataset and avoids accidentally relearning the erased concept; and Directional Gradient Rectification (DGR), which projects away the destructive component of the retention gradient only when it conflicts with the erasure direction (cos phi < 0), with a dynamically adapted rectification strength omega.

4. Empirical and theoretical validation. Extensive experiments across multiple concepts and three adversarial attacks (P4D, UnlearnDiffAtk, Ring-A-Bell) show large ASR reductions — the paper reports approximately 5.31% reduction on the nudity concept and approximately 24% on the Van Gogh style concept under P4D and UnlearnDiffAtk — while maintaining or improving FID/CLIP relative to advanced baselines, plus two theorems (4.1 and 4.2) characterizing erasure convergence and retention benefit.

Main Findings

• The vulnerability comes from the erasure target, not the attack. In a motivation study on "nudity" (c_e^0) and its related terms naked, sexual, erotic, impure, the authors define three concept groups C0 (all five terms), C1 (nudity, naked), and C2 (sexual, erotic, impure). Noise predictions for C1 terms are close together while C2 terms diverge strongly, so erasing "nudity" via ESD enlarges distances for C1 but leaves C2 essentially unchanged — meaning class-level information survives.

• Predicted noise distances expose the failure. The paper introduces d0 (intra-model separation), d1 (cross-model erasure strength), and d2 (distance between original and adversarial prompts in the fine-tuned model, reflecting adversarial vulnerability), measured at 10^-4 scale. After ESD fine-tuning, d2 stays consistently small for all instances of C0, which explains the high attack success rates.

• AET substantially improves robustness. Compared with AdvUnlearn, the ESD+AET variant achieves a 9.07% and 59.16% reduction in pre-attack success rate and attack success rate respectively (Table 1: ESD w/ AET reaches 6.38% Pre-ASR and 26.06% ASR, versus ESD at 20.42% and 76.05%, and AdvUnlearn at 15.45% and 64.79%).

• AEGIS leads on nudity erasure robustness. In Table 2, AEGIS attains ASR1 = 12.06% (P4D), ASR2 = 8.45% (UnlearnDiffAtk), and ASR3 = 3.52% (Ring-A-Bell) on SD v1.4, against base-model scores of 100%, 100%, and 83.10%. Strong baselines such as SalUn reach 19.86%, 11.27%, and 7.04%, and STEREO 45.77%, 14.08%, and 7.04%.

• Utility is largely preserved. AEGIS reports FID = 17.43 on the nudity task, close to the base SD v1.4 FID of 16.7 and lower than SalUn (33.62), AdvUnlearn (19.34), STEREO (18.27), and ESD (18.18). CLIP scores are referenced in the paper's metric description and Table 2, but the numeric CLIP row is not present in the truncated content provided.

• Gradient conflict between erasure and retention is real and measurable. The cosine similarity between the erasure gradient g_e and retention gradient g_r is negative throughout many training iterations, motivating the DGR projection with lambda = 0 when cos phi ≥ 0 and a projected lambda when cos phi < 0.

• Theory supports both goals. Theorem 4.1 shows GRP updates either converge to a degenerate point where cos phi = -1 or to a local minimum of the erasure loss (for step size alpha ≤ 2/L), and Theorem 4.2 shows the rectified parameter yields retention loss no greater than the unrectified one under stated curvature conditions — with the caveat that the learning rate must not be excessively small.

Methodology in Plain English

The starting point is the standard concept-erasure recipe: take a diffusion model, and fine-tune it so that when it sees a forbidden prompt such as "nudity," it behaves as if it saw a harmless prompt instead. The authors first show empirically that this recipe only removes the specific wording used, leaving other words that mean similar things ("sexual," "erotic") still able to trigger the concept.

Their fix has two moving parts. First, instead of picking a fixed safe target by hand, they optimize a small learnable text embedding — the Adversarial Erasure Target — by repeatedly nudging it with sign-gradient steps so that it lands near the middle of the whole cluster of related words, rather than on the edge. Erasing toward that central point removes the entire neighborhood instead of just one word. They speed this up by updating the target once per training epoch rather than running many inner iterations.

Second, they notice that the retention objective and the erasure objective often pull the model's weights in opposite directions. Rather than forcing both, they check the angle between the two gradients. If they point the same way, they simply skip retention for that step. If they conflict, they remove only the component of the retention gradient that opposes erasure, using a strength factor that itself is adapted during training so that erasure gets priority early on. Because retention is expressed purely as a penalty on how far the weights drift from the original model, no retention images or prompts are needed at all.

All experiments fine-tune the U-Net with the Adam optimizer at a learning rate of 10^-5, batch size 1, for 1000 epochs, with AET step size 10^-3, one AET iteration per epoch, and a step size of 0.1 for updating the rectification strength. Stable Diffusion v1.4 and v2.1 are the foundation models; evaluations cover nudity, artistic style, and object-related concepts using P4D, UnlearnDiffAtk, and Ring-A-Bell attacks, with ASR, FID, and CLIP as metrics.

Why This Matters

Concept erasure sits directly on the safety-critical path for deploying image generators: if erasure can be defeated by a rephrased or adversarially crafted prompt, the safety guarantee is mostly cosmetic. This paper matters because it reframes the problem — arguing the weak point is the choice of erasure target rather than the amount of fine-tuning — and because it removes the need for retention datasets, which are expensive, hard to curate, and a source of bias or accidental concept relearning.

Real-world applications include:

• Content moderation and safety filters in commercial image generators, where a moderation stack needs to hold up against users who actively search for bypass phrasing.
• Copyright and artist-style protection, letting style owners request removal of an artistic style (the paper evaluates the Van Gogh style concept) without degrading the model's other styles.
• Regulatory compliance for model providers required to demonstrate that harmful categories such as nudity are reliably suppressed, with auditable attack success rates.
• Creative and design tools, where unrelated generation quality (measured by FID and CLIP) must stay intact after safety fine-tuning so the product remains useful.
• On-device or privacy-sensitive deployments, where a retention-data-free method avoids shipping or curating extra datasets.

Industry relevance is straightforward: the method is drop-in on top of existing erasure pipelines, the paper reports only small retention costs, and it is evaluated against a broad set of twelve baseline methods (ESD, FMN, AC, UCE, SalUn, AGE, RECE, RECELER, SH, ED, STEREO, SPM), which is the comparison a practitioner would need before adoption. Source code is released at https://github.com/Feng-peng-Li/AEGIS.

Future Directions

• Verifying AET's semantic-center approximation at scale. The AET is optimized over K = 10 sign-gradient steps from a randomly initialized embedding of length m' = 1; it is open how this scales to concepts with many more related terms or to abstract concepts where a semantic center is ill-defined.
• Extending beyond the three concept groups tested. Results are reported in detail for nudity, with Van Gogh (style) and Church (object) named as evaluated categories; full robustness and retention numbers for those categories are not present in the truncated content and would clarify generalizability.
• Characterizing the conditions in Theorem 4.2. The retention benefit depends on curvature conditions (ell ≥ L·||g_e − g_r||²/||g_e + g_r||², and a stricter version involving 2/alpha) and on the learning rate not being too small — practical guidance on when these hold for real U-Net fine-tuning would strengthen the method.
• Adapting the framework to newer architectures and modalities. The work targets latent diffusion U-Nets in Stable Diffusion v1.4 and v2.1; whether the target-guided erasure and gradient projection transfer to transformer-based or video diffusion models is untested here.
• Analyzing the rectification strength schedule. Whether the dynamically updated omega, which starts small, needs retuning across concepts or can be set universally remains an open practical question.

Target Audience

This paper suits machine learning researchers and graduate students working on generative model safety, machine unlearning, and adversarial robustness, as well as applied scientists and engineers responsible for deploying text-to-image systems who need to understand why prompt-level safety filters fail and what a more robust erasure pipeline looks like. Readers should already be comfortable with diffusion model training objectives and gradient-based optimization; the theorems and the adversarial-target formulation are not beginner material.

Authors’ abstract

Concept erasure helps stop diffusion models (DMs) from generating harmful content; but current methods face robustness retention trade off. Robustness means the model fine-tuned by concept erasure methods resists reactivation of erased concepts, even under semantically related prompts. Retention means unrelated concepts are preserved so the model's overall utility stays intact. Both are critical for concept erasure in practice, yet addressing them simultaneously is challenging, as existing works typically improve one factor while sacrificing the other. Prior work typically strengthens one while degrading the other, e.g., mapping a single erased prompt to a fixed safe target leaves class level remnants exploitable by prompt attacks, whereas retention-oriented schemes underperform against adaptive adversaries. This paper introduces Adversarial Erasure with Gradient Informed Synergy (AEGIS), a retention-data-free framework that advances both robustness and retention.