ClasSAE: Class-Aligned Sparse Autoencoders via Differentiable Feature-Class Affinity

Jakub Stępień; Marcin Mazur; Jacek Tabor; Przemysław Spurek

cs.CVOct 2, 2026

ClasSAE: Class-Aligned Sparse Autoencoders via Differentiable Feature-Class Affinity

Authors: Jakub Stępień, Marcin Mazur, Jacek Tabor, Przemysław Spurek

Organizations: Jagiellonian University · IDEAS Research Institute

Abstract

Sparse Autoencoders (SAEs) began as an unsupervised tool for decomposing neural representations into sparse, interpretable features, and are increasingly used not only for passive analysis but also for active interventions such as unlearning, bias mitigation, and concept editing. A central challenge for these editing and steering methods is reliably matching features to target concepts; most current approaches address this by computing post-hoc scores over an already-trained, frozen dictionary. We instead introduce ClasSAE, a novel method that both automatically assigns classes to features and guides the encoder toward class-separable representations during training. Specifically, we apply a differentiable top- $k$ operator to a trainable feature--class affinity matrix with per-feature budgets, coupling the features selected for each sample to the classes they are trained to represent. Because gradients flow through the selection of active features rather than only through their magnitudes, the encoder and the affinity matrix co-adapt rather than being fit in separate stages. The result is a dictionary that is both class-separable and class-annotated, with no need for post-hoc probing. We propose three variants for enforcing sparsity within this framework, which achieve comparable overall performance with slightly different trade-offs. Using CLIP ViT-L/14 embeddings on ImageNet, we show that the learned affinity matrix agrees closely with an independently estimated post-hoc feature-class matrix computed on held-out data. The model also supports direct class prediction from the encoder and affinity matrix alone, without a separately fitted classifier, and its more class-aligned encoder yields improved separation in Targeted Probe Perturbation evaluations. https://github.com/St0pien/ClasSAE.

Figures & tables

Figure 1: ClasSAE jointly learns a sparse dictionary and a class affinity matrix, enabling classifier-free class prediction and targeted editing: removing only the features exclusive to a class erases that class from the reconstruction while shared features are preserved.

Figure 2: Feature activation heatmaps for ClasSAE trained on DINOv3 patch tokens, using whole-image class labels only. Top: features claimed by “Siamese cat, Siamese.” Bottom: features claimed by “American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier.” Each panel reports the feature’s budget $k_{i}$ , its claim strength for the labeled class, and its mean activation. Details in Appendix F .

Panel of Figure 3 — Figure 3: (a) Each dictionary coordinate holds a row of trainable class logits $\Lambda_{i}$ (grey). A trainable vector $\kappa$ splits a global association budget $K_{\mathrm{tot}}=\rho d$ into per-feature budgets $k_{i}$ (orange), and Soft Top- $k_{i}$ turns each row into a claim $M_{i}$ carrying exactly $k_{i}$ units of mass (green). Budgets need not be equal: row 1 is a specialist ( $k=1$ ), row 7 a generalist ( $k=12$ ), and what one row takes another must give up. (b) For a sample of class $c$ , the SAE selects $\hat{k}$ coordinates (blue) while column $M_{:,c}$ marks the coordinates that claimed $c$ (green). The score is the fraction of the selected set that is claimed (shaded); the double arrow indicates that minimizing $-s_{c}$ pushes from both sides at once, since the encoder must rank claimed features into the top- $\hat{k}$ while the table must claim features that actually fire. Both panels are drawn in the hard limit; during training $S(x)$ and $F_{c}$ are the soft weights $\pi(x)$ and $M_{:,c}$ .

Architecture	$L_{0}$	FVE $\uparrow$	Alive $\uparrow$
BatchTopK	381.055 ±0.073	0.973 ±0.000	0.713 ±0.013
JumpReLU	383.346 ±3.158	0.986 ±0.000	0.579 ±0.019
Matryoshka	381.001 ±0.199	0.968 ±0.000	0.377 ±0.027
Standard	345.936 ±1.021	0.899 ±0.000	0.671 ±0.005
Top-AFA †	1184.456 ±24.535	0.991 ±0.001	0.662 ±0.029
TopK	380.000 ±0.000	0.970 ±0.000	0.455 ±0.010

Table 1: Reconstruction quality. Comparison at matched target

L_{0}\approx 380

L_{0}

is the realized

L^{0}

-norm of the code; FVE is the fraction of input variance explained by the reconstruction; Alive is the fraction of dictionary coordinates that fire on at least one validation sample. Bold marks the best result among architectures matched to the target

L_{0}

; † Top-AFA selects its own

L_{0}

and is excluded from this comparison.

