Title: Attribution Guided Distillation for Matryoshka Sparse Autoencoders
Primary area: interpretability and explainable AI
Keywords: Mechanistic Interpretability, Sparse Autoencoders, Matryoshka Sparse Autoencoders

Abstract: Sparse autoencoders (SAEs) aim to uncover monosemantic, interpretable features from deep networks but they face key limitations such as feature splitting and absorption. Matryoshka SAEs (MSAEs) partially mitigate these issues with a nested dictionary architecture, yet they still suffer from redundancy, hyperparameter sensitivity, and higher reconstruction loss. To address these limitations, we propose a Distilled Matryoshka SAE (DMSAE) method that uses an attribution-guided selection strategy to iteratively identify and preserve high attribution neurons. In our method, these neurons are promoted into a frozen "core" that serves as the innermost prefix shared across all outer prefix reconstructions. When trained at layer 12 of Gemma-2-2B, DMSAE outperforms the MSAE baseline on most SAEBench benchmarks. DMSAE is model agnostic and may extend to other models where distilling a minimal core of features is desirable for interpretability and control.

TL;DR: We address limits of sparse autoencoders for mechanistic interpretability by iteratively distilling high-attribution neurons into a frozen core that is shared across training cycles.