🤖 AI

SAEmnesia: Erasing Concepts in Diffusion Models with Supervised Sparse Autoencoders

SAEmnesia introduces a supervised sparse autoencoder framework that enforces one-to-one concept-neuron mappings to achieve feature centralization, enabling highly efficient, scalable, and precise concept erasure in diffusion models while significantly reducing hyperparameter search and outperforming state-of-the-art methods on multiple benchmarks.

Original authors: Enrico Cassano, Riccardo Renzulli, Marco Nurisso, Mirko Zaffaroni, Alan Perotti, Marco Grangetto

Published 2026-06-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Enrico Cassano, Riccardo Renzulli, Marco Nurisso, Mirko Zaffaroni, Alan Perotti, Marco Grangetto

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 giant, incredibly talented artist (a Diffusion Model) who can draw anything you describe. But sometimes, this artist has a bad habit: they keep drawing things you don't want, like copyrighted characters, inappropriate content, or specific objects you've asked them to forget.

The problem is that this artist doesn't have a neat filing cabinet where "Bears" are in one drawer and "Sandwiches" are in another. Instead, the concept of a "Bear" is scattered across thousands of different mental notes, mixed up with "Cats," "Fur," and "Forest." This is called feature splitting. Trying to erase a "Bear" is like trying to remove a single thread from a tangled ball of yarn without unraveling the whole sweater. It's messy, slow, and often ruins the picture.

SAEmnesia is a new tool designed to fix this mess. Here is how it works, using simple analogies:

1. The Problem: The Tangled Yarn

In previous methods, if you wanted the artist to forget "Bears," you had to guess which thousands of mental notes to tweak. Because the notes were mixed up, you might accidentally erase "Fur" or "Animals" along with the bear, or you might have to try hundreds of different combinations to find the right ones. It's like trying to find a specific needle in a haystack by guessing which handful of hay to pull out.

2. The Solution: The "One-Concept, One-Neuron" Rule

SAEmnesia changes the training rules. Instead of letting the artist mix concepts randomly, it forces a strict one-to-one mapping.

The Analogy: Imagine giving the artist a new rule: "Every single concept gets its own dedicated, labeled sticky note."
How it works: The system uses a special "Supervised Sparse Autoencoder" (think of it as a smart librarian). During training, it looks at the artist's notes and says, "Okay, the concept 'Bear' goes on sticky note #11,979. The concept 'Van Gogh Style' goes on sticky note #4,200."
The Result: Now, "Bear" isn't scattered across the whole library; it's locked in a single, specific drawer. This is called Feature Centralization.

3. The Eraser: A Precise Scalpel

Once the concepts are neatly organized on their own sticky notes, erasing them becomes incredibly easy.

The Analogy: If you want the artist to stop drawing bears, you don't need to tear out pages of the book. You just find sticky note #11,979 and tell the artist, "Ignore this note."
The Benefit: This makes the process 96.67% faster at finding what to erase because you aren't searching through combinations anymore. You just turn off one specific switch.

4. What the Paper Actually Found

The authors tested this on a standard test called UnlearnCanvas (a benchmark for testing how well models can forget things). Here are their specific claims:

Better Accuracy: SAEmnesia improved the ability to erase objects by 9.22% compared to the previous best method (SAeUron).
Sequential Learning: If you ask the model to forget 9 different things in a row (like Bears, then Cats, then Flowers), SAEmnesia got 28.4% better at remembering everything else while forgetting the new targets.
Safety: It successfully suppressed "NSFW" (Not Safe For Work) content, specifically nudity, better than previous methods.
Robustness: Even when someone tried to trick the system with "adversarial attacks" (trying to force the model to draw the forbidden thing anyway), SAEmnesia held up much better than the competition.
Reversibility: Because the tool is attached to the model like a plug-in and doesn't permanently change the model's brain, you can unplug it, and the model goes back to being exactly how it was before.

Summary

Think of SAEmnesia as a precision editor for AI artists. Instead of hacking away at a messy, tangled web of ideas, it organizes the artist's mind so that every concept has its own unique address. To erase something, you just send a letter to that specific address saying, "Stop showing this." It's faster, cleaner, and more effective than trying to untangle the whole web.

Technical Summary: SAEmnesia: Erasing Concepts in Diffusion Models with Supervised Sparse Autoencoders