Architecture	Mean CS $\uparrow$	Median CS $\uparrow$	Mean Precision $\uparrow$	Mean Recall $\uparrow$	Mean F1 $\uparrow$
BatchTopK	0.410 ±0.009	0.269 ±0.043	0.348 ±0.007	0.347 ±0.008	0.346 ±0.008
JumpReLU	0.287 ±0.010	0.000 ±0.000	0.465 ±0.016	0.463 ±0.016	0.463 ±0.016
Matryoshka	0.092 ±0.002	0.000 ±0.000	0.178 ±0.018	0.178 ±0.017	0.178 ±0.017
Standard	0.555 ±0.004	0.848 ±0.000	0.683 ±0.004	0.681 ±0.004	0.682 ±0.004
Top-AFA	0.266 ±0.016	0.127 ±0.048	0.180 ±0.005	0.179 ±0.005	0.179 ±0.005
TopK	0.195 ±0.003	0.000 ±0.000	0.177 ±0.004	0.176 ±0.004	0.176 ±0.004

Table 2:

M

-honesty. Mean/median cosine similarity (CS) and precision/recall/F1 between the train-side matrix (

M

for ClasSAE, post-hoc

\hat{M}

for baselines) and the test-side

\hat{M}

. ClasSAE variants demonstrate that

M

carries a real, class-faithful association, not just an artifact of training.

Architecture	Acc $\uparrow$	Acc Top 5 $\uparrow$	Acc Top 10 $\uparrow$	F1 $\uparrow$	Margin $\uparrow$
BatchTopK	0.002 ±0.000	0.007 ±0.000	0.015 ±0.001	0.000 ±0.000	-0.148 ±0.001
JumpReLU	0.678 ±0.003	0.855 ±0.004	0.893 ±0.004	0.682 ±0.002	0.023 ±0.000
Matryoshka	0.002 ±0.000	0.006 ±0.000	0.015 ±0.000	0.000 ±0.000	-0.137 ±0.002
Standard	0.760 ±0.000	0.933 ±0.000	0.963 ±0.000	0.757 ±0.000	0.013 ±0.000
Top-AFA	0.004 ±0.002	0.012 ±0.002	0.022 ±0.005	0.003 ±0.001	-0.044 ±0.007
TopK	0.611 ±0.008	0.779 ±0.010	0.821 ±0.008	0.607 ±0.004	0.006 ±0.000

Table 3: Out-of-the-box classification. Acc , Acc Top 5 , Acc Top 10 : top-1, top-5, and top-10 accuracy. F1 : macro F1. Margin : gap between the top two class scores. ClasSAE

L^{1}

is the top scorer on every metric, and all ClasSAE variants trained with true labels achieve a larger margin than any baseline. Standard is the strongest baseline: it outperforms the soft-gradient ClasSAE on accuracy and F1 and slightly exceeds ClasSAE Hard on Acc Top 10.

Architecture	TPP 20 $\uparrow$	TPP 200 $\uparrow$	TPP 600 $\uparrow$	TPP PMI $\uparrow$	Edit Strength $\downarrow$
BatchTopK	0.450 ±0.003	0.266 ±0.011	0.051 ±0.003	0.277 ±0.016	70.850 ±4.841
JumpReLU	0.102 ±0.009	0.463 ±0.001	0.281 ±0.006	0.453 ±0.003	87.082 ±5.024
Matryoshka	0.431 ±0.008	0.113 ±0.009	0.015 ±0.003	0.104 ±0.004	162.848 ±2.378
Standard	0.428 ±0.003	0.465 ±0.001	0.432 ±0.003	0.430 ±0.001	62.379 ±0.196
Top-AFA	0.174 ±0.014	0.464 ±0.004	0.356 ±0.019	0.012 ±0.002	31.282 ±3.376
TopK	0.423 ±0.010	0.304 ±0.007	0.067 ±0.008	0.309 ±0.010	124.047 ±1.017

Table 4: Targeted probe perturbation. TPP at

k\in\{20,200,600\}

and under the PMI-graded edit (TPP PMI); Edit Strength is the mean

L^{1}

norm of the TPP PMI edit. ClasSAE Hard leads at

k{\geq}200

and ClasSAE

L^{1}

under TPP PMI, with the two smallest edit strengths.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Architecture	Graded vs. Shuffled Forget Drop $\uparrow$	Retain Drop $\downarrow$	Edit Strength $\downarrow$
BatchTopK	-0.284 ±0.010	0.075 ±0.009	70.850 ±4.841
JumpReLU	0.170 ±0.017	0.034 ±0.003	87.082 ±5.024
Matryoshka	0.000 ±0.000	0.421 ±0.014	162.848 ±2.378
Standard	0.224 ±0.046	0.055 ±0.000	62.379 ±0.196
Top-AFA	-0.442 ±0.032	0.001 ±0.000	31.282 ±3.376
TopK	0.000 ±0.000	0.199 ±0.006	124.047 ±1.017

Appendix

Table 5: Joint-classifier unlearning. Graded-vs-shuffled forget-accuracy drop, retain-accuracy drop, and edit strength (mean

L^{1}

norm); the two drops are measured relative to baseline classifier accuracy. ClasSAE variants dominate on forget drop while keeping edit strength small.

Sparse Autoencoders

Sparse autoencoders decompose neural network activations into sparse features to study, interpret, and control learned representations.

See all 218 papers

Jun 23, 2026

Evaluating the Interpretability of Sparse Autoencoders with Concept Annotations

