PaperlayerView exact revision on arXiv

Source-linked AI summary

SAeUron: Interpretable Concept Unlearning in Diffusion Models with Sparse Autoencoders

Bartosz Cywiński, Kamil Deja

arXiv:2501.18052v3cs.LGcs.AI

TL;DR

Diffusion models can generate harmful content, while existing unlearning methods often obscure how they alter the base model. SAeUron uses interpretable sparse-autoencoder features for targeted activation interventions, achieving strong unlearning results while preserving overall model performance and supporting multiple concepts and adversarial robustness.

  • Problem

    Diffusion models may generate harmful or undesirable content, while existing machine-unlearning methods often provide limited transparency into changes to the base model.

  • Method

    SAeUron trains sparse autoencoders on Stable Diffusion activations across denoising timesteps, then selects and ablates interpretable concept-specific features during inference.

  • Results

    SAeUron achieves state-of-the-art UnlearnCanvas performance, effectively removes nudity on I2P, and supports multiple-concept unlearning with robustness against adversarial attacks.

  • Takeaways & Limitations

    Feature-level interventions make the unlearning process more transparent while enabling precise concept removal without modifying the base model’s weights.

  • Takeaways & Limitations

    Inference-time intervention adds 1.92% overhead, and SAE training requires activations from a reasonable number of varied samples to retain remaining concepts well.

Abstract

from arXiv · show

Diffusion models, while powerful, can inadvertently generate harmful or undesirable content, raising significant ethical and safety concerns. Recent machine unlearning approaches offer potential solutions but often lack transparency, making it difficult to understand the changes they introduce to the base model. In this work, we introduce SAeUron, a novel method leveraging features learned by sparse autoencoders (SAEs) to remove unwanted concepts in text-to-image diffusion models. First, we demonstrate that SAEs, trained in an unsupervised manner on activations from multiple denoising timesteps of the diffusion model, capture sparse and interpretable features corresponding to specific concepts. Building on this, we propose a feature selection method that enables precise interventions on model activations to block targeted content while preserving overall performance. Our evaluation shows that SAeUron outperforms existing approaches on the UnlearnCanvas benchmark for concepts and style unlearning, and effectively eliminates nudity when evaluated with I2P. Moreover, we show that with a single SAE, we can remove multiple concepts simultaneously and that in contrast to other methods, SAeUron mitigates the possibility of generating unwanted content under adversarial attack. Code and checkpoints are available at https://github.com/cywinski/SAeUron.

1. Introduction

SAeUron addresses harmful or undesirable content in diffusion models with an interpretable, feature-level unlearning approach. It uses sparse autoencoder features to target concepts while preserving model performance, and reports strong results across concept, style, nudity, multi-concept, and adversarial evaluations.

  • Diffusion models can generate harmful or undesirable content, while retraining from scratch is costly and impractical for their enormous datasets.
  • Fine-tuning-based unlearning methods provide limited understanding of base-model changes and may leave targeted concepts vulnerable to adversarial recovery.
  • SAeUron trains sparse autoencoders on Stable Diffusion activations from all denoising timesteps to learn sparse, semantically meaningful features for concept blocking.
  • Interpretable, sparse features enable targeted analysis before unlearning and limit interference with remaining concepts, improving transparency over other methods.
  • SAeUron achieves state-of-the-art UnlearnCanvas performance across 20 objects and 50 styles, removes nudity on I2P, and supports multi-concept and adversarially robust unlearning.

2. Related Work

