- The paper introduces SAEParate, a SAE-based framework that employs a supervised contrastive loss to explicitly disentangle latent concept representations in diffusion models.
- It demonstrates state-of-the-art performance with joint unlearning accuracy over 86% and minimized latent overlap, ensuring minimal interference with non-target concepts.
- The method enables efficient, inference-time concept suppression without modifying base model weights, promising robust deployment in safety-critical generative applications.
Disentangled Sparse Representations for Concept-Separated Diffusion Unlearning
Overview and Motivation
The problem of selective concept unlearning in text-to-image (T2I) diffusion models is increasingly critical due to the proliferation of sensitive, potentially harmful, or undesired generations. Traditional methods have relied on model parameter modification or post-hoc suppression mechanisms. Sparse autoencoder (SAE)-based approaches, such as SAeUron, are attractive due to their ability to excise targeted concepts by manipulating interpretable latent features at inference time, preserving base model weights. However, conventional SAEs, trained with sparsity objectives alone, do not guarantee concept separation in their latent representations—shared features across concepts undermine precise, surgical unlearning, especially for compositional or closely related concepts.
The paper "Disentangled Sparse Representations for Concept-Separated Diffusion Unlearning" (2605.12122) introduces SAEParate, an SAE-based framework for T2I diffusion models that enforces explicit concept disentanglement in the learned latent space via a supervised contrastive objective. This enables more precise suppression of designated concepts with minimal collateral interference, even in demanding settings such as joint style-object unlearning where conventional latent representations are highly entangled.
Methodological Contributions
1. Concept-Aware Contrastive SAE Training
SAEParate modifies the standard unsupervised SAE training procedure by introducing a supervised, concept-aware contrastive loss in the latent space. Using a multiview batching scheme derived from the temporal structure of diffusion models, the approach pulls together latent representations of the same concept (across closely related denoising timesteps) and pushes apart representations of different concepts, including compositional (style-object) cases. Crucially, partial matches—samples sharing only one attribute (e.g., object or style)—are designated as "hard negatives" and receive extra repulsion within the contrastive loss. This induces separation not just at the class level, but at the level of fine-grained compositional concepts.
2. GeLU-Enhanced Encoder Architecture
To support the increased discriminative pressure from the contrastive objective, the SAE encoder is augmented with an additional hidden layer and a GeLU activation. This deeper, nonlinear architecture improves representational expressivity and enables more effective clustering of concepts in the latent space compared to vanilla single-layer, ReLU-based SAEs.
3. Inference-Time Concept Unlearning in Disentangled Space
After training, concept-specific feature sets are identified via feature-importance scores that reflect exclusivity and discriminativity in the contrastively trained latent space. Unlearning is performed at inference time by suppressing only the latents associated with the target concept, minimizing degradation in non-target generations.
Experimental Evaluation
Experiments are conducted on the UnlearnCanvas benchmark, which provides rigorous evaluations for both single-attribute (object or style) and joint attribute (style-object) unlearning, as well as quantitative metrics reflecting both the degree of unlearning and retention of unaffected concepts. Key baselines include state-of-the-art parameter-editing and non-parametric approaches, notably SAeUron.
Results indicate clear, strong performance gains for SAEParate:
- State-of-the-art joint style-object unlearning: SAEParate achieves average joint unlearning accuracy (UA) of 86.2%, style consistency (SC) of 81.9%, and object consistency (OC) of 91.4%. Previous best baselines, including SAeUron, show severe trade-offs between suppression and retention, with retention metrics often collapsing for non-target concepts.
- Minimized latent overlap: Quantitative overlap and cluster analyses show a complete resolution of entangled or redundant feature selection—pairwise overlaps between concept feature sets are reduced to near zero, and the maximal cosine similarity between distinct concept groups is sharply decreased.
- Compact concept encoding: k-sparse probing reveals that concept-discriminative information is localized to a minimal set of latents, often just k=1 or 2 per concept, compared to higher-k requirements and greater redundancy in vanilla SAEs.
- Ablation and generalization: Removal of either the contrastive objective or GeLU enhancement degrades both performance (in unlearning/retention trade-off) and clustering. SAEParate demonstrates generalization to unseen concepts and compositional settings, outperforming baselines in out-of-distribution and sequential unlearning scenarios.
- Efficiency: The addition of the contrastive objective and enhanced encoder only marginally increases memory and storage requirements over SAeUron, maintaining practical deployment viability.
Theoretical and Practical Implications
From a theoretical perspective, SAEParate substantiates the hypothesis that conventional sparsity constraints are insufficient for robust concept unlearning: explicit supervision via contrastive loss is essential for disentangling compositional representations in the high-capacity intermediate activations of large diffusion models. The method operationalizes a structured form of supervised contrastive learning, extending its domain of application beyond instance discrimination to structural/compositional disentanglement.
Practically, the approach enables lightweight, post-hoc, inference-time concept unlearning without base model modification. This is particularly salient for commercial or safety-critical deployments where fine-grained control of generative behaviors is demanded, and costly retraining or parameter editing is undesirable. The clear separation in latent space also enhances interpretability and future maintainability of unlearning pipelines.
Limitations and Future Directions
The principal limitation is reliance on concept labels for supervised contrastive pretraining. Current generalization, while markedly improved over baselines, is intrinsically conditioned on the available labeled concepts and their combinatorial structure. Extension to unsupervised or weakly-supervised contrastive objectives, as well as application to other modalities (video, multimodal generation), is a promising avenue. Incorporation of domain adaptation or continual learning strategies could further improve robustness under non-stationary or privacy-preserving training regimes.
Conclusion
SAEParate establishes a new paradigm for sparse, interpretable, and fine-grained concept unlearning in T2I diffusion models by combining a supervised contrastive latent objective with an enhanced nonlinear encoder. The resulting disentangled latent space enables surgical suppression of specified concepts with state-of-the-art precision, particularly in the challenging setting of joint attribute unlearning. This work clarifies both the representational pathologies of prior SAE-based approaches and practical pathways toward reliable, scalable unlearning mechanisms in large generative models (2605.12122).