Sparse autoencoders (SAEs) are increasingly used to extract interpretable concepts from vision and vision language models, yet existing evaluation methods largely rely on proxy metrics or qualitative inspection rather than measuring semantic correspondence. We present a human-grounded evaluation framework that quantifies alignment between SAE latents and human-annotated concepts, without requiring user studies, and validate this matching through targeted attribute perturbations. To enable this intervention-style evaluation in vision, we construct synCUB and synCOCO, synthetic benchmarks of paired images that differ in exactly one attribute. We introduce Fully-Binary Matching Pursuit (FBMP), a coalition-based matching procedure that supports many-to-one mappings between SAE latents and annotated concepts, and consistently outperforms one-to-one baselines. For functional validation, we propose a Targeted Attribute Perturbation Alignment Score (TAPAScore), which tests whether matched concepts respond selectively and in the expected direction under targeted image-level attribute perturbations. Under sanity checks, our matching and TAPAScore are the only evaluated metrics that reliably distinguish trained SAEs from untrained ones. Across SAEs trained on CLIP and DINOv2 embeddings, we find that increased overcompleteness can reduce perturbation alignment, indicating a reduction in interpretability. Our evaluation framework suggests that moderate dictionary sizes provide the best trade-off, yielding the most interpretable SAEs. Code and datasets are available at https://github.com/JonasKlotz/sae-concept-eval.

Jonas Klotz, Cassio F. Dantas, Pallavi Jain +2

The Berlin Institute for the Foundations of Learning and Data (BIFOLD) · Technische Universität Berlin, Germany · INRAE +2

VLM Interpretability Sparse Autoencoders Concept-Based Explanations Explainability Evaluation

Ask PDF

Jul 27, 2026

Sparse Autoencoders Encode Both Concepts and Functions: The Downstream Geometry of Feature Effects

The wide-scale use of sparse autoencoders (SAEs) as interpretability tools is limited by inconsistent links between SAE features and model behavior. Features with clear activation descriptions may have weak or unexpected causal effects; steering can vary across prompts or oppose the intended direction; and activation-based feature selection can miss features that produce the desired output change. Prior work has studied feature geometry inside the model, where features are computed. We instead study the geometry of changes in model logits caused by feature interventions. We introduce Feature-Effect Geometry Analysis (FEGA), an unsupervised framework that removes the same active SAE feature across contexts and analyzes the resulting cloud of logit changes. Across SAE variants, consistent one-dimensional effects are rare: few features behave like reusable directions. To interpret this variation, we distinguish value-like features, tied to static information such as factual attributes, from pointer-like features, associated with context-dependent operations. Value-like features more often exhibit structured, low-dimensional effects, although these effects typically span several directions. Pointer-like features, by contrast, predominantly exhibit diffuse effects. Our results show that a feature can be interpretable and causally relevant without providing a stable direction for steering.

Phu Gia Hoang, Anwoy Chatterjee, Tanmoy Chakraborty +2

UKP Lab, Technical University of Darmstadt · Department of Electrical Engineering Indian Institute of Technology Delhi, India

Sparse Autoencoders Mechanistic Interpretability Causal Interventions in Language Models

Ask PDF

May 18, 2026

Aligned Training: A Parameter-Free Method to Improve Feature Quality and Stability of Sparse Autoencoders (SAE)

Sparse autoencoders (SAEs) are one of the main methods to interpret the inner workings of deep neural networks (DNNs), decomposing activations into higher-dimensional features. However, they exhibit critical shortcomings where a large fraction of features are never activated and are unstable. Despite variants of SAEs that attempt to mitigate these issues, they require additional data, resampling, or training. We propose the aligned training, a parameter-free reparameterization of SAEs that simultaneously improves reconstruction quality, eliminates dead features, and significantly enhances stability across training seeds. Our approach is motivated by an overlooked observation that SAE feature quality, measured by the inner product between encoder and decoder directions (which we call the alignment score), follows a bimodal distribution across all modern architectures. The proposed aligned training enforces a geometric constraint between the encoder and decoder such that their inner product equals one for every feature, which removes a source of degeneracy in the SAE training without adding any hyperparameters. Across multiple models, dictionary sizes, and sparsity levels, the aligned training shows Pareto improvements on the SAEBench benchmarks. Beyond improving dead features, stability and reconstruction, our method readily integrates with techniques in mechanical interpretability such as Top/BatchTop-K architectures and p-Annealing. Overall, the aligned training substantially improves feature quality and stability of SAE without computational complexity or cost.

Michał Brzozowski, Neo Christopher Chung

Samsung AI Center, Warsaw, Poland · University of Warsaw, Poland

Sparse Autoencoders Mechanistic Interpretability Representation Alignment

Ask PDF

Abstract

Sparse Autoencoders

Explore similar work

Evaluating the Interpretability of Sparse Autoencoders with Concept Annotations

Sparse Autoencoders Encode Both Concepts and Functions: The Downstream Geometry of Feature Effects

Aligned Training: A Parameter-Free Method to Improve Feature Quality and Stability of Sparse Autoencoders (SAE)