Related work spans sparse autoencoders, diffusion-model machine unlearning, and interpretability research. SAeUron adapts feature-based unlearning to diffusion models without modifying prompt embeddings or base-model weights.

  • Sparse Autoencoders: Sparse autoencoders learn compact, interpretable representations by encouraging only a small fraction of latent neurons to activate for each input.
  • Machine Unlearning in Diffusion Models: Machine unlearning can transform model outputs into independent feature contributions, allowing selected concepts to be removed by blocking corresponding weights or nodes.
  • Machine Unlearning in Diffusion Models: Most diffusion-model unlearning methods fine-tune models using objectives such as negative classifier-free guidance, re-steering losses, or saliency-selected parameter updates.
  • Machine Unlearning in Diffusion Models: Non-fine-tuning approaches modify text embeddings or cross-attention weights, whereas SAeUron changes neither prompt embeddings nor the base model’s weights.
  • Interpretability of Diffusion Models: Interpretation of specific functions and features in text-to-image diffusion components remains limited despite studies of semantic directions and cross-attention mechanisms.
  • SAeUron’s Position: SAeUron trains one SAE across multiple denoising steps of standard Stable Diffusion and uses its disentangled features for downstream unlearning.

3. Sparse Autoencoders for Diffusion Models

SAeUron adapts sparse autoencoders to cross-attention activations collected throughout Stable Diffusion’s denoising process. Its extensions and inference procedure preserve non-targeted features while ablating selected concept features.

  • SAEs are trained on cross-attention feature maps extracted at every denoising timestep, with each spatial position represented as a d-dimensional activation vector.
  • Each feature map has shape F_t ∈ R^h×w×d, where h and w index spatial dimensions and d is the feature-vector dimension.
  • Each spatial position becomes an independent SAE training sample, yielding h × w samples per feature map after discarding spatial-position information.
  • The standard single-layer ReLU SAE encodes d-dimensional activations into n latent features and decodes them with learned encoder and decoder weight matrices.
  • TopK activation retains only the k largest latent activations per vector, while BatchTopK selects B × k largest activations across each training batch.
  • During inference, selected SAE features are ablated after encoding and then decoded back, while remaining features stay unchanged to limit performance impact.

4. Method

SAeUron selects concept-specific SAE features using importance scores, then intervenes during diffusion inference to ablate those features while limiting unnecessary changes. The procedure inserts the SAE between U-Net blocks, modifies selected activations, and decodes them before passing them onward.

  • 4.1. Selection of SAE features for unlearning: SAeUron identifies SAE features for a target concept by computing importance scores at each denoising timestep.The score uses activations from target-concept data and broader diffusion-model activation data.
  • 4.1. Selection of SAE features for unlearning: Only τc highest-scoring features are targeted, while frequency filtering excludes dead features and features activating above the 99th percentile.Most features receive near-zero scores, so the intervention focuses on a small set of concept-specific features.
  • 4.2. SAE-based concept unlearning: During inference, the SAE is inserted between a selected U-Net cross-attention block and the subsequent block to process intermediate activations.The encoder decomposes each activation vector before selected features are modified and decoded.
  • 4.2. SAE-based concept unlearning: Selected features are ablated by scaling their activations with a negative multiplier normalized by their average activation on concept samples.This intervention removes the targeted concept’s influence from the activation vector.
  • 4.2. SAE-based concept unlearning: The method modifies a feature only when its activation exceeds the average activation threshold, preventing random ablation when scores are low.Modified representations are decoded with the error term preserved and passed to the next diffusion block.

5. Experiments

Experiments assess whether SAE features are meaningful, interpretable, and useful for targeted diffusion-model unlearning. SAeUron achieves strong unlearning and preservation performance while enabling feature inspection and reuse across concepts.

  • Technical details: SAEs are applied to selected cross-attention blocks using activations collected across denoising timesteps.Style filtering uses block up.1.2, while object filtering uses up.1.1; training collects feature maps across 50 denoising timesteps.
  • Interpreting SAE features: The experiments validate feature relevance through 5-nearest-neighbors classification on SAE activations, using 40 features with two features per class.Activations are extracted from the unconditional generation component to exclude text-embedding influence.
  • Interpreting SAE features: Score-selected SAE features show high discriminative accuracy across timesteps, while randomly selected features still outperform random guessing.Accuracy improves during denoising as object-relevant visual attributes emerge.
  • Interpreting SAE features: Style features activate on characteristic backgrounds while object features activate on the targeted object regardless of background or style.Activation heatmaps are normalized to [0, 1] and overlaid on generated image patches.
  • Interpreting SAE features: Selected unlearning features are often interpretable monosemantic concepts such as ears or paws, with activation patterns varying across diffusion frequencies.Visualizing removed features supports analysis of failure cases and improves transparency.
  • Interpreting SAE features: VLM-based annotation further provides text descriptions of SAE features, supporting their interpretation beyond manual visualization.The paper presents this as a general interpretability use of SAEs for diffusion models.
  • Concept unlearning with SAeUron: On UnlearnCanvas, SAeUron significantly outperforms other methods in style unlearning, ranks second in object unlearning, and achieves the best overall average performance.It maintains stable unlearning and preservation metrics, while requiring SAE training on only two cross-attention blocks and preserving overall model performance; it also removes nudity on I2P while preserving overall quality.

