Supervised sparse auto-encoders as unconstrained feature models for semantic composition

TL;DR

The paper tackles the limitations of unsupervised sparse auto-encoders by introducing a decoder-only, supervised SSAE grounded in unconstrained feature model theory. It defines a predefined sparse concept dictionary and learns a decoder to reconstruct feature spaces from sparse latent codes, enabling compositional generalization and feature-level editing without modifying prompts or employing $L_1$ regularization. Applied to Stable Diffusion 3.5 prompt embeddings, the approach demonstrates that unseen concept combinations can be composed and edited modularly, suggesting a scalable path toward interpretable, structured interfaces for foundation models. While promising, the authors acknowledge limitations such as reliance on a predefined concept set and limited experimental scope, outlining clear directions for broader evaluation and encoder integration. Overall, the work advances interpretable, semantically grounded interventions in large models by combining theoretical insights from unconstrained feature models with practical, sparse supervision.

Abstract

Sparse auto-encoders (SAEs) have re-emerged as a prominent method for mechanistic interpretability, yet they face two significant challenges: the non-smoothness of the $L_1$ penalty, which hinders reconstruction and scalability, and a lack of alignment between learned features and human semantics. In this paper, we address these limitations by adapting unconstrained feature models-a mathematical framework from neural collapse theory-and by supervising the task. We supervise (decoder-only) SAEs to reconstruct feature vectors by jointly learning sparse concept embeddings and decoder weights. Validated on Stable Diffusion 3.5, our approach demonstrates compositional generalization, successfully reconstructing images with concept combinations unseen during training, and enabling feature-level intervention for semantic image editing without prompt modification.

Figures (7)

• Figure 1: Editing workflow: We train weights of a decoder-only SSAE $\mathbf{W}_2$ and the dictionary of concepts (i.e. trainable parameters of $\mathbf{Y}$) to reconstruct prompt embeddings for image generation. We use the trained dictionary (sparse latent space) and $\mathbf{W}_2$ to add/modify/remove to guarantee safer image generation at inference time.
• Figure 2: Initial prompt: "A brune girl with blue eyes on horseback across a plain, wearing a red t-shirt and a hat, holding a gun, looking in front of her."; then we perform the swap between "brune" and "blond" via our transformation in the sparse latent space learnt by our decoder-only SSAE.
• Figure 3: Initial prompt: "A brune girl with blue eyes on a boat, wearing a red t-shirt and a cap, holding a gun, looking in front of her."; then we perform the swap between "brune" and "blond" via our transformation in the sparse latent space learnt by our decoder-only SSAE.
• Figure 4: Initial prompt: "A blong girl with brown eyes sitting at a bar, wearing a blue t-shirt and a baseball cap, holding a gun, looking in front of her." (top-left), then we applied our transformation in the sparse latent space learnt by our decoder-only SSAE to successively: remove the concept of "holding a gun" (top-right), insert the concept of "holding a coffee" (bottom-left), swap it with "holding a coca-cola" (bottom-right), showcasing compositional generalisation.
• Figure 5: Initial prompt: "A blond girl with blue eyes in a car, wearing a red t-shirt and a cap, holding a gun, looking in front of her." (left), then we applied our transformation in the sparse latent space learnt by our decoder-only SSAE to successively: remove the concept of "holding a gun" (middle), swap the concepts of "blond hair" and "brune hair"(right), showcasing compositional generalisation.