Learning Multi-Level Features with Matryoshka Sparse Autoencoders

Thank you for your positive review of our paper and thoughtful feedback! We are grateful for your recognition that our work is well-written with clear language, presents relevant and extensive experiments, and effectively addresses known limitations of SAEs.

Regarding different model scales and architectures: While computational constraints limited testing on a wider range of the largest models, we have reasons to believe the benefits of Matryoshka SAEs will generalize and potentially even increase with model scale. Feature absorption and splitting are fundamental challenges in SAEs rather than model-specific issues. Based on theoretical considerations, larger models with more complex feature hierarchies might actually benefit more from our approach, as they contain richer hierarchical structures that Matryoshka SAEs are designed to preserve. Different architectures might show varying degrees of improvement, with models having stronger hierarchical representations potentially showing more pronounced benefits. We acknowledge the current experimental scope as a limitation and will add a note to this effect in the camera-ready version.

Thank you again for your positive assessment and valuable question. We hope our response reinforces your support of our paper.

Thank you for your very positive review of our paper! We were happy to read your comment about the experimental design being excellent and your appreciation for the simplicity and novelty of the idea.

Furthermore, we appreciate your suggestion to include a few feature examples, which we will do in the camera-ready version. We will also make sure to minimize the whitespace in this final version. Regarding the PCA baseline: Your intuition is exactly correct. The high variance explained and low CE loss are indeed due to the lack of a sparsity constraint, allowing it to use its full capacity for reconstruction unlike the SAEs. We will add a sentence clarifying this in the camera-ready version.

Thank you again for your strong support and enthusiastic review!

Thank you for your positive review and insightful questions. We appreciate your assessment that our paper is clearly written and the idea is novel. We address your key questions below:

Regarding the toy model demonstration:

1. On d < L setup: You correctly note the toy model uses a non-overcomplete set-up, for simplicity and illustrative purposes. The toy model is not meant to be a realistic depiction of features in LLMs, but rather a simplified demonstration to illustrate a specific scenario where vanilla SAEs struggle with feature absorption while Matryoshka SAEs excel. It provides a controlled environment to highlight the core mechanism of our approach. We chose to use a small number of latents to make it easier to visualize and show the resulting learned features, but expect that similar results would hold for toy models with more ground-truth features.

2. On adaptive l1 penalty: Thanks for pointing this out, we should have been clearer here. The adaptive sparsity control is not a Matryoshka SAE-specific component, but rather a component we use in both SAEs. The only difference between the two SAEs is the nested reconstruction objective of the Matryoshka SAE. We will update this in the camera-ready version.

3. On orthogonal features: You raise an important point about the correlation between parent and child features. Although the parent and child features are orthogonal, the resulting combined activations of the "parent + child" are not orthogonal to each other, due to the shared parent component. This is why we observe a small correlation between parent and child features in the Matryoshka SAE (Figure 3), as the disentanglement is, although close, not perfect. Given that there is some earlier work by Park et al [1] that suggests that in real-world settings, parent and child features are often orthogonal, we think it is reasonable to use orthogonal features in our toy model as well.

