---
title: Distilled Matryoshka Sparse Autoencoders
url: https://www.emergentmind.com/topics/distilled-matryoshka-sparse-autoencoders-dmsaes
type: topic
---

# Distilled Matryoshka Sparse Autoencoders

Distilled Matryoshka Sparse Autoencoders (DMSAEs) constitute a training pipeline for sparse autoencoders that extracts a compact, transferable core of robust, human-interpretable features. By iteratively distilling and freezing the directions most consistently useful for a base model’s next-token loss, DMSAEs address the instability and redundancy of standard sparse feature learning. The methodology revolves around an attribution-guided selection process, transferring only the distilled core encoder weight vectors across cycles, while reinitializing the decoder and non-core latents. Empirical evaluation on the Gemma-2-2B model demonstrates that DMSAEs yield a strongly reusable core—comprising 197 stabilized features after seven cycles—which improves consistency, interpretability, and downstream SAEBench metrics relative to conventional Matryoshka Sparse Autoencoders [2512.24975].

## 1. Matryoshka Sparse Autoencoders: Hierarchical Feature Learning

A typical sparse autoencoder represents an overcomplete dictionary through an encoder–decoder architecture, subject to a sparsity constraint:
\[
f(x) = \mathrm{ReLU}(W_{\text{enc}} x + b_{\text{enc}}) \in \mathbb{R}^m,\quad
\hat{x} = W_{\text{dec}} f(x) + b_{\text{dec}}
\]
where $f(x)$ enforces sparsity (commonly through TopK thresholding).

Matryoshka Sparse Autoencoders (MSAEs) extend this framework by introducing a hierarchy of prefix sizes $M = \{ m_1 < m_2 < ... < m_L \}$, with reconstruction objectives over all prefixes:
\[
\mathcal{L}_{\text{MSAE}}(x) = \sum_{m\in M} \| x - \hat{x}_m \|_2^2 + \alpha \mathcal{L}_{\text{aux}}
\]
where $\hat{x}_m$ reconstructs with only the first $m$ latents. Early latents must encode high-frequency, generalizable content, while later ones specialize. In vanilla MSAEs, no explicit “core” features are preserved across runs, resulting in significant variability and challenging feature reuse [2512.24975].

## 2. Attribution Metric: Gradient × Activation

DMSAEs introduce an attribution-driven basis selection methodology for identifying high-value features. For a given token position $u$, let $x_u \in \mathbb{R}^d$ denote the residual stream activation, and $g_u = \partial \mathcal{L}_{\text{NT}} / \partial x_u$ the gradient of next-token loss. Encoding and masking to the smallest prefix, for each latent $j$:
- Compute activation $a_{u, j}$
- Normalize the decoder vector $\bar{w}^{\text{dec}}_j = w^{\text{dec}}_j / \| w^{\text{dec}}_j \|_2$
- Compute $s_{u, j} = g_u^\top \bar{w}^{\text{dec}}_j$
- Attribution score: $GxA_{u,j} = | a_{u,j} \cdot s_{u,j} |$

Due to the heavy-tailed attribution distribution, DMSAEs aggregate $GxA_{u, j}$ over positions via a high quantile (e.g., $q = 0.99$):
\[
A_j = \text{Quantile}_u ( GxA_{u, j}; q )
\]
Latents are sorted by $A_j$; the smallest subset whose cumulative attribution exceeds a threshold $\tau$ is retained as the distilled core:
\[
C = \operatorname*{argmin}_{S \subseteq P} |S| \quad \text{s.t.} \quad \sum_{j \in S} A_j \geq \tau \sum_{j \in P} A_j
\]

## 3. Iterative Distillation and Transfer Protocol

The DMSAE pipeline operates as a multi-cycle, iterative distillation process. The key steps are:
- **Initialization:** Score the released SAEBench model to select the initial core $C^{(0)}$.
- **Per-Cycle Training:**
  - Freeze encoder rows corresponding to $C^{(t-1)}$; reinitialize all other parameters.
  - Train a two-group MSAE, letting the core latents remain dense while constraining non-core sparsity.
  - Upon convergence, compute quantile-based GxA attributions for core and prefix-0 latents.
  - Select the smallest core $C^{(t)}$ achieving $\tau$-coverage as above.
- **Core Stabilization:** After $T$ cycles, the final distilled core is $C^* = C^{(T)} \cap C^{(T-1)}$, containing only latents persisting through the last two cycles.

This procedure is formalized in the DMSAE high-level pseudocode provided in [2512.24975].

## 4. Distilled Core Convergence and Empirical Evaluation

Applied to Gemma-2-2B, DMSAEs were trained at layer 12 residual streams with dictionary size $K = 65{,}000$ latents, $T = 7$ distillation cycles, and a non-core sparsity of $k = 320$ on 500M tokens. Prefix sizes $M$ and coverage threshold $\tau = 0.9$ were used throughout [2512.24975].

Across cycles, the selected core stabilized between 200–400 latents. The intersection of the last two cycles, $C^{(6)} \cap C^{(7)}$, yielded a distilled core of 197 features persisting across restarts. Empirical evidence indicates that a randomly chosen core of equal size becomes inactive (core $\ell_0 \rightarrow 0$), whereas the distilled core remains active and contributes significantly to loss reduction throughout training. This demonstrates that the distilled directions are systematically high-value.

## 5. Performance on SAEBench and Downstream Metrics

DMSAEs were benchmarked by transferring the distilled core $C^*$ to new SAEs at multiple sparsity regimes ($k \in \{20, 40, 80, 160, 320, 640\}$), freezing only $C^*$ encoder rows and reinitializing all else. SAEBench evaluations included:
- Reconstruction loss
- Fraction of variance explained
- Feature absorption
- RAVEL
- Targeted concept removal
- Spurious correlation removal
- AutoInterp

Results show that across sparsity levels, DMSAEs match or exceed the vanilla MSAE baseline on reconstruction, absorption, and RAVEL, maintaining stable downstream metrics except for a decrease in AutoInterp at the lowest $k$. A sparse-core ablation (global TopK mask optionally including the core) reveals qualitatively similar trends [2512.24975].

## 6. Implications: Interpretability, Transfer, and Model Compression

Distilled Matryoshka Sparse Autoencoders yield several key advantages:
- **Feature Transferability:** Freezing only the most consistently useful encoder directions enables reliable reuse of core features across sparsity budgets, restarts, and tasks.
- **Interpretability:** The distilled core produces a compact, monosemantic, and non-redundant basis, stabilizing feature semantics and simplifying manual or automated feature annotation. It also reduces feature splitting and absorption.
- **Compression:** Only the $|C^*| \approx 197$ encoder weight rows require preservation for transfer, facilitating a two-stage compression approach: a stable core “backbone” and a lightweight, re-trainable residual dictionary.

A plausible implication is that DMSAEs enable more modular and interpretable model analysis pipelines by decoupling the stable, attribution-maximizing core from non-core features adaptively specialized for different downstream requirements [2512.24975].

Source: https://www.emergentmind.com/topics/distilled-matryoshka-sparse-autoencoders-dmsaes