6. Additional experiments

Additional experiments test SAeUron on sequential multi-concept unlearning and adversarial prompting. The method retains performance as more concepts are removed and shows minimal degradation under attack, unlike competing approaches.

  • Unlearning of multiple concepts: SAeUron maintains superior unlearning effectiveness and retaining accuracy during sequential unlearning of six styles.Competing methods lose retention as the number of targeted concepts increases.
  • Unlearning of multiple concepts: When unlearning 49 of 50 styles simultaneously, SAeUron shows almost no performance degradation across three randomly selected combinations.The remaining style is used to evaluate preservation quality.
  • Robustness to adversarial attacks: The I2P evaluation uses nudity-unlearning metrics to assess removal of unwanted content alongside preservation of model quality.The supplied table caption identifies the benchmark and metric comparison but provides no numerical values.
  • Robustness to adversarial attacks: Under UnlearnDiffAtk adversarial prompts, SAeUron shows minimal unlearning-accuracy degradation, whereas competing methods suffer significant drops.The attack optimizes a 5-token prefix for 40 iterations with learning rate 0.01.

7. Limitations

SAeUron has inference, data, deployment, and concept-scope limitations. These trade-offs include overhead, activation-storage demands, restricted access requirements, and weaker performance for similar or abstract concepts.

  • Inference cost: 1.92% inference overhead slightly slows generation.The method operates during inference rather than only during training.
  • Training cost: Training requires activations from many diverse samples to retain remaining concepts well, adding computational overhead for single-concept unlearning.
  • Deployment: The blocking mechanism is practical only when users lack direct model access because open-source users could remove it.
  • Training cost: SAE training demands significant activation storage for large datasets, although the method has low GPU and storage requirements relative to other techniques.
  • Scope: Performance is limited when unlearning concepts while preserving similar concepts or targeting abstract concepts without distinct visual characteristics.

8. Conclusions

SAeUron uses sparse, interpretable SAE features to localize and ablate unwanted concepts in text-to-image diffusion models. The method achieves state-of-the-art unlearning while maintaining overall performance, with robustness to adversarial attacks and support for sequential multi-concept unlearning.

  • SAeUron leverages sparse autoencoders to unlearn concepts from text-to-image diffusion models.
  • Sparse, interpretable features enable precise concept-specific interventions while maintaining overall model performance.
  • Interpretable features make the unlearning process more transparent by clarifying which features are targeted.
  • SAeUron achieves state-of-the-art UnlearnCanvas results, mitigates adversarial attacks, and can sequentially unlearn multiple concepts.

Impact Statement

The paper’s impact spans safer and more transparent concept removal, while also identifying deployment and misuse considerations. Its experiments connect localized cross-attention interventions with object, style, and sequential unlearning behavior.

  • Impact Statement: SAeUron is designed to block selected unwanted, biased, or harmful content, but it could also be misused to promote such content.
  • SAE behavior: BatchTopK SAEs allocate more active features to detailed or central image regions and fewer to background areas.
  • SAE behavior: SAEs distribute active features unevenly across image samples, indicating differing sample importance for reconstruction.
  • Cross-attention localization: Cross-attention ablations identify separate blocks whose removal most affects object quality or original image style.
  • Sequential unlearning: Sequential evaluation requires progressively removing concepts while retaining previously unlearned targets and overall generative performance.
  • Sequential unlearning: SAeUron maintains high unlearning accuracy and outperforms competing methods in retaining non-targeted concepts across sequential phases.

