Source-linked AI summary
Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Bart Bussmann, Noa Nabeshima, Adam Karvonen, Neel Nanda
TL;DR
Standard SAE sparsity objectives can distort features as dictionaries grow, undermining high-level concepts and downstream interpretability. Matryoshka SAEs address this by jointly training nested dictionaries whose smaller prefixes reconstruct inputs independently, preserving general features alongside specialized ones. Experiments report reduced absorption and improved interpretability-related performance, with trade-offs in reconstruction and training time.
Problem
As SAE dictionaries grow, sparsity can split or absorb general concepts, causing downstream performance and interpretability to deteriorate.
Method
Matryoshka SAEs jointly train nested dictionaries, requiring each smaller prefix of latents to reconstruct the input independently.
Results
Matryoshka SAEs reduce feature absorption, improve disentanglement, probing, and targeted concept erasure, and often maintain or improve downstream performance as dictionaries grow.
Takeaways & Limitations
The results support using nested training objectives to mitigate SAE scaling pathologies and retain interpretable features at multiple abstraction levels.
Takeaways & Limitations
Matryoshka SAEs slightly reduce reconstruction performance and increase training time; five nested SAEs increased training time by circa 50%.
Abstract
from arXiv · showhide
Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting neural networks by extracting the concepts represented in their activations. However, choosing the size of the SAE dictionary (i.e. number of learned concepts) creates a tension: as dictionary size increases to capture more relevant concepts, sparsity incentivizes features to be split or absorbed into more specific features, leaving high-level features missing or warped. We introduce Matryoshka SAEs, a novel variant that addresses these issues by simultaneously training multiple nested dictionaries of increasing size, forcing the smaller dictionaries to independently reconstruct the inputs without using the larger dictionaries. This organizes features hierarchically - the smaller dictionaries learn general concepts, while the larger dictionaries learn more specific concepts, without incentive to absorb the high-level features. We train Matryoshka SAEs on Gemma-2-2B and TinyStories and find superior performance on sparse probing and targeted concept erasure tasks, more disentangled concept representations, and reduced feature absorption. While there is a minor tradeoff with reconstruction performance, we believe Matryoshka SAEs are a superior alternative for practical tasks, as they enable training arbitrarily large SAEs while retaining interpretable features at different levels of abstraction.
1. Introduction
Standard SAEs face interpretability failures that worsen with larger dictionaries, while Matryoshka SAEs train nested reconstructions to preserve general and specific features. Experiments report improved downstream interpretability outcomes, with a small reconstruction-fidelity trade-off.
- Challenges and approach: Sparsity-driven feature splitting can replace broad concepts with specialized latents, leaving high-level categories missing from the dictionary.For example, punctuation may fragment into separate period, question-mark, and comma features.
- Challenges and approach: Larger SAE dictionaries can lower training loss while downstream performance deteriorates because sparsity exploits added capacity for distortions.These failures create practical challenges for reliable probing, steering, and analysis.
- Experimental findings: Matryoshka SAEs reduce feature absorption and composition while sacrificing a small amount of reconstruction fidelity at average sparsity L0=40.The figure compares these metrics between Matryoshka and standard SAEs.
- Challenges and approach: Matryoshka SAEs simultaneously train nested dictionaries, requiring each smaller dictionary to reconstruct inputs using only its own initial latents.This prevents later specialized latents from absorbing earlier general features and adds only a modest training-time increase.
- Experimental findings: Across synthetic and real-world datasets, Matryoshka SAEs improve disentanglement, probing, and targeted concept erasure while often maintaining or improving downstream performance as dictionaries grow.The reported trade-off is a small increase in reconstruction error.
- Implications: The authors interpret these results as evidence that SAE scaling pathologies can be mitigated through changes to the training procedure.They present this as a positive sign for continued SAE interpretability research.
2. Background
SAEs map neural-network activations into sparse features and reconstruct them, but sparsity objectives can distort hierarchical and independent concepts. These distortions become more severe with larger dictionaries, motivating nested representations such as Matryoshka learning.
- Sparse autoencoders: Sparse autoencoders encode high-dimensional neural activations into nonnegative sparse latent features and decode them into approximate reconstructions.A latent is typically active when its activation is greater than zero.
- Scaling challenges: SAE training balances reconstruction quality against sparsity, but this objective can produce feature splitting, absorption, and composition.These are sparsity-driven distortions rather than merely reconstruction failures.
- Scaling challenges: Feature splitting fragments unified concepts into specialized latents, causing high-level concepts such as punctuation marks to disappear from the dictionary.The general concept may instead be implicit in each specialized decoder direction.
- Scaling challenges: Feature absorption gives general features systematic blind spots where specialized latents respond, making the general feature less reliable and harder to interpret.The paper illustrates this with a female-names feature that develops a hole for “Lily.”
- Scaling challenges: Feature composition encourages single latents for combinations such as “red triangle” instead of separately representing independent color and shape features.This distortion arises when features naturally co-occur.
- Scaling challenges: Larger dictionaries provide more capacity but also more opportunities for sparsity-driven distortions that reconstruction objectives alone do not prevent.Their persistence suggests limitations in standard training objectives rather than optimization alone.
- Matryoshka representation learning: Matryoshka Representation Learning encodes varying levels of granularity within one embedding so representations can adapt to downstream computational constraints.Its loss evaluates performance using only parts of the embedding.
3. Matryoshka Sparse Autoencoders
Matryoshka SAEs train nested dictionaries that must reconstruct inputs at multiple scales, organizing early latents toward general features and later latents toward specific features. The method also uses BatchTopK activation for Gemma-2-2B experiments, with inference replacing batch selection by a calibrated global threshold.
- Nested dictionaries: Matryoshka SAEs simultaneously train nested autoencoders of increasing size, each reconstructing inputs from only the first m_i latents.The nested dictionary sizes satisfy m_1 < m_2 < ... < m_n = m.
- Variants: The method includes random-prefix training, which samples dictionary prefixes per batch, and fixed-prefix training, which sets them before training.These are the two Matryoshka SAE variants evaluated in the method description.
- Training objective: The training objective enforces good reconstruction at multiple scales, requiring each prefix of latents to reconstruct the input independently.The first m_1 latents reconstruct alone, followed by progressively larger prefixes.
- Training objective: Early latents are pressured to capture general features, while later latents can specialize in more specific features because they participate only in larger dictionaries.This creates a hierarchy of representations at different scales.
- BatchTopK activation: BatchTopK retains the B × K largest activations across a batch and sets the others to zero, while inference uses a calibrated global threshold.The threshold is calibrated on training data to maintain the desired average sparsity.
4. Experiments
Experiments across synthetic, TinyStories, and Gemma-2-2B settings show that Matryoshka SAEs preserve hierarchical features while improving disentanglement and downstream concept-related performance. They incur slightly worse reconstruction but comparable downstream language-model loss, and generally maintain or improve performance as dictionary size grows.
- Experimental settings: Experiments span a synthetic toy model, a 4-layer TinyStories model, and Gemma-2-2B benchmarks.The settings increase in complexity from controlled feature absorption to larger-scale validation.
- Toy model: Matryoshka SAEs recover parent features alongside active children, whereas Vanilla SAEs absorb parent features when child features activate.The toy-model activation comparison directly tests whether hierarchical structure is preserved.
- Toy model: Matryoshka decoder vectors show minimal overlap with ground-truth features, while Vanilla SAEs show redundant parent-child representations.The Matryoshka similarity matrix has a clear diagonal structure, unlike the Vanilla matrix.
- TinyStories: A 25k-latent Matryoshka SAE preserves broad female-token coverage while retaining specialized Lily and Sue latents without holes.The larger reference SAE develops holes when specialized latents absorb cases from the general female-words feature.
- Gemma-2-2B: At L0 = 40, Matryoshka SAEs explain 70% of activation variance versus 72% for BatchTopK, but achieve comparable downstream cross-entropy loss at larger L0s.Nested reconstruction constraints account for the mild reconstruction gap while downstream loss remains comparable.
- Gemma-2-2B: At L0 = 40, Matryoshka SAEs reach 0.05 absorption versus 0.49 for BatchTopK and use one latent per first letter versus three.They also outperform benchmark architectures on TPP and SCR, perform best on sparse probing at lower sparsities, and show lower decoder-vector cosine similarity.
- Gemma-2-2B: Matryoshka SAEs improve or maintain performance across increasing dictionary sizes, while alternative architectures often degrade.The reported scale range includes dictionary sizes of 4k, 16k, and 65k.
5. Limitations
The evaluation identifies three limitations: slightly reduced reconstruction performance, increased training time, and reliance on quantitative and automated interpretability measures.
- Slightly reduced reconstruction performance may be problematic for applications requiring precise activation reconstruction.Downstream loss remains competitive despite this trade-off.
- Five nested SAEs increased training time by circa 50% in the experiments.The increase depends on the number of nested dictionary sizes, M.
- The evaluation relies heavily on quantitative feature-quality measures and automated interpretability metrics.These measures may not fully capture human-relevant interpretability or practical utility for manual downstream analysis.
6. Discussion and Conclusion
The paper presents Matryoshka SAEs as a way to address scaling-related SAE pathologies through hierarchical feature learning. It reports lower feature absorption and improved concept isolation, while suggesting that scalable interpretability methods are increasingly important for understanding larger neural networks.
- Matryoshka SAEs enforce hierarchical feature learning through nested dictionaries.The approach is presented as a modification to the SAE training process.
- Matryoshka SAEs achieve lower feature absorption rates and improved concept isolation compared to standard approaches.
- The reported results suggest that SAE scaling pathologies can be overcome through changes to training objectives.The paper characterizes these pathologies as artifacts of myopic training objectives rather than fundamental limitations.
- Scalable interpretability approaches are described as essential for understanding increasingly large and capable neural networks.
Impact Statement
The work advances neural-network interpretability, which could support failure-mode identification, behavior verification, and safer AI development. However, improved interpretability may also accelerate language-model development and deployment, with potentially mixed societal effects.
- Better understanding of model internals could help identify failure modes, verify intended behaviors, and guide more reliable and safe AI systems.
- Interpretability advances may accelerate large-language-model development and deployment, potentially amplifying both positive and negative societal impacts.
A. SAE Training Details
The training details describe baseline SAE comparisons, comparable training conditions, and an example Matryoshka implementation. The implementation computes sparse features, incrementally reconstructs inputs across dictionary-size cutoffs, and sums reconstruction losses.
- The SAE Bench baseline suite contains six proposed SAE variants trained on layer 12 of Gemma-2-2B and layer 8 of Pythia-160M.
- Matryoshka SAEs were trained with identical data ordering and hyperparameters for direct comparison.
- The example implementation computes ReLU sparse features using batch_top_k.The comment identifies these as the same features used by a standard BatchTopK SAE.
- Reconstructions are incrementally accumulated from feature groups corresponding to the dictionary-size cutoffs.
- Mean-squared reconstruction losses are stored for each accumulated output, while the function returns summed loss and final reconstruction.
C.1. Toy Model Training Implementation
The toy-model Matryoshka SAE uses nested prefix reconstructions to organize latent activations across dictionary sizes. Compared with matched BatchTopK SAEs, it trades some input-variance reconstruction for a smaller downstream cross-entropy gap as more sub-SAEs are used.
- Model and training configuration: The Matryoshka training loss sums MSE and ℓ1 losses across sampled prefixes, with the full SAE always included.Training uses 10 sampled prefixes per batch and truncated Pareto(0.5) prefix lengths.
- Latent activation patterns: Early Matryoshka latents activate more frequently because they participate in multiple reconstruction objectives.The activation-rate distribution is shown across different nested dictionary groups.
- Sub-SAE reconstruction and downstream performance: BatchTopK SAEs explain more input variance, but the downstream language-model cross-entropy degradation gap narrows as more Matryoshka sub-SAEs are used.The comparison uses matched nested dictionary sizes of 2304, 4608, 9216, 18432, and 36864 latents, with corresponding average L0 norms of 22, 25, 27, 29, and 32.
D.1. Investigating Composition with Meta-SAEs
Meta-SAEs explain less variance in Matryoshka decoder directions than in BatchTopK directions, indicating less shared information among Matryoshka latents. This disentanglement advantage persists as dictionary size increases.
- Meta-SAE setup: Meta-SAEs use one-quarter of the input SAE dictionary size and four active latents, with explained variance serving as a proxy for shared information.They are trained on the decoder matrices of both architectures.
- Results: At the largest dictionary size, Meta-SAEs explain 55% of variance in BatchTopK decoder directions versus 42% in Matryoshka directions.The comparison is shown in Figure 14.
- Results: Matryoshka latents are more disentangled, while BatchTopK develops increasing shared structure with dictionary size and Matryoshka maintains relatively constant disentanglement.The interpretation follows the decoder-direction variance comparison.
E. Evaluations with Board Game Models
Board-game evaluations compare SAE architectures using board reconstruction and feature coverage. Matryoshka performs comparably in the low-L0 regime but shows degradation at higher L0 values in the small ChessGPT model.
- Evaluation metrics: The board-game evaluation measures board reconstruction and feature coverage.Board reconstruction assesses recovery of the game board, while feature coverage measures alignment with predefined board-game features.
- Results: At L0 ≤150, Matryoshka performs comparably to ReLU, P-Anneal, Gated, and TopK architectures on both metrics.Peak scores across architectures are fairly comparable and generally occur in this lower-L0 range.
- Results: Matryoshka performance degrades at higher L0 values, partly because L0 ≥150 occupies a substantial fraction of ChessGPT’s hidden dimension of 512.This setting differs from Gemma-2B, where the largest L0 values represent a much smaller fraction of the hidden dimension.
- Scope: The study emphasizes Gemma-2-2B for its main analysis because the authors prioritize large-scale language models that better reflect real-world applications.ChessGPT results are included for completeness and comparison with prior work.
- Hierarchical visualization: The hierarchical visualization modifies masked cosine similarity into a directed, activation-scaled measure and recursively links latents across wider SAE models.A threshold of 0.6 is used to identify parent-child relationships.
- Hierarchical visualization: The resulting latent structure can be a DAG rather than a strict tree because deeper latents may exceed the similarity threshold for multiple parents.Visualization assigns each multi-parent latent to its highest-similarity parent, obscuring some underlying complexity.
G.1. Ablation: Loss Weighting
Ablations show that equal loss weighting preserves Matryoshka’s feature-quality advantages, while changing gradient flow, group count, or dictionary-size sampling produces trade-offs. Five nested groups provide a reasonable overall balance in the reported metrics.
- Loss weighting: Weighted Matryoshka assigns loss weights proportional to the number of new latents, emphasizing full-dictionary reconstruction while retaining hierarchical pressure.For [25%, 50%, 100%] nested dictionaries, the weights are [0.25, 0.25, 0.5] instead of [0.33, 0.33, 0.33].
- Loss weighting: With 200M training tokens, Weighted Matryoshka matches BatchTopK downstream cross-entropy more closely than equally weighted Matryoshka.The weighted variant performs worse on feature absorption and spurious correlation removal, indicating a feature-quality trade-off.
- Gradient flow: The stop-gradient variant improves feature absorption at low L0 but degrades loss recovered, sparse probing, and spurious correlation removal.The authors did not pursue this variant further.
- Number of nested dictionaries: The three-group variant improves loss recovered but worsens feature absorption and SCR, while the ten-group variant slightly improves feature absorption but performs significantly worse on loss recovered and automated interpretability.Five groups provide a reasonable balance across these metrics.
- Dictionary-size sampling: Randomly sampled dictionary sizes perform similarly to fixed sizes, with minor degradation in spurious correlation removal and sparse probing.The authors use fixed dictionary sizes for the main experiment because random sampling may add distributed-training complexity.