Research map / Label-efficient and robust learning / Recognition and long-tail
Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Matryoshka SAEs, which train multiple nested sparse autoencoder dictionaries of increasing size jointly so that early latents must reconstruct inputs alone, reducing feature absorption and splitting while scaling dictionary size.
Problem
Standard sparse autoencoders trained with larger dictionaries suffer from feature splitting, feature absorption, and feature composition driven by sparsity, leaving high-level concepts missing or warped and degrading downstream usefulness.
Approach
Train several nested SAE dictionaries (prefixes of one latent vector) simultaneously, each required to reconstruct the input using only its own prefix of latents, summing the reconstruction losses so early latents learn general features and later latents specialize.
Key results
- At L0=40 on Gemma-2-2B, Matryoshka SAE absorption rate is 0.05 versus 0.49 for BatchTopK SAE, and it needs about one latent per first letter versus three for BatchTopK.
- Matryoshka SAE explains 70% of variance at L0=40 versus 72% for BatchTopK, a small reconstruction trade-off, while downstream CE loss is comparable at larger L0.
- Meta-SAEs explain 42% of variance in Matryoshka decoder directions versus 55% for BatchTopK at the largest dictionary size, indicating less feature composition.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use nested reconstruction objectives instead of a single flat dictionary objective | SAE on Gemma-2-2B · Gemma-2-2B layer 12, L0=40 | feature absorption rate: 0.49 -> 0.05 | Helped |
| Use nested reconstruction objectives instead of a single flat dictionary objective | SAE on Gemma-2-2B · Gemma-2-2B layer 12, L0=40 | variance explained: 72% -> 70% | Hurt |
| Use nested reconstruction objectives instead of a single flat dictionary objective | SAE on Gemma-2-2B · Gemma-2-2B layer 12, L0=40 | latents per first letter: three -> one | Helped |
| Use nested reconstruction objectives instead of a single flat dictionary objective | SAE on Gemma-2-2B · Gemma-2-2B decoder directions, largest dictionary size | variance explained by meta-SAE: 55% -> 42% | Helped |
| Weight nested reconstruction losses by the number of new latents in each group rather than equally (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | downstream cross-entropy loss | Helped |
| Weight nested reconstruction losses by the number of new latents in each group rather than equally (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | feature absorption and spurious correlation removal | Hurt |
| Stop gradients between nested dictionary reconstructions (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | feature absorption | Helped |
| Stop gradients between nested dictionary reconstructions (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | loss recovered, sparse probing, and spurious correlation removal | Hurt |
| Use three nested dictionaries instead of five (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | loss recovered | Helped |
| Use three nested dictionaries instead of five (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | feature absorption and SCR | Hurt |
| Use ten nested dictionaries instead of five (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 65k dictionary | automated interpretability | Hurt |
| Randomly sample dictionary prefix sizes during training instead of fixing them (ablation) | Matryoshka SAE · Gemma-2-2B layer 12, 16k dictionary | spurious correlation removal and sparse probing | Hurt |
Surprising findings
- Stop-gradient variant improved absorption at low L0 but degraded loss recovered, sparse probing, and SCR, so it was not pursued.
- Matryoshka SAEs slightly underperform other architectures at high L0 on the ChessGPT board-game evaluation.
Limitations
- Slightly reduced reconstruction fidelity (variance explained) relative to BatchTopK SAEs.
- Training time increases with the number of nested dictionaries (about 50% with five).
- Evaluation relies mainly on automated and quantitative metrics rather than human interpretability studies.
Experimental setup
Datasets: The Pile TinyStories SAE Bench evaluation suite synthetic hierarchical toy model data
Benchmarks: SAE Bench sparse probing (35 tasks) spurious correlation removal targeted probe perturbation first-letter feature absorption ChessGPT board evaluations
Metrics: variance explained downstream CE loss recovered feature absorption rate sparse probing accuracy SCR TPP max decoder cosine similarity automated interpretability score
Baselines: BatchTopK SAE ReLU SAE TopK SAE Gated SAE JumpReLU SAE P-Annealing SAE
Builds on: Matryoshka Representation Learning BatchTopK SAE TopK SAE JumpReLU SAE
Closest papers by meaning
- Incorporating Hierarchical Semantics in Sparse Autoencoder Architectures (NeurIPS 2026 · similarity 0.90)
- Projecting Assumptions: The Duality Between Sparse Autoencoders and Concept Geometry (NeurIPS 2025 · similarity 0.90)
- Towards Principled Evaluations of Sparse Autoencoders for Interpretability and Control (ICLR 2025 · similarity 0.90)
- Are Sparse Autoencoders Useful? A Case Study in Sparse Probing (ICML 2025 · similarity 0.90)
- Mechanistic Interpretability with Sparse Autoencoder Neural Operators (NeurIPS 2026 · similarity 0.90)
- Archetypal SAE: Adaptive and Stable Dictionary Learning for Concept Extraction in Large Vision Models (ICML 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.