Regarding the first-letter experiment: To find the relevant latents for the first-letter task, we use a train and test set of single token inputs sampled directly from the tokenizer, unweighted by token frequency. We first train logistic regression probes on residual stream activations to define the ground truth direction for each letter. We then use k-sparse probing on the SAE latents to identify the main latent(s) representing that feature. Specifically, we start with k=1 to find the primary latent, then incrementally increase k, adding additional latents to our "main latents" set when they improve the F1 score by more than a threshold (τ=0.03). For each token, we then measure when these main SAE latents fail to fully activate on tokens starting with that letter (their projection onto the ground truth direction is less than the model's projection) while other latents (the absorbing latents) compensate on these tokens. Since we use the SAEBench implementation for our evaluation, we refer to their paper for more details [2].

Regarding sparse probing performance: The performance drop at higher sparsity in Figure 7 is indeed interesting, but it's important to note that this phenomenon occurs across many other baseline architectures as well, not just in Matryoshka SAEs. This appears to be a general characteristic of SAEs rather than a limitation specific to our approach. We want to note that the performance drop in sparse probing when scaling up the dictionary size from 16k to 65k (Figure 10) is very minor (0.775 -> 0.772) and we don't believe this is a significant effect.

Regarding BatchTopK batch size: We chose a batch size of 2048 to maintain consistency with baselines as you suspected. While we didn't perform extensive ablations on this hyperparameter, we suspect that for reasonably large batch sizes it is not a critical hyperparameter, and that standard BatchTopK SAEs and Matryoshka BatchTopK SAEs will be similarly affected by changes in batch size.

We hope these explanations address your questions. We would appreciate your consideration if these clarifications strengthen your support of the paper.

[1] Park, Kiho, et al. "The geometry of categorical and hierarchical concepts in large language models." arXiv preprint arXiv:2406.01506 (2024).

[2] https://www.neuronpedia.org/sae-bench/info#feature-absorption

Thank you for your detailed feedback and valuable suggestions. We are particularly encouraged that you found the paper well-written, addressing an important problem, and believe it is a valuable contribution that should be accepted.

We address your key questions below:

1. You asked why specialized concepts emerge without explicit constraints. As you suggest, the driving force is indeed the reconstruction loss. The overall objective (summed over all dictionary sizes) incentivizes the full SAE to minimize reconstruction error. Specialized latents emerge naturally because they capture finer details present in some inputs that cannot be adequately represented by the limited number of coarser features alone, thus minimizing the total reconstruction error across the dataset. Our experiments (Appendix C) confirm that later latents, while activating less frequently, significantly contribute to reconstruction quality.

2. We appreciate this suggestion! Explicitly enforcing the number of active latents per group might indeed further improve granularity control. We'll include this in our discussion as a promising future research direction.

3. Thanks for the suggestion! We agree that it would be beneficial to include qualitative examples of parent-child relationships found in Matryoshka SAEs. We will include these in the camera-ready version.

4. You noted the absorption metric relies on the first-letter task. This is correct and follows the methodology established by SAE Bench [1] and Chanin et al. [2], benchmarks widely used in the community. While we acknowledge that broadening the tasks for measuring absorption could be beneficial (and will note this in our limitations), developing improved absorption metrics is beyond the scope of this work. Crucially, however, we want to emphasize that our other key evaluations, such as sparse probing and concept isolation, do utilize a diverse range of classification tasks.

5. Thanks for pointing out this omission. We sample the prefix lengths from a discrete distribution inspired by the Pareto distribution. Specifically, each possible prefix length $l \in \{1, \ldots, m\}$, where $m$ is the total number of latents, is assigned a probability proportional to:

where $\alpha > 0$ controls how heavily the distribution favors shorter prefixes (we use $\alpha = 0.5$ in our experiments). This yields a monotonically decreasing probability distribution over prefix lengths. We then normalize this distribution and sample prefix lengths without replacement, always including the full prefix length $m$ to ensure the complete SAE is trained. We will add this to the camera-ready version.

6. Figure 4 shows absorption in vanilla SAEs. As requested, we will add a corresponding figure for the Matryoshka SAE in the camera-ready version, demonstrating how the same feature (e.g., "female tokens") remains distinct (see https://sparselatents.com/tree_view?loc=residual+pre-attn&layer=3&sae_type=S%2F2&latent=65 for the example). We did not experiment with varying $K$ in this experiment, but we extensively tested this quantitatively in section 4.3.

7. Thank you for pointing out that the metrics in Figure 6 are not clearly defined. The metrics in Figure 6 are calculated the same as in SAEBench [1]. We will clarify them in the camera-ready version. The metrics are defined as follows:

   • The splitting metric counts how many latents are needed to represent a single feature. We measure this by training k-sparse probes and detecting when increasing k by one causes a jump in F1 score by more than threshold τ = 0.03. This indicates that additional latents contain significant information about the feature, suggesting the feature has been split across multiple latents.

   • The absorption rate measures the fraction of tokens where a latent corresponding to a first-letter feature fails to activate despite the token starting with that letter due to the feature being absorbed by another latent. First, we find false-negative tokens where the main SAE latents fail to fully activate on tokens starting with that letter (their projection onto the ground truth direction is less than the model's projection) while other latents (the absorbing latents) compensate on these tokens. Since we use the SAEBench implementation for our evaluation, we refer to their paper for more details [1].

Thanks again for your valuable suggestions. We hope these clarifications and our planned revisions address your concerns. Given your positive assessment of our work's contribution and importance, we would be grateful if you would consider strengthening your support.

[1] https://www.neuronpedia.org/sae-bench/info#feature-absorption

[2] Chanin, David, et al. "A is for absorption: Studying feature splitting and absorption in sparse autoencoders." arXiv preprint arXiv:2409.14507 (2024).