SAeUron: Interpretable Concept Unlearning in Diffusion Models with Sparse Autoencoders
Bartosz CywinskiKamil Deja
Introduces SAeUron, a mechanistically interpretable method that trains sparse autoencoders across denoising timesteps to precisely identify and ablate unwanted visual concepts from text-to-image diffusion models without degrading overall generation quality or succumbing to adversarial attacks.
Modern text-to-image diffusion models can inadvertently create copyrighted, explicit, or harmful imagery, creating substantial safety, ethical, and legal risks. Retraining foundation models from scratch to remove unwanted content is computationally prohibitive, making machine unlearning an essential alternative. However, current unlearning techniques primarily rely on fine-tuning model weights or applying negative gradients. These methods function as opaque "black boxes" that frequently degrade overall image generation quality, struggle when removing multiple concepts sequentially, and often only superficially mask concepts, leaving models vulnerable to adversarial prompt attacks.
The article evaluates and demonstrates SAeUron, a transparent and interpretable unlearning framework that identifies and blocks specific visual concepts within text-to-image diffusion models using sparse autoencoders (SAEs), neural networks that decompose complex internal representations into distinct, human-interpretable components without modifying the base model's core weights.
To establish credibility, the evaluation was conducted across rigorous, standardized benchmarks using the Stable Diffusion architecture. The researchers trained lightweight SAEs on internal activations gathered across multiple denoising steps from key cross-attention blocks—specifically targeting an object-specialized block and a style-specialized block. Using an importance-scoring mechanism on validation prompt activations, SAeUron pinpoints concept-specific features and scales them down during inference. The system was benchmarked against leading unlearning methods on the UnlearnCanvas dataset (spanning 20 objects and 50 styles), the I2P inappropriate content benchmark, and the UnlearnDiffAtk adversarial evaluation framework.
The experimental findings demonstrate substantial improvements in both unlearning precision and operational stability. First, on the UnlearnCanvas benchmark, SAeUron achieved the highest overall average score of 94.03%, outperforming all existing methods in style unlearning (95.80% unlearning accuracy, 99.10% in-domain retain accuracy, and 99.40% cross-domain retain accuracy) while maintaining competitive performance in object unlearning (78.82% unlearning accuracy and over 95% retention). Second, SAeUron excelled at removing explicit content on the I2P benchmark, reducing total detected nudity instances to 18 (compared to 743 in the base model and 28 to 838 across competing baselines) while maintaining strong image quality. Third, under adversarial prompt attacks designed to force models into generating unlearned content, SAeUron exhibited high robustness, maintaining unlearning integrity with a post-attack success rate of only 1.40% on nudity prompts, whereas competing methods saw attack success rates surge up to 70–98%. Finally, SAeUron proved uniquely scalable for multi-concept erasure; in a stress test removing 49 out of 50 artistic styles simultaneously, the model retained 99.29% unlearning accuracy and over 95% preservation of the remaining style, avoiding the catastrophic degradation observed in fine-tuning approaches.
These findings indicate that generative concept removal can be achieved efficiently and transparently without retraining or destabilizing underlying foundation model parameters. By intervening only on interpretable latent features during inference, organizations can significantly reduce compute costs, storage demands (requiring only 0.2 GB of storage compared to 4 GB for fine-tuned checkpoints), and compliance risks. Furthermore, because target features can be inspected visually or annotated with language models prior to ablation, stakeholders gain mechanistic transparency over why and how content is blocked.
Based on these results, deploying inference-time sparse activation blocking is a highly effective strategy for closed-system image generation APIs and hosted platforms requiring rigorous safety filters and multi-concept unlearning. Organizations should prioritize mapping concept-specific feature dictionaries rather than maintaining separate fine-tuned model checkpoints for each safety constraint. Further work and pilot evaluations are recommended before full deployment, particularly to refine feature separation for closely related visual categories (such as visually overlapping animal classes) and to establish complementary guardrails for broad, abstract harms like hate or harassment that do not correspond to concrete visual geometries.
Decision-makers should note certain operational boundaries. SAeUron introduces a minor 1.92% latency overhead during image generation, and because it operates at inference time, it is primarily suitable for hosted or access-controlled platforms; open-source model releases could have the activation-blocking module bypassed by end users. Additionally, because the SAEs are trained on visual activations, the method's effectiveness is exceptionally high for well-defined visual elements (such as nudity, artistic styles, and objects) but diminishes for abstract concepts lacking consistent visual features.
- Paper: Sparse Autoencoders Find Highly Interpretable Features in Language Models, Hoagy Cunningham et al. (2023). This paper establishes how sparse autoencoders decompose neural activations into interpretable features—the core representation SAeUron later uses to identify concepts for removal.
- Paper: What the DAAM: Interpreting Stable Diffusion Using Cross Attention, Raphael Tang et al. (2023). Its method for tracing prompt-word influence through Stable Diffusion cross-attention provides useful groundwork for understanding the activations SAeUron analyzes.
- Paper: Cones: Concept Neurons in Diffusion Models for Customized Generation, Zhiheng Liu et al. (2023). Cones shows how concept-specific signals can be localized and suppressed in diffusion-model attention layers, preparing readers for SAeUron’s activation-level concept ablation.
- Paper: Safe Latent Diffusion: Mitigating Inappropriate Degeneration in Diffusion Models, Patrick Schramowski et al. (2023). Safe Latent Diffusion introduces inference-time suppression of harmful generations and the I2P evaluation context that help situate SAeUron’s safety intervention.
No sufficiently relevant recommendations were found.