E. SAE trainings details

The appendix details SAE training, feature localization, steering, and timestep analyses. These experiments characterize feature sparsity, generalization, spatial allocation, and the temporal emergence of style-related signals.

  • Training setup: BatchTopK SAEs use k = 32, expansion factor 16, Adam optimization, learning rate 0.0004, batch size 4096, and decoder-weight normalization.
  • Training setup: The object SAE is trained for 5 epochs and the style SAE for 10 epochs, with dead latents defined by no activation over 10M samples.
  • Feature density: Very few features activate very frequently, while the object SAE contains dead latents and the style SAE does not.
  • Training cost: Both SAEs were trained on one NVIDIA RTX A5000 GPU; object and style training took 27 hours 40 minutes and 59 hours 1 minute, respectively.
  • Generalization: Training on half of the styles yields over 50% unlearning accuracy on out-of-distribution styles, demonstrating generalization to unseen concepts.
  • Activation steering: Feature steering increases selected SAE feature activations during diffusion and produces images exhibiting the corresponding artistic styles.
  • Timestep analysis: Style-classification accuracy remains similar with all features and score-based selection, increasing from approximately the 30th timestep.
  • Timestep analysis: Score thresholds decrease during denoising, indicating that fewer features remain highly relevant to specific concepts later in generation.

K. UnlearnDiffAtk evaluation of object unlearning

Under UnlearnDiffAtk, SAeUron’s object-unlearning accuracy appears to drop, but image-level and CLIP-based checks indicate that the targeted object generally remains absent. The evaluation therefore distinguishes classifier-based attack success from actual target-object generation.

  • Evaluation interpretation: UnlearnDiffAtk attack success is partly driven by evaluating each iteration through the classifier’s argmax prediction.This explains why reported unlearning accuracy can fall even when the attack does not restore the unlearned object.
  • Target similarity: CLIPScore between the target object name and generated images remains nearly unchanged before and after attack.The result indicates that adversarial prompting rarely leads to generation of the targeted object.
  • Object-level robustness: Despite successful attacks, the targeted object remains absent from SAeUron-generated images.The attack can cause the classifier to make largely random predictions when no replacement object appears.

L. Adversarial attack evaluation for nudity unlearning

SAeUron remains robust on adversarial nudity prompts, although its performance is constrained for abstract concepts and highly similar classes. The additional analyses connect robustness and efficiency with feature-level interpretability and reveal where unlearning causes collateral degradation.

  • Adversarial nudity evaluation: 1.4% Post-ASR is achieved on 143 adversarial nudity prompts, with the smallest difference between Pre-ASR and Post-ASR.Pre-ASR measures nudity generation before attack, whereas Post-ASR measures it after prompts are adversarially tuned.
  • Efficiency and scalability: SAeUron achieves good unlearning with training sets of 100–1000 images and is more efficient than most compared methods.The comparison evaluates five SAE training-set sizes while keeping the training setup fixed at five epochs.
  • Collateral effects: Highly similar classes can impair one another because score-based selection overlaps features during early denoising timesteps.The overlapping feature identified for Cats and Dogs relates to generating both animals’ heads.
  • Scope limitations: SAeUron underperforms on the full I2P benchmark, performing on par only with FMN for broad inappropriate-content unlearning.The authors attribute this to SAE features corresponding mainly to individual visual objects rather than abstract notions such as hate, harassment, or violence.
  • Preservation of other concepts: Across most UnlearnCanvas classes, SAeUron removes the target while preserving remaining-class accuracy, with degradation concentrated among similar classes.The Cats–Dogs interaction is specifically reported as a case where removing one class harms the other.
  • Feature interpretability: Feature annotations successfully identify visual characteristics associated with targeted concepts, including rabbit facial features and ears.The annotations are generated from activation overlays interpreted in the context of original images.
Loading 2501.18052v3…