Attribution-Guided Distillation of Matryoshka Sparse Autoencoders

TL;DR

This work tackles interpretability and transferability challenges in sparse autoencoders by introducing Distilled Matryoshka Sparse Autoencoders (DMSAEs), an attribution-guided distillation framework that iteratively trains Matryoshka SAEs and selects a compact core of high-value latent features. By transferring and freezing only the core encoder directions across cycles and scoring latents via gradient × activation with respect to the next-token loss, DMSAEs produce a stable, reusable core (e.g., 197 latents) that improves SAEBench metrics on Gemma-2-2B activations. The method demonstrates that consistent latent features can be transferred across sparsity levels, offering a practical route toward reproducible, monosemantic feature discovery, albeit at higher computational cost due to multi-cycle distillation. Future work aims to reduce distillation overhead and extend the approach to other architectures and domains, enabling more robust interpretability pipelines for large-scale models.

Abstract

Sparse autoencoders (SAEs) aim to disentangle model activations into monosemantic, human-interpretable features. In practice, learned features are often redundant and vary across training runs and sparsity levels, which makes interpretations difficult to transfer and reuse. We introduce Distilled Matryoshka Sparse Autoencoders (DMSAEs), a training pipeline that distills a compact core of consistently useful features and reuses it to train new SAEs. DMSAEs run an iterative distillation cycle: train a Matryoshka SAE with a shared core, use gradient X activation to measure each feature's contribution to next-token loss in the most nested reconstruction, and keep only the smallest subset that explains a fixed fraction of the attribution. Only the core encoder weight vectors are transferred across cycles; the core decoder and all non-core latents are reinitialized each time. On Gemma-2-2B layer 12 residual stream activations, seven cycles of distillation (500M tokens, 65k width) yielded a distilled core of 197 features that were repeatedly selected. Training using this distilled core improves several SAEBench metrics and demonstrates that consistent sets of latent features can be transferred across sparsity levels

Figures (7)

• Figure 1: DMSAE distillation schematic. Each cycle trains a Matryoshka SAE where the smallest reconstruction uses core + prefix-0 and larger reconstructions use progressively larger non-core prefixes. After training, we score latents in core + prefix-0 using a gradient $\times$ activation score for next-token loss, and then choose the smallest set that accounts for a target fraction ($\tau$) of the total attribution. This set becomes the core for the next cycle. Across cycles, we copy and freeze only the core encoder weight vectors, while core encoder biases, core decoder weights and biases, and all non-core latents are randomly reinitialized.
• Figure 2: Evolution of core latents across distillation cycles (target sparsity $k{=}320$, selection threshold $\tau{=}0.9$). After each training cycle $t$ (including $t=0$, the SAEBench checkpoint), we rank latent features by their contribution to the language model’s next-token loss and select the smallest set $C^{(t)}$ whose cumulative score reaches 90% of the total attribution. The y-axis shows the number of latents selected into the core at the end of cycle $t$. Colors indicate when a feature was first selected (the earliest cycle in which it appears in any $C^{(t)}$), showing how later cycles combine previously discovered latents with newly identified ones. For benchmarking, we defined the distilled core as the intersection of the final two cycle cores, yielding 197 latents.
• Figure 3: Distillation uncovers useful latents. We trained two DMSAEs ($k=320$), one that uses a core derived from our distillation procedure (blue; Fig. \ref{['fig:core_latent_origins_0-1-2-3-4-5-6-7']}, cycle 7) and another that uses a randomly initialized core of equivalent size (red). The y-axis shows the mean number of core latents active per sample during training ($L_{0,\mathrm{core}}$). The model with a randomly initialized core drives $L_{0,\mathrm{core}}$ towards zero during training, while the model with the distilled core continues to rely on the core throughout training. This suggests that distillation identifies core latents that are systematically useful for reconstruction.
• Figure 4: SAEBench performance across sparsities (dense core DMSAE). SAEBench metrics versus measured sparsity (mean number of active latents) for Distilled MSAE (DMSAE) with a dense core, compared with Matryoshka SAE (MSAE) and BatchTopK, evaluated on Gemma-2-2B layer 12 activations (all SAEs have a dictionary size of 65k). Here, the DMSAE core (197 latents) is taken from the distilled core produced by the $k{=}320$ distillation run (\ref{['fig:core_latent_origins_0-1-2-3-4-5-6-7']}). We then used this core to train new DMSAEs from scratch at each target sparsity ($k\in\{20,40,80,160,320,640\}$). Core encoder weights are frozen while sparsity is enforced on the remaining non-core latents. DMSAE matches or improves several metrics, particularly absorption and RAVEL, but AutoInterp drops sharply at the lowest sparsities. For a more direct comparison to the SAEBench baselines, Appendix E reports a sparsity-matched setting in which BatchTopK is applied globally (including the core); these runs exhibit similar overall trends (Fig. \ref{['fig:saebench_core_sparsity_ablation']}).
• Figure 5: Choosing the attribution coverage threshold for core promotion. We trained DMSAEs ($k{=}320$) using different attribution coverage thresholds ($\tau$). Promoting more latents steadily reduces reconstruction loss (top) and increases fraction of variance explained (bottom) throughout training, with no clear saturation up to $\tau{=}1.0$. The $\tau{=}0.9$ setting nearly overlaps $\tau{=}1.0$ in fraction of variance explained and achieves similar reconstruction loss. These results justify using $\tau{=}0.9$ as our default attribution threshold.