Problem Statement
Concept unlearning in text-to-image diffusion models is currently hindered by "feature splitting," a phenomenon where a single concept is distributed across many latent features (neurons) rather than being localized to a single unit. This distributed representation, known as polysemanticity, makes identifying and removing specific concepts challenging. Existing methods often require exhaustive searches across numerous latent combinations to find the correct subset for modification, rendering the process computationally expensive. Furthermore, overlapping latents blur concept boundaries, increasing the risk of unintended side effects on related concepts during intervention.

Methodology
The authors introduce SAEmnesia, a framework that utilizes a supervised sparse autoencoder (SAE) to enforce one-to-one mappings between concepts and latent neurons. Unlike unsupervised SAE approaches that discover features post-hoc, SAEmnesia integrates supervised labeling directly into the training process using the text prompts used to generate the images.

The core technical components include:

Feature Centralization: By systematically labeling concepts during training, the method achieves "feature centralization," binding each concept to a single, interpretable neuron. This prevents the splitting of concepts across multiple neurons.
Supervised Loss Function: The training objective combines standard unsupervised reconstruction loss with a supervised component ( $L_{supSAE}$ $L_{s u pS A E}$ ). This supervised loss consists of:
- Concept Assignment (CA) Loss: Encourages specific assigned latents to activate strongly when their corresponding concepts are present in the input.
- Decorrelation (DC) Loss: Penalizes correlations between activation patterns of latents assigned to different macro-groups of concepts (e.g., objects vs. styles) to promote disentanglement.
- Sparsity Regularization: An $L_1$ term to ensure most latents remain inactive for any given input.
Inference and Erasure: To erase a target concept, the method requires steering only a single latent. The activation of this specific latent is multiplied by a negative value (normalized by average activation) to suppress the concept's influence. This process is reversible; detaching the SAE restores the original model exactly.

Key Contributions
The paper outlines three primary contributions:

Supervised SAE Framework: The introduction of a supervised SAE that explicitly enforces one-to-one concept-latent mappings, eliminating feature splitting and enabling transparent control over concept representations.
Efficiency in Erasure: Demonstrating that this structure allows for concept erasure by steering a single latent, which substantially reduces the need for hyperparameter search during inference.
Performance Improvements: Achieving superior results on the UnlearnCanvas benchmark compared to state-of-the-art mechanistic approaches, improved scalability in sequential unlearning, stronger robustness to adversarial attacks, and effective suppression of NSFW content.

Experimental Results
Evaluated on the UnlearnCanvas benchmark using Stable Diffusion v1.5, SAEmnesia demonstrates significant improvements over the current state-of-the-art SAE-based method (SAeUron) and other unlearning approaches:

Benchmark Performance: SAEmnesia achieves a 91.51% average score on object unlearning, a 9.22% absolute improvement over SAeUron (82.29%). It shows consistent gains across Unlearning Accuracy (UA), In-domain Retain Accuracy (IRA), and Cross-domain Retain Accuracy (CRA).
Computational Efficiency: By reducing the search space from exploring multiple latent combinations to tuning a single multiplier per concept, the method reduces hyperparameter search requirements by 96.67% (from 210 evaluations to 7).
Sequential Unlearning: In tasks involving the sequential removal of nine objects, SAEmnesia improves unlearning accuracy by 28.4% compared to the baseline, demonstrating superior scalability.
Robustness: Under white-box adversarial attacks (UnlearnDiffAtk), SAEmnesia retains 57.50% unlearning accuracy compared to 34.20% for SAeUron. It also shows resilience in black-box attacks (Ring-A-Bell).
NSFW Suppression: On the I2P benchmark for nudity removal, the method achieves state-of-the-art results (9 detections vs. 18 for SAeUron) while preserving generation quality.

Significance and Claims
The paper positions SAEmnesia as a step forward toward "precise and controllable concept erasure" in generative models. By transforming distributed representations into concentrated, monosemantic features, the method addresses the fundamental interpretability challenges that have hampered previous unlearning techniques. The authors claim that this approach enables fine-grained control over generative behavior without modifying the base model weights, offering a fully reversible and transparent mechanism for removing harmful, inappropriate, or copyrighted content. The work emphasizes that while it serves as a research tool for safer deployment, it must be applied alongside rigorous evaluation and governance mechanisms to ensure reliability.

1. The Problem: The Tangled Yarn

2. The Solution: The "One-Concept, One-Neuron" Rule

3. The Eraser: A Precise Scalpel

4. What the Paper Actually Found

Summary

More like this