SAEmnesia: Erasing Concepts in Diffusion Models with Supervised Sparse Autoencoders

Concept unlearning in diffusion models is hampered by feature splitting, where concepts are distributed across many latent features, making their removal challenging and computationally expensive. We introduce SAEmnesia, a supervised sparse autoencoder framework that overcomes this by enforcing one-to-one concept-neuron mappings. By systematically labeling concepts during training, our method achieves feature centralization, binding each concept to a single, interpretable neuron. This enables highly targeted and efficient concept erasure. Compared to the state-of-the-art sparse autoencoder-based unlearning approach, SAEmnesia reduces hyperparameter search by 96.67\% and achieves a 9.22\% improvement on the UnlearnCanvas benchmark for objects. Our method also shows superior scalability in sequential unlearning, improving accuracy by 28.4\% when removing nine objects, establishing a step forward for precise and controllable concept erasure. Moreover, SAEmnesia effectively suppresses nudity on the I2P benchmark and remains robust to adversarial attacks. Source code available at https://github.com/EIDOSLAB/SAEmnesia.

AI systems that generate images from text descriptions can also produce harmful or copyrighted content, so we want ways to make them "forget" specific things - a particular artist's style, say, or nudity - without retraining the whole system or degrading everything else it can do. The difficulty is that, inside these systems, a single idea like "dog" isn't stored in one tidy place; it's smeared across many internal components, so removing it cleanly is hard and requires a slow, expensive trial-and-error search. We developed a method, SAEmnesia, that reorganizes the system's internal memory during training so that each concept is bound to a single dedicated component, like giving every idea its own labeled drawer instead of scattering it through a messy filing cabinet. To make the system forget something, we simply switch off the one drawer that holds it. This makes erasing a concept far more precise, roughly 30 times cheaper to tune, and harder for adversaries to undo. As AI-generated images become widespread, tools like this offer a transparent, controllable way to keep these systems safer and more trustworthy.