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

Evaluating the Interpretability of Sparse Autoencoders with Concept Annotations
Jun 23, 2026Jonas Klotz, Cassio F. Dantas, Pallavi Jain +2VLM Interpretability Sparse Autoencoders
The Berlin Institute for the Foundations of Learning and Data (BIFOLD) · Technische Universität Berlin, Germany · INRAE +2
Sparse Autoencoders Encode Both Concepts and Functions: The Downstream Geometry of Feature Effects
Jul 27, 2026Phu Gia Hoang, Anwoy Chatterjee, Tanmoy Chakraborty +2Sparse Autoencoders Mechanistic Interpretability
UKP Lab, Technical University of Darmstadt · Department of Electrical Engineering Indian Institute of Technology Delhi, India
Aligned Training: A Parameter-Free Method to Improve Feature Quality and Stability of Sparse Autoencoders (SAE)
May 18, 2026Michał Brzozowski, Neo Christopher ChungSparse Autoencoders Mechanistic Interpretability
Samsung AI Center, Warsaw, Poland · University of Warsaw, Poland

Jun 23, 2026J/K move · Enter open · S save

Evaluating the Interpretability of Sparse Autoencoders with Concept Annotations

Jonas Klotz, Cassio F. Dantas, Pallavi Jain, Diego Marcos+1

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

Open paper AskPDF arXiv

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.

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)