Source-linked AI summary
Archetypal SAE: Adaptive and Stable Dictionary Learning for Concept Extraction in Large Vision Models
Thomas Fel, Ekdeep Singh Lubana, Jacob S. Prince, Matthew Kowal, Victor Boutin, Isabel Papadimitriou, Binxu Wang, Martin Wattenberg, Demba Ba, Talia Konkle
TL;DR
Existing SAEs can produce sharply different dictionaries across otherwise similar training runs, limiting their reliability for interpretability. The paper introduces convex-hull-constrained Archetypal SAEs and relaxed variants, finding improved stability and more structured, meaningful representations while retaining reconstruction quality.
Problem
Existing SAEs exhibit severe instability, with similar training runs producing divergent concept dictionaries that undermine their reliability as an interpretability protocol.
Method
A-SAE constrains dictionary atoms to the convex hull of data, while RA-SAE adds a small relaxation to preserve flexibility; the paper also introduces dictionary-quality metrics and benchmarks.
Results
Across evaluations, archetypal SAEs improve stability and representation structure, while RA-SAEs balance reconstruction quality with meaningful concept discovery.
Takeaways & Limitations
Archetypal constraints provide a more reliable basis for concept discovery in large-scale vision models and may extend to other structured modalities.
Takeaways & Limitations
Directly optimizing archetypal representations over all tokens is often infeasible at very large scale, motivating reduced-point approximations and relaxation.
Abstract
from arXiv · showhide
Sparse Autoencoders (SAEs) have emerged as a powerful framework for machine learning interpretability, enabling the unsupervised decomposition of model representations into a dictionary of abstract, human-interpretable concepts. However, we reveal a fundamental limitation: existing SAEs exhibit severe instability, as identical models trained on similar datasets can produce sharply different dictionaries, undermining their reliability as an interpretability tool. To address this issue, we draw inspiration from the Archetypal Analysis framework introduced by Cutler & Breiman (1994) and present Archetypal SAEs (A-SAE), wherein dictionary atoms are constrained to the convex hull of data. This geometric anchoring significantly enhances the stability of inferred dictionaries, and their mildly relaxed variants RA-SAEs further match state-of-the-art reconstruction abilities. To rigorously assess dictionary quality learned by SAEs, we introduce two new benchmarks that test (i) plausibility, if dictionaries recover "true" classification directions and (ii) identifiability, if dictionaries disentangle synthetic concept mixtures. Across all evaluations, RA-SAEs consistently yield more structured representations while uncovering novel, semantically meaningful concepts in large-scale vision models.
1. Introduction
The paper frames concept extraction as an interpretability approach for understanding internal visual concepts, then identifies instability in current SAEs and proposes archetypal variants to address it.
- Concept-based methods link neural-network representations to human-interpretable visual concepts.
- Concept extraction methods can be understood as instances of dictionary learning.
- Current SAE training is unstable, motivating the paper’s focus on more reliable concept dictionaries.
- A-SAE constrains dictionary atoms to the convex hull of training data, while RA-SAE adds a small relaxation for greater modeling flexibility.The geometric constraint is intended to improve consistency across training runs.
- The paper evaluates dictionary quality through metrics and benchmarks targeting classification-direction recovery and disentanglement of synthetic image mixtures.
2. Related Work
The related work situates SAEs within dictionary learning and explainable vision, emphasizing their scalability while noting unresolved interpretability challenges.
- Sparse Coding & Dictionary Learning: Dictionary learning seeks sparse representations that reconstruct inputs using a small subset of learned dictionary atoms.
- Sparse Coding & Dictionary Learning: Classical dictionary-learning approaches include clustering, matrix factorization, and related sparse-coding methods.
- Vision Explainability: Attribution methods identify input regions influencing predictions but do not by themselves provide the concepts used by the network.
- Vision Explainability: Concept-based interpretability instead seeks semantically meaningful directions connected to latent activations.
- Vision Explainability: SAEs are presented as a highly scalable form of dictionary learning, while prior work reports overly specific features and compositionality challenges.
3. (In)Stability of SAEs
This section formalizes concept extraction and compares dictionary-learning approaches, showing that SAEs offer strong reconstruction fidelity but suffer from instability across runs and perturbations.
- Concept Extraction as Dictionary Learning: Concept extraction approximates activation data with sparse linear combinations of learned dictionary atoms and corresponding codes.
- Concept Extraction as Dictionary Learning: SAEs, NMF, Semi-NMF, and related methods can be expressed within a common dictionary-learning framework.
- Concept Extraction as Dictionary Learning: SAEs are especially suited to large-scale applications because they use backpropagation and inherent batch-learning capabilities.
- Concept Extraction as Dictionary Learning: The main comparison balances sparsity, measured by the ℓ0 norm of Z, against fidelity, measured by ℓ2 reconstruction error.
- 3. (In)Stability of SAEs: SAE methods outperform optimization-based dictionary-learning methods in reconstruction fidelity at fixed sparsity levels.
- 3. (In)Stability of SAEs: The stability metric measures optimal average cosine similarity after best dictionary alignment via the Hungarian algorithm.
- 3. (In)Stability of SAEs: Cosine stability around 0.5 was observed for TopK SAE on DinoV2 after changing only the random seed across four runs.Dataset perturbations of 5% or 10% produced a similar trend.
4. Towards Archetypal SAEs
Archetypal SAEs constrain dictionary atoms to convex combinations of data points, anchoring representations within the data geometry. A relaxed variant preserves this anchoring while allowing controlled deviations, improving scalability and reconstruction flexibility.
- Archetypal SAE formulation: A-SAE constructs dictionary atoms as convex combinations of data points, making each atom originate from the observed representations.The dictionary is parameterized by a row-stochastic matrix multiplying the data matrix.
- Geometric interpretation: Constraining atoms to the data convex hull prevents out-of-sample directions and empirically improves dictionary stability.With nonnegative codes, reconstructions remain within the conic hull of the data.
- Scalability: Scaling the unconstrained archetypal formulation is infeasible for very large token sets, motivating reduced representative sets and relaxation.The method uses a much smaller centroid set C selected with K-Means and optimizes only the combination weights for A-SAE.
- Relaxed Archetypal SAE: RA-SAE learns convex-combination weights together with a small-norm deviation term, allowing controlled movement beyond the convex hull.The deviation is constrained by ||Λ||_2 ≤ δ, while the weights remain constrained.
- Empirical tradeoff: RA-SAE achieves reconstruction performance on par with conventional Top-K SAEs while retaining the stability benefits of archetypal constraints.The relaxation parameter enables exploration beyond conv(C) without abandoning the geometric anchoring.
5. Experiments
The experiments evaluate dictionary-learning methods across reconstruction, consistency, dictionary and code structure, plausibility, identifiability, and qualitative concept discovery. Archetypal methods trade some reconstruction advantage for greater consistency and structure, while RA-SAE improves alignment with classification directions and yields organized, semantically specific concepts.
- Experimental Setup: The experiments compare five vision backbones and dictionary-learning methods using overcomplete representations trained on ImageNet activations.The evaluation covers DINOv2, SigLip, ViT, ConvNeXt, and ResNet50, with dictionaries generally sized at five times the feature dimension.
- Dictionary Learning Metrics: The evaluation measures sparse reconstruction, consistency, dictionary structure, and code structure to characterize learned solutions beyond reconstruction quality.Consistency includes stability across runs and proximity of atoms to real data, while structure metrics include rank, coherence, connectivity, and negative interference.
- Dictionary Learning Metrics: Archetypal SAEs match existing SAEs in reconstruction and outperform NMF methods in R2 at comparable sparsity levels.SAE methods achieve the strongest reconstruction overall, whereas C-NMF, S-NMF, and Archetypal methods excel in consistency and structural measures.
- Plausibility Benchmark: RA-SAE achieves significantly higher plausibility than the TopK SAE baseline by aligning learned concepts more closely with classification-head directions.The plausibility benchmark measures whether dictionary concepts recover directions represented by the classifier’s final-layer weights.
- Soft Identifiability Benchmark: The Soft Identifiability Benchmark tests whether an SAE assigns distinct concepts to objects generated from synthetic mixtures with known ground-truth structure.The benchmark uses twelve datasets of collaged objects and scores whether a concept activates above a threshold for each object class.
- Qualitative Examples: Qualitative analyses show RA-SAE discovering structured fine-grained and context-dependent concepts, including rabbit-part distinctions, shadows, barber identity, and flower-petal edges.RA-SAE also forms clusters of semantically related features and separates rabbit ears, body, face, and paws more clearly than TopK counterparts.
6. Conclusion
The paper identifies instability in Sparse Autoencoders and introduces archetypal constraints to improve dictionary reliability. Its benchmarks and qualitative analyses indicate that these constraints support more structured, plausible, and semantically meaningful representations.
- Conclusion: Identical SAE training runs can produce divergent concept dictionaries, limiting their reliability for interpretability.The paper frames this instability as a fundamental limitation of current Sparse Autoencoders.
- Conclusion: A-SAE constrains dictionary atoms to the data’s convex hull, while RA-SAE relaxes that constraint to balance stability with reconstruction quality.The relaxed variant is reported to support meaningful concept discovery in large-scale vision models.
- Conclusion: Novel metrics and identifiability-inspired benchmarks provide a systematic framework for assessing dictionary quality and disentanglement.The paper reports that archetypal constraints improve the structure and plausibility of learned representations.
Impact Statement
Archetypal SAEs address instability in sparse interpretability by producing more consistent, data-grounded dictionaries and structured concepts in large vision models.
- Archetypal SAEs yield dictionaries that are more consistent across runs and better aligned with true classification and generative directions.
- These benchmarks and representations support more reliable and scientifically useful sparse concept discovery.
- RA-SAE uncovers shadow-based, context-dependent barber, and fine-grained petal-edge concepts in DinoV2.
- The DinoV2 Archetypal SAE has low dictionary auto-correlation, a long-tail firing-rate distribution, and uniformly distributed reconstruction error except at CLS.
B. Extended Related Work on Archetypal Analysis
Archetypal SAEs build on archetypal analysis while adapting its geometric constraints and optimization strategy to modern sparse autoencoders.
- Archetypal analysis represents observations as convex combinations of extremal archetypes located on the boundary of the data convex hull.
- Prior work extended archetypal analysis through kernel, multilayer, supervised, scalable, and deep-learning variants.
- The proposed method constrains only decoder atoms to the data convex hull while retaining a linear encoder and end-to-end SAE training.
- Unlike archetypal analysis, the method targets SAE stability and semantic consistency rather than approximation or generation.
C. Formal Definitions of Metrics
The paper evaluates sparse autoencoders through metrics covering reconstruction, consistency, dictionary structure, code structure, and scalable archetypal distillation.
- C. Formal Definitions of Metrics: The evaluation spans sparse reconstruction, consistency, dictionary structure, and code structure.
- C.1. Sparse Reconstruction: Reconstruction Error R2 measures how well the learned dictionary explains input activations relative to a single-bias predictor, while Dead Codes measure unused atoms.
- C.2. Consistency: Stability compares dictionaries across independently initialized runs, with 1 indicating perfect alignment and 0 indicating seed-specific concepts.
- C.2. Consistency: Max Cosine Similarity gives an upper-bound alignment but does not enforce global dictionary consistency.
- C.2. Consistency: OOD Score measures dictionary atoms’ deviation from real activations, with 0 meaning every atom exactly matches a data point.
- C.3. Structure in the Dictionary (D): Stable Rank and Effective Rank assess dictionary dimensionality and singular-value concentration, while coherence measures redundancy between atoms.
- C.4. (iv) Structure in the Codes (Z): Connectivity measures concept co-activation diversity, whereas Negative Interference counts destructive interactions between co-activated, negatively correlated atoms.
- D. Distilling A into C: K-Means is the most effective tested distillation method for reducing the activation matrix to a compact set while producing convex-hull points.
E. Implementation Details
The implementation constructs relaxed archetypal dictionaries from distilled centroids and uses convex-hull geometry to control stability, rank, and data grounding.
- RA-SAE forms dictionary atoms as convex combinations of distilled centroids plus controlled additive perturbations.
- Standard SAE gradient updates can drift dictionaries away from the data because code correlations and seed-dependent terms influence the update.
- Archetypal parameterization guarantees dictionary atoms lie in the data convex hull and reconstructions lie in its conic hull.
- Using a distilled set C remains expressive when C contains the extreme points of the original data hull.
- Row-stochastic weights limit how much dictionary changes can grow under small data perturbations, providing a geometric stability control.
- The dictionary rank cannot exceed the data rank, preventing arbitrarily complex bases and favoring structured or low-rank representations.
F.4. Bounding OOD score with Archetypal Constraints
Under non-interfering archetype assumptions, the OOD measure of archetypal dictionary atoms is bounded by their mixture weights and is minimized when an atom aligns with a data point. The section also reports that Silverman-kernel JumpReLU gave the most stable and accurate reconstruction results.
- OOD bound: The OOD measure is tied directly to row-stochastic dictionary weights, reaching its maximum of 1 when an atom perfectly aligns with a data point.This construction keeps dictionary atoms grounded in the data.
- Assumptions: Under non-interfering archetypes, active data points contributing to each dictionary atom are orthogonal, preventing their contributions from interfering.The assumption permits non-orthogonal archetypes in the dictionary bank but not simultaneous activation.
- Special case: OOD(D_i) = 0 when some weight W_ij = 1, meaning the dictionary atom aligns perfectly with a data point.The sparsity of W helps limit interference and supports the derived bounds.
- JumpReLU kernel choice: The Silverman kernel produced the most stable and accurate JumpReLU reconstruction results and was selected with bandwidth 10^-2.Kernel choice appeared to have only a modest impact on performance.
H. Soft Identifiability Benchmark
The Soft Identifiability Benchmark tests whether SAEs recover distinct known concepts from synthetic image mixtures. It uses controlled datasets, multiple pretrained vision models, and an adaptive activation threshold to measure concept-class assignment accuracy.
- Benchmark purpose: The benchmark assesses whether SAEs recover distinct concepts from synthetic image mixtures whose underlying generative factors are known.This directly tests identifiability against known ground-truth factors.
- Synthetic datasets: Twelve synthetic datasets contain 4,000 collaged images each, with four objects per image and 9–20 unique objects per dataset.Dictionary size equals the number of true generative factors.
- Evaluation setup: Each dataset is split into 2,000 training and 2,000 test images, then processed through DinoV2, ResNet50, SigLIP, or ViT before SAE training.Pooled activations from the pretrained vision models serve as SAE input representations.
- Metric: Accuracy measures whether each object class is assigned a distinct SAE dictionary concept, using an activation threshold selected from empirical concept-activation percentiles.Thresholds range from the 1st to the 100th percentile to adapt to activation distributions.
H.1. Complete Results
The complete benchmark results are reported as accuracy scores for various methods across models and classes. The section points readers to a full table and additional per-dataset analysis.
- Complete results: The accompanying analysis includes per-dataset accuracy scores and additional breakdowns.These results supplement the aggregate table.
- Complete results: Table 4 reports accuracy scores for various methods across models and classes.The table is presented as the complete set of benchmark results.