Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Authors: Bart Bussmann
Noa Nabeshima
Adam Karvonen
Neel Nanda
Venue: ICML 2025 poster (2025)
Keywords: sparse autoencoders
mechanistic interpretability
Primary area: social_aspects->accountability_transparency_and_interpretability
TL;DR: We introduce Matryoshka SAEs to mitigate feature splitting and absorption in large SAEs. By training nested dictionaries to enforce a feature hierarchy, our method improves feature disentanglement and performance on downstream interpretability tasks.

Abstract: Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting neural networks by extracting the concepts represented in their activations. However, choosing the size of the SAE dictionary (i.e. number of learned concepts) creates a tension: as dictionary size increases to capture more relevant concepts, sparsity incentivizes features to be split or absorbed into more specific features, leaving high-level features missing or warped. We introduce Matryoshka SAEs, a novel variant that addresses these issues by simultaneously training multiple nested dictionaries of increasing size, forcing the smaller dictionaries to independently reconstruct the inputs without using the larger dictionaries. This organizes features hierarchically - the smaller dictionaries learn general concepts, while the larger dictionaries learn more specific concepts, without incentive to absorb the high-level features. We train Matryoshka SAEs on Gemma-2-2B and TinyStories and find superior performance on sparse probing and targeted concept erasure tasks, more disentangled concept representations, and reduced feature absorption. While there is a minor tradeoff with reconstruction performance, we believe Matryoshka SAEs are a superior alternative for practical tasks, as they enable training arbitrarily large SAEs while retaining interpretable features at different levels of abstraction.

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Bart Bussmann"
Noa Nabeshima"
Adam Karvonen Neel Nanda
Abstract
extract meaningful, cumaaninterprretbbie features that pro- vide insight into a model's decision-making process. In practice, however, sparsity is only an imperfect proxy for interpretability; optimizing represent inputs using as few active latents possible can lead to pathological solutions.
Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting neural networks by extracting the concepts represented in their acti- vations. However, choosing the size of the SAE dictionary (i.e. number of learned concepts) cre- ates tension: as dictionary size increases cap- ture more relevant concepts, sparsity incentivizes features to be split or absorbed into more spe- cific features, leaving high-level features missing or warped. We introduce Matryoshka SAEs, a novel variant that addresses these issues by simul- taneously training multiple nested dictionaries of increasing size, forcing the smaller dictionaries to independently reconstruct the inputs without using the larger dictionaries. This organizes fea- tures hierarchically- the smaller dictionaries learn general concepts, while the larger dictionaries learn more specific concepts, without incentive to absorb the high-level features. We train Ma- tryoshka SAEs on Gemma-2-22 and TinyStories and find superior performance on sparse prob- ing and targeted concept erasure tasks, more dis- entangled concept representations, and reduced feature absorption. While there is minor trade- off with reconstruction performance, we believe Matryoshka SAEs are a superior alternative for practical tasks, as they enable training arbitrarily large SAFs while retaining interpretable features at different levels of abstraction.
A core issue lies in the mismatch between the flat sparsity ob- jective and the inherent hierarchical structure of real-world features. For example, a broad, general concept (e.g. "punc- tuation marks") might be a parent feature to more specific ones ("punctuation mark and period", "punctuation mark and question mark," and "punctuation mark and comma"), When an SAE has the capacity to learn many latents, the sparsity penalty can push it to replace each general concept with a set of narrowly specialized features without retain- ing the original high-level category - phenomenon called feature splitting (Bricken et al., 2023). Even more prob- lematic is feature absorption (Chanin et al., 2024), where a parent feature only partially splits and number of specific instances of general feature absorbed by more spe- cialized latents, leaving "holes" in the representation of the more general feature (e.g. a latent that activates on all tokens starting with an F,, except if the word is Elephant). Finally, feature composition occurs when the SAE merges distinct concepts (like "red" and "triangle") into single composite features ("red triangle") to minimize the number of active latents, instead of learning the underlying features (Anders et al., 2024; Wattenberg & Viégas, 2024; Leask et al., 2025).
These problems become more severe as SAES scale to larger dictionary sizes (i.e., more learned concepts), since the model exploits the additional capacity to further minimize active latents through mechanisms like feature splitting (Kar- vonen al., 2024a). This presents stark contrast to typical deep learning, where scaling consistently improves both training loss and model capabilities (Kaplan et al., 2020; Hoffmann et al., 2022). While larger SAES do achieve lower training loss (Gao et al., 2024; Templeton, 2024), their per- formance on downstream tasks often deteriorates (Karvonen et al., 2024a). This fundamental tension creates significant practical challenges for deploying SAEs in interpretabil- ity research they become less reliable for probing model behavior, steering models, analyzing feature circuits, or un- derstanding how models process information. The current recommended approach is to train a sweep of different sizes and evaluate each of them, but this is costly (Lieberum et al.,
1. Introduction
Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting the internal representations of neural networks (Bricken et al., 2023; Cunningham et al., 2023). By training autoencoders map dense activations of neu- ral networks into sparse latent space, researchers aim to
'qqual contribution. Author order determined by coin flip. Cor- respondence to: Bart Bussmann <bartbussmann& gomm.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Calculate prefix loss Ly for each sampled prefix
Sample SAE prefixes
Sum losses and update all SAE weights
Ls
flower to
033
L L
La
L12
Unerplamed
Hale
Courre Similanty
Figure 1. Architecture and Performance of Matryoshka SAEs. (Left) The model leamss multiple nested reconstructions simultancously, with each reconstruction using only subset of the total latents. This creates pressure for carly latents capture general features while later latents can specialize in more specific concepts. (Right) Comparative metrics between SAEs with average sparsity of 40 active latents per token (L0=40) showing that while Matryoshka SAEs sacrifice a small amount reconstruction fidelity (higher variance unexplained). they achieve significantly lower feature absorption rates and less feature composition (lower decoder cosine similarity).
2024). Ideally, we would like to have single SAE that can capture both high-level and fine-grained features.
in interpretability research. We find that concerning failure modes that were thought to be fundamental barriers to us- ing SAEs, can be mitigated through simple changes in the training procedure. This indicates that other problems with SAEs might be easier to solve than thought, and illustrates the promise of more research this area.
Therefore, we introduce Matryoshka SAEs, a novel hierar- chical approach to SAE training, named after the Russian nesting dolls and inspired by Kusupati et al. (2024). Ma- tryoshka SAEs learn a single feature space that retains both abstract and specific features by training multiple nested SAES of increasing size simultaneously. Each nested sub- SAE is optimized to reconstruct the input using only subset of the total latents (including the sub-SAEs nested within itself). This nested structure prevents later, more specialized latents from absorbing the roles of earlier, more general ones, effectively regularizing the SAE to learn features at multiple levels of abstraction. Despite the need for recon- structions at multiple scales, Matryoshka SAEs can be im- plemented efficiently with modest increase in training time compared to traditional SAEs, depending on the number of nested sub-SAEs.
2. Background
2.1. Interpretabilityy with Sparse Autoencoders
Early mechanistic interpretability work sought to under- stand neural networks by examining individual architectural components, but faced challenges due to polysemanticity where components seem to exhibit multiple unrelated be- haviors (Olah et al., 2020; Elhage et al., 2022). In recent years, SAEs emerged promising solution for extract- ing monosemantic and interpretable features from model activations (Bricken et al., 2023; Cunningham et al., 2023; Templeton, 2024; Gao et al., 2024; Rajamanoharan et al., 2024a:b), which researchers can analyze through manual inspection (Bricken et al., 2023; Lin, 2023) or automated interpretability techniques (Gao et al., 2024; Juang et al., 2024). Their effectiveness has been demonstrated across var- ious applications, such as analyzing attention mechanisms in GPT-2 (Kissane al., 2024), circuit analysis (Marks et al., 2024; Makelov et al., 2024), and model-diffing (Lindsey et al., 2024).
Through extensive experiments on both synthetic and real- world datasets, we demonstrate that Matryoshka SAEs miti- gate feature absorption across wide range of model sizes and sparsity levels. We find that Matryoshka SAEs have more disentangled latent representations (as measured by maximum cosine similarity of decoder vectors) and also im- prove performance on probing and targeted concept erasure tasks compared to standard SAE baselines. Furthermore, we find that with increasing dictionary size, Matryoshka SAES often improve or maintain their performance on downstream tasks, where other architectures deteriorate in performance. While Matryoshka SAES exhibit small trade-off in recon- struction error due to the additional constraints, we argue that the improved quality of the learned latents, which also get better with scaling, make them an attractive alternative for practical applications.
Technically, an SAE consists of an encoder and a decoder:
f(x) oWWnn x enn
& wolff(x) + dec
We consider our results a positive sign for the use of SAEs
The encoder trained to map high- high-dimensional activations or hidden states from inside a neural network into sparse,
2

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
overcomplete representation f(x), called the SAE latents or features. The activation function a enforces non-negativity and sparsity in f(x), with a latent i typically considered active when fx) > 0.
fragment, absorption creates specific "holes" in otherwise intact general features, making the general feature less reli- able for downstream tasks and harder to interpret.
Feature composition represents another sparsity- -riven dis- tortion that occurs when features naturally co-occur. When presented with a direct product space of independent fea- tures (e.g. colors and shapes), the sparsity objective incen- tivizes learning single latents that capture specific combi- nations (like "red triangle") rather than representing the underlying independent features ("red" and "triangle") sep- arately (Anders et al., 2024; Wattenberg & Viégas, 2024), Even when these features are conceptually and functionally independent in the model, combining them allows the SAE to achieve the same reconstruction with fewer active latents, Leask al. (2025) have shown using "meta-SABs" that the latents of sparse autoencoders are often composed of more fundamental "meta-latents".
The decoder maps the sparse features back to the original input space through a linear transformation, producing an approximate reconstruction x. This reconstruction aims to minimize the mean squared error with respect to the original input x while maintaining the sparsity constraints the intermediate representation.
Recent architectural improvements include JumpReLU SAEs (Rajamanoharan et al., 2024b), which learn individual activation thresholds, and Topk SAEs (Gao et al., 2024), which enforce exactly K active features per sample. For our Gemma-2-2B experiments, we use the BatchTopK (Buss- mann et al., 2024) activation function, which we further describe in Section 3.2.
These issues become more severe as dictionary size in- creases. While larger dictionaries allow for representing more relevant features, they also enable more opportunities for sparsity- distortions. The reconstruction objective alone does not prevent these distortions since they achieve similar reconstruction with fewer active latents. Their per- sistence suggests they arise from fundamental limitations in standard SAE training objectives rather than optimization challenges. This creates a pressing need for new training approaches that can preserve features at multiple levels of abstraction while maintaining the benefits of sparsity. Ide- ally, we want SAEs that can capture both high-level concepts and their refinements without letting specialized features absorb or fragment their general counterparts.
2.2. Challenges in Scaling SAEs
At the heart of SAE training lies a fundamental tension between reconstruction quality and sparsity. The sparsity objective, typically implemented via LI regularization or explicit LO constraints such as TopK, encourages SAES to represent inputs using as few active latents as possible. While some degree of sparsity is necessary to extract in- terpretable features from superposition, optimizing for the reconstruction + sparsity loss often leads to pathological solutions. This manifests in several interrelated phenomena:
Feature splitting (Bricken et al., 2023) causes unified con- cepts to fragment into many specialized latents for exam- plc, feature responding to punctuation marks may split into separate features for question marks, periods, and commas. Even though each individual, specialized latent might be interpretable, the high-level concept of "punctuation mark" which might be functionally used by the LLM is now missing from the dictionary. Instead, the "punctuation mark" decoder direction is now implicitly represented in the de- coder vector of each specific latent, as for example "comma" also implies "punctuation mark". Leask al. (2025) found that when doubling the dictionary size of SAE, only one-third of the latents in the larger SAE represent novel concepts, while the other latents represent concepts similar to those of the smaller SAE in a sparser way.
2.3. Matryoshka Representation Learning
Matryoshka Representation Learning (MRL) (Kusupati et al., 2024) is an approach to representation learning that encodes information at varying levels of granularity within a single embedding vector. This allows the representation to be adaptable to the computational constraints of diverse downstream tasks. MRL is designed to modify existing representation learning pipelines with minimal overhead during training and no additional costs during inference and deployment. The key distinction of MRL lies in its loss function, which is modified to incorporate the perfor- mance when using only part of the embedding vector. The name "Matryoshka" is inspired by the nested Russian dolls, reflecting the nested nature of the representations.
Feature absorption (Chanin et al., 2024) occurs when a la- tent representing a general feature develops systematic blind spots in cases handled by more specialized latents. For ex- ample, consider latent that activates female names like "Mary", "Jane". "Sarah", and "Lily". : If specialized latent splits for the name "Lily"", the latent for female names might instead become an "all female names except Lily" latent. Unlike feature splitting where general concepts completely
3. Matryoshka Sparse Autoencoders
Matryoshka SAES extend traditional sparse autoencoders by simultancously training multiple nested autoencoders of increasing size, as illustrated in Figure 1. Given a maximum
3

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
dictionary size m, we define sequence of nested dictionary sizes M m,; -,,, My where m < m22 < ... < mh m. Each size mi corresponds to a sub-SAE that must reconstruct the imput using only the first m, latents.
BatchTopk X o 1[X T())
(6)
where 7 -(X) is the (B x K K)th largest value in X, @ denotes element-wise multiplication, and 1[-] is the indicator func- tion. This allows the number of active latents per example to vary naturally while maintaining an average sparsity of K' across the batch. During inference, BatchTTopK is re- placed with a global threshold to ensure consistent behavior independent of the batch:
We experiment with two types of Matryoshka SAE: random prefix, where M is randomly sampled from a distribution per batch, and fixed prefix where M is set as a hyperparam- eter before training. Formally, for an input x R², the encoder and decoder are defined as:
f(x) = a (Weeeccy x benc)
(3)
BatchTopk(x))
1[x 2 0)
(f) Wdde
Oam... .oom; + bder for m, E
(4)
where I is calibrated on the training data to maintain the desired average sparsity.
where Wene E R=**" is the encoder matrix, Wale E Raxm the decoder matrix, bone RT is the encoder bias, bdle E IR" is the decoder bias, and the subscript notation 0: 0:m; indicates taking the first m; rows elements. Each nested decoder Wdee 0:m.:: must learn to reconstruct the input using only a subset of the latents, creating hierarchy of representations at different scales.
4. Experiments
To understand how Matryoshka SAEs prevent feature ab- sorption and fragmentation, we conduct experiments across three settings of increasing complexity: a synthetic toy model designed to exhibit feature absorption, a qualita- tive investigation of the features learned on a small 4-layer TinyStories model, and finally benchmarks on the larger Gemma-2-2B architecture.
3.1. Training Objective
The key innovation in Matryoshka SAEs is the training ob- jective that enforces good reconstruction at multiple scales simultaneously:
4.1. Toy Model Demonstration of Feature Absorption
We first demonstrate the ability of Matryoshka SAES to avoid feature absorption in a controlled, synthetic setting, with a toy model following the model introduced in Chanin et al. (2024). Concretely, we construct tree-structured set of L binary features, each mapped tion in Rd (we use d 20). The root feature is always present (though it is excluded from L). Every other fea- ture is sampled conditional on its parent being present, with an associated edge probability P(parent-+child) E (0, 1). If a feature is active, we add its direction vector to produce the final d-dimensional input. This induces a hierarchical de- pendency analogous to "comma" = "punctuation mark" in real text: child features always appear with their parent features.
[(x)
Wider Ozma + d 112 + aCLxxx Σ ||- -((())) MEA
reconstruction using first m. latents
unique direc-
(5)
where Laux the standard auxilary loss as used in Gao et al. (2024). The first my latents must learn to reconstruct the input as well as possible on their own, the first m2 latents must do the same, and so on. Early latents must be able to reconstruct the input when m, is smaller, creating pressure for them to capture general, widely applicable features. Later latents only participate in reconstructing larger dictionaries, allowing them to specialize in more specific features.
We train two types of autoencoders on data generated from this model: a standard sparse autoencoder (Vanilla SAE) with L latents, matching the true number of features and a Matryoshka SAE with the same total latents L, but with additional nested reconstruction objectives on sub-prefixes of the latents. For the Matryoshka SAE, we sample the prefix lengths from truncated Pareto distribution at each training step, and use a ReI.U activation function (see Ap- pendix C for more training details). Both models have the same encoder and decoder architecture
3.2. Bop Activation Function
For our Gemma-2-2B language model experiments, we use the BatchTopK activation function (Bussmann et al., 2024) to determine which latents are active. The BatchTopK tivation function improves upon standard element-wist ap- proaches by considering sparsity across batches rather than individual examples. For a batch X = [x],..., xB). Batch- TopK retains the x K largest activations across the entire batch while setting others to zero:
Figure shows the ground-trui truth features against the closest-
4

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Feature Absorption Toy Model
Matryoshka Decoder, Ground Truth Cosine Simllrity
Varilla Decader, Ground- Truth Costne Similrity
Marry poothio Checosite Vettere
Vieila Drocoder Worter
Figure 2. Toy model decoder vector similarity. Graphical representation of the toy model (left). The heatmaps show the cosine similarity between learned latent vectors and ground -ruth feature vectors for the Matryoshka SAE (middle) and Vanilla SAE (right). The Matryoshka SAE shows a clear diagonal structure, which demonstrates disentanglement of the hierarchical features and learning the ground truth. The Vanilla SAE, however, exhibits high similarity between parent and child latents, indicating feature absorption.
parent feature and its three children
that the model is learning redundant representations and fail- ing to disentangle the hicrarchical features. The Matryoshka SAE, on the other hand, shows a clear diagonal structure in its similarity matrix, demonstrating that it is able to recover the true feature vectors with minimal overlap, with only small nonzero off-diagonal terms.
4.2. TinyStories Investigation
Batch
We now examine feature absorption in 4-layer Trans- former' (hidden size 768) trained on TinyStories (Eldan & Li, 2023), dataset of simple English children's 's stories. To systematically study how features evolve as SAE dictio- nary size increases, we train a family of small "reference" SAEs with varying dictionary sizes (30, 100, 300, 1k, 3k, and 10k latents) on three model locations: attention block outputs, MLP outputs, and the residual stream.
Ground-Truth Features
Matryoshka SAE Latent Acts
Vanilla SAE Latent Acts
These reference SAES reveal consistent pattern: smaller SAES (like S/1 with 100 latents) often capture broad, general concepts, while larger ones tend to fragment these concepts into specialized cases. To visualize these relationships, we developed an interactive tool (sparsel atentss....... www.///w.e///t.eeee latents evolve across SAE sizes. Technical details about the visualization methodology can be found in Appendix E,
Figure Toy model activations. Ground-truth feature activations alongside Matryoshka SAE activations and Vanilla activations for one of the parents and its children on the toy model. Notice that in the Vanilla SAE the parent latent (column bracketed in blue) does not fire when its children (bracketed in red) fire. For the activations of all features in the toy model, see Figure 11.
that traces how
matching Matryoshka and Vanilla activations on single batch for single parent and its children features (for plot of the activations of all features see Figure 11). The Vanilla SAE exhibits feature absorption - when one of the child features is active, the corresponding parent feature does not activate. contrast, the Matryoshka SAE recovers the true underlying feature structure, with parent features remaining active regardless of which child features are active.
"Female tokens" Absorption Example. In Figure 4, we highlight latent 65 in our 300-latent reference SAE (S/2/65), which consistently fires on female-coded tokens: "she," "her," "girl," and female names like "Lily" and "Sue." When we examine the corresponding latents in the 1000-latent SAE (S/3), we find that the general female-word. latent (S/3/66) has developed systematic "holes" stops fir- ing "Lily" because specialized "Lily-specific" latent (S/3/359) has absorbed that case. The original concept be- comes fragmented, obscuring that the model treats "Lily" as a female name.
To further see whether the autoencoders have found the ground truth features, we measure the cosine similarity be- tween the learned decoder vectors and the ground-truth feature vectors (Figure 2). For the Vanilla SAE, we observe high similarity between parent and child latents, indicating
Ihtppa: / wtggtththhheeg.g com. com//oaaaabehiimm/ t nymodel
5

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Latent 5/2, 300-latent Vanilla SAE
The SAEs were trained on the residual stream activations from layer 12 of Gemma 2-2B using 500M tokens sampled from The Pile (Gao et al., 2020). We use the Adam optimizer with a learning rate of 3 x 10-4 and batch size of 2048. No additional regularization terms are used beyond the implicit regularization from the multiple reconstruction objectives.
sisti extleciiins Enety
the poodd.
Female Tokens
deg. Rax (ull lavid to play with
Latent 65
enterr Mer Mi pobke to let crees cara des labe ure Holin the nowe
INGi leoe der, 1002 co wes dlartig la the pick 190 tu twe
We evaluate Matryoshka SAEs against six alternative SAE architectures (see Appendix A) to assess both basic perfor- mance and effectiveness in addressing feature absorption and composition challenges. We use the evaluations pro- vided by SAE Bench (Karvonen et al., 2024a) to compare Matryoshka SAEs to the baseline SAES across a wide vari- cty of tasks. In these evaluations, we use all latents of the Matryoshka SAE, In Appendix F, we study the performance of the different nested sub-SAEs.
Latents in 5/3, 1800-latent Vanilla SAE
siEtil retted Evemy mnde
starts
Female Tokens with Holes Latent 66
wetthe dey. Nax (and loee paay Mhe
posket cilla dorte dny Delyy wilking Aeael
IMC day. Sane wnnd dlayies is the par 190 Mi row
witer viteiin
plesing the aare with thit ses, They Jitdd to silde
Reconstruction and Downstream Performance. We first examine basic performance metrics of the SAEs: reconstruc- tion quality (measured by variance of the input explained by the reconstruction) and downstream performance (measured by cross cross-entropy loss when feeding reconstructed activa- tions back into the language model). Figure 5 shows these metrics across different sparsity levels.
Lily
ele may aod vttecking the blacte tupeties wne towes
day.
Latent 359
peccautee day Heett and site Aoart bes, bay, there are are
(BiS) the eat Luked skay dn pisk. They Libad stise
sut
( bay, ans Meed ald baly nome hat My youe
hear
Taued
Sue
(MCor1) 1tits gir) Mooel toaer tiph shitr I1 wer pertr fie
Latent 861
(Rcustin dar, Jinc gid sesel gall thand aas ds the sloer the plose 25 nnd
(M02 day, jittle for need Fue and sistt Mis Sccedd 1908 1900 cod
1.00
3.20
BelL
JampiteLL
3:10
Tooit
Figure 4. Feature absorption in a TinyStories model. Example activations showing how a general "female words" latent (S/2/65) develops holes a larger SAE (S/3/66) as specialized latents for "Lily" (S/3/359) and "Sue" (S/3/861) I) absorb specific cases. Feature absorption locations are circled red.
Batch Toak
ReLU with p-Annealing
Matryoshik Betch Topk
0.21
Goted
PCA Baseline
45
08
150
Figure 5. Reconstruction performance. The variance explained of Matryoshka SAEs is slightly worse than some competing ar- chitectures. However. looking at the downstream LLM CE loss, Matroshka SAES perform comparable, especially at larger LOs.
In contrast, when we train a 25k-latent Matryoshka SAE with the same average sparsity, it maintains the broad cover- age of female-coded tokens while still developing special- ized child latents (e.g. a Sue latent, Lily latent, and a female names latent without holes, see here). This preserves the hierarchical relationship between general and specific fea- tures, making the model's representations easier to interpret. In Appendix D. we show a few examples of features learned by the Matryoshka SAE.
Matryoshka SAEs consistently show slightly worse recon- struction compared standard BatchTopK SAES. For ex- ample, at L0 of 40 the reconstruction of Matryoshka SAES explains 70% of the variance in model activations, whereas BatchTopK SAEs explain 72% of the variance. This per- formance gap is expected, as the nested reconstruction con- straints prevent Matryoshka SAES from splitting and absorb- ing broad latents to optimize sparsity. Therefore, it would be surprising if there was not mild degradation here. The PCA baseline, included for reference, is not subjected to any sparsity constraints and can therefore utilize its full capacity for reconstruction, unlike the SAFS.
4.3. Larger-Scale Validation with Gemma-2-2B
In our Gemma-2-2B language model experiment, we use fixed prefix Matryoshka over random as we found that they attained slightly better eval metrics (see Ap- pendix H for details). We use a dictionary size of D = 65536 with five nested sub-SAEs where M = (2048, 6144, 14336, 30720, 65536). We train five Ma- tryoshka SAES with an average sparsity of respectively 20, 40, 80, 160, and 320 active latents per token.
However, downstream model loss tells a different story. De- spite worse reconstruction, at larger LOs, Matryoshka SAEs achieve comparable cross-entropy loss when their recon-
6

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
structed activations are fed back into the language model. This suggests that Matryoshka SAEs may be learning more meaningful features that better capture the language model's internal representations, even though they reconstruct the raw activations less accurately.
ture absorption and splitting compared the baseline ar- chitectures. For example, an average sparsity (LO) of 40 active latents per token, Matryoshka SAEs exhibit an ab- sorption rate of just 0,05, in contrast to 0.49 for BatchTopK SAFs. This indicates that Matryoshka SAEs are far more successful at preserving general first-letter features in the presence of specialized latents.
Feature Absorption and Splitting. We evaluate the tent to which Matryoshka SABs reduce two key pathologies that emerge when scaling sparse autoencoders: feature ab- sorption and feature splitting. Following the methodology of Chanin et al. (2024), we use first-letter classification tasks as a probe to measure these phenomena.
Similarly, Matryoshka SAEs require fewer latents to capture first-letter features, demonstrating reduced feature splitting. At an LO of 40, Matryoshka SAEs need only one latent per first letter on average, while BatchTopK SAEs split this information across three latents.
These results demonstrate that the hierarchical structure imposed by Matryoshka SAEs effective combating the feature absorption and splitting that typically emerge when increasing SAE dictionary size. Matryoshka SAEs are able to maintain coherent high-level features even in the presence of specialized low-level features.
ReLU
JumpletU
Topx
Batch TopK
RoLU with p-Annaaling
Matryoshks Uetch lepk
Gated
PCA Baseline
L3 LLOKET (3333877
10 (LCke 33303377
Sparse Probing and Targeted Concept Removal. We evaluate the quality of the features learned by the SAEs through three complementary analyses that measure differ- ent aspects of feature quality and disentanglement.
Figure 6. Feature absorption and splitting. Matryoshka SAEs significantly reduce feature absorption and splitting compared to baselines.
First, we perform sparse probing to assess whether the learned latents encode semantically meaningful concepts (Gao et al., 2024). We use suite of 35 binary classification tasks spanning five domains: professions, sentiment, lan- guage, code, and news. For each task, we encode the inputs through the SAF, apply mean pooling over non-padding tokens, select the most relevant latent, and train logistic regression probe. Strong performance on these tasks indi- cates that the SAE has Iearned relevant latents and are not missing, entangled, or split across multiple latents.
The absorption rate, reported in Figure 6 and calculated following SAFBench (Karvonen et al., 2024a), measures the fraction of tokens where latent corresponding to a first-letter feature fails to activate despite the token starting with that letter, because the feature is absorbed by another latent. To quantify feature absorption, we first train logis- tic regression probes on the raw activations to establish ground truth for which directions encode first-letter infor- mation. The absorption rate is then determined by finding felse-negative tokens where main SAE latents (identified via k-sparse probing for the letter feature) fail to fully ac- tivate (their projection onto the ground truth direction is less than the model's projection) while other "absorbing" la-
Spure Prokeng top Kenurry
0.78
0.0
6.15
tents compensate on these tokens. For further details on this metric, we refer to Chanin al. (2024) and the SAEBench implementation.
6,24
4.10
120
20
300
LD
LO
Secondly, we also analyze feature splitting. The splitting metric in Figure 6, also from SAEBench, counts how many latents are needed to represent a single first-letter feature. We measure this by training k-sparse probes on SAE latents and detecting when increasing & (the number of latents used by the probe) by one causes jump in F1 score by more than a threshold = 0.03. Such jump indicates that additional latents contain significant information about the feature, suggesting the feature has been split across those multiple latents.
ReLU
JumpReLL
6.12
The Mebe
Topk
0.10
Batch Topk
0.0
ReLU with p-Annealing
R008
Matryashika Batch Topk
Gated
PCA Baseline
20
160
100
LO (LOWor's sparsor)
Figure 7. Probing and concept isolation. Matryoshka SAEs score considcrably better at SCR and TPP than baseline models, and better at sparse probing at lower sparsitics.
As shown in Figure 6, Matryoshka SAES reduces both fea-
7

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Second, test the SAE's ability to separate dataset biases, such the spurious correlation between a profession like "nurse" always being associated with "female", using spuri- ous correlation removal (SCR) (Karvonen et al., 2024b). We train classifiers on intentionally biased data, then ablate the latents that encode the spurious signal. A higher resulting accuracy suggests the SAE was more effective at isolating the spurious correlation (e.g. gender), enabling the classifier to focus on the intended task (e.g. profession classification).
investigate this phenomenon further using meta-SAEs and find that the same holds.
Automated Interpretability. We evaluate feature inter- pretability using gpt4o-mini as an LLM judge. Following Bills et al. (2023) and Paulo et al. (2024), we generate fea- ture explanations for 1000 latents from range of activating dataset examples and score the explanation by asking an LLM judge to predict on which inputs a latent will be ac- tive. Figure shows that the interpretability of latents of Matryoshka SAEs comparable to standard BatchTopp SAEs and more interpretable than many alternative baseline architectures.
Third, we evaluate how well the SAE isolates individual concepts using Targeted Probe Perturbation (TPP) (Karvo- nen et al., 2024b). We ablate latents specific to one class and measure the change in accuracy of a classifier trained for that class as well as classifiers for other classes. If the ablated latents are properly disentangled, it should have an isolated causal effect- reducing accuracy on the relevant class probe while leaving other class probes unaffected.
ReLU
JumpReLU
0.85
TopK
Batch Tapk
autsimeer
ReLU p-Annealing
As shown in Figure 7, Matryoshka SAEs significantly out- perform all benchmark architectures on the TPP and SCR metrics, demonstrating they learn more disentangled and isolated concept representations. On the sparse probing tasks, Matryoshka SAEs achieve the best performance at lower sparsity levels and consistently outperform standard BotchTTppK SAEs across all sparsity levels. These results strongly suggest that the hierarchical structure imposed by Matryoshka SAFS leads to higher quality and more relevant learned features.
0.75
Matryoshka Batch Topk
Gated
PCA Baseline
20
40
160
320
-
LL
(Lower is sporser)
Figure 9, Automatic interpretabilityy score. Matryoshka SAE latents are among the most interpretable.
Scaling SAES To evaluate how Matryoshka SAEs perfor- mance changes when increasing dictionary size, train Matryoshka SAEs with dictionary sizes of 4k, 16k, and 65k. Figure 10 shows that on almost all metrics Matryoshka SAES improve or are stable with scale, whereas alterna- tive architectures often degrade. These results indicate that as one increases the dictionary size of the SAE to capture more of features of the LLM,, Matryoshka SABs may be the superior choice for many downstream tasks.
Feature Composition Analysis. To quantifly feature com- position and shared information between latents, we look at the average maximum cosine similarity between two latents in an SAE. If the SAE has many latents with high cosine similarity, this means that multiple latents are representing similar information, indicating composition of information.
0.60
ReLU
Coeine Siniliity
JumpReLU
0.35
TopK
0.8
0.30
Batch Topk
0.25
ReLU with p-Annealing
0.20
Matryoshka Batch TopK
835
Gated
0.70
LG
0.25
R022
Tairr
Figure 8. Average maximum cosine similarity. The maximium cosine similarity is considerably lower for Matryoshka SAEs, indi- cating less feature composition.
0.15
Maltyohhia Satch Typk Datee FCA Daseline
In Figure 8, we see that the average maximum cosine sim- ilarity between decoder vectors is substantially lower for Matryoshka SAEs, indicating feature composition is less of problem for Matryoshka SAEs. In Appendix F.2, we
Figure 10. Scaling evaluations. Performance across three differ- ent dictionary sizes. Average performance across LOs 40, 80, 160 are reported.
8

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
5. Limitations
Impact Statement
While Matryoshka SAEs demonstrate improved perfor- mance of downstream tasks, this comes the cost of slightly reduced reconstruction performance compared to standard SABs. Although downstream loss remains competitive, ap- plications requiring precise activation reconstruction may find this trade-off problematic.
This work advances mechanistic interpretability techniques for neural networks, Better understanding of model internals could help identify failure modes, verify intended behaviors, and guide the development of more reliable and safe Al sys- tems. However, advances in model interpretability may also accelerate progress in large language model development and deployment, potentially amplifying both positive and negative societal impacts of these technologies.
Furthermore, the multiple nested reconstruction objectives increases training time compared to standard SAEs. This increase is dependent on the length of the set of nested dictionary sizes (M). In our experiments with nested SAEs, this increased training time with circa 50%.
Acknowledgments
We extremely grateful for feedback, advice, edits, help- ful discussions, and support from Joel Becker, Gytis Daujo- tas, Julian D' D'Costa, Leo Gao, Collin Gray, Dan Hendrycks, Benjamin boffner-Brodsky, Mason Krug, Patrick Leaskm Hunter Lightman, Mark Lippmann, Charlie Rogers-Smith Logan R. Smith, Glen Taggart, Adly Templeton, and Matthew Wearden.
While we conducted several ablations (Appendix H), many design choices specific to Matryoshka SAES remain to be explored. For instance, explicitly enforcing the number of active latents for different latent groups within the Ma- tryoshka SAE has not been studied and might offer even finer control over the granularity of learned features. Such variations could be promising avenues for future research.
This research was made possible by funding from Light- speed Grants and Al Safety Support
Finally, our evaluation relies heavily on imperfect quantita- tive measures of feature quality and automated interpretabil- ity metrics and are mostly conducted a single LLM. For example, our evaluation of feature absorption relies primar- ily on first-letter classification tasks. While these standard- ized metrics provide valuable insights, they may not fully capture human-relevant aspects of feature interpretability or the practical utility of the learned representations for manual downstream analysis. Future work investigating the human interpretability of Matryoshka SAE latents and their applicability to practical interpretability tasks would be valuable.
References
Anders, E., Neo, C., Hoelscher- Sobeeheer--Obemairr, J., and Howard, J. N. Sparse autoencoders find composed features in small toy models. https: //ww. wttpss//www.esswrrr.t.. com//poots//50wga2222333113 sparse-autoencodee
lesswrong.
f find- cod--compoee---tturrss---- in 1lltoy, 2024.
Anthropic Interpretability Team. Training sparse au- toencoders.
6. Discussion and Conclusion
html# training saes, 2024. [Accessed January 20, 2025].
We have presented Matryoshka SAEs, a novel variant of sparse autoencoders that addresses key challenges with spar- sity when scaling SAEs. By enforcing hicrarchical feature learning through nested dictionaries, our approach achieves lower feature absorption rates and improved concept isola- tion compared to standard approaches. This demonstrates that the pathologies of scaling SAEs are not fundamental limitations, but rather artifacts of myopic training objectives that can be overcome through simple modifications to the training process.
Biderman, S., Schoelkopf, H., Anthony, Q., Bradley, H.,
O' O'Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., Skowron, A., Sutawika, L., and van der Wal, O. Pythia: A suite for analyzing large language models across training and scaling, 2023. URL https://arxiv. orgg abs 23004 01373.
Bills, S., Cammarata, N., Mossing, D., Tillman, H., Gao, L., Goh, G., Sutskever, I., Leike, J., Wu, J., and Saunders, W. Language models can explain neurons in language models. URL https://openaipublic. blob. core. windows. nel/neuron-explainer//ppnnnr 14.05. 2023), 2, 2023.
Looking forward, the success of this approach suggests that many apparent problems with SAEs may have surprisingly simple solutions waiting to be discovered. As we work to understand increasingly large and capable neural networks, approaches to interpretability that scale as well will become essential to better understanding these systems.
himl. accessed:
Bricken, T., Templeton, A., Batson, J., Chen, Jermyn, A., Conerly, T., Turner, N., Anil, C., Denison, C., Askell, A., et al. Towards monosemanticity: Decomposing language

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
models with dictionary learning. Transformer Circuits Thread, 2023.
Karvonen, A., Rager, C., Marks, S., and Nanda, N. Eval- uating sparse autoencoders on targeted concept erasure tasks. arXiv preprint arXiv:2411.1888 2024b.
Bussmann, B., Leask, P., and Nanda, N. Batchtopk sparse autoencoders. arXiv preprint arXiv:2412.064/1 2024.
Karvonen, A., Wright, B., Rager, C., Angell, R., Brinkmann, J., Smith, L., Verdun, C. M., Bau, D., and Marks, Measuring progress in dictionary learning for language model interpretability with board game models, 2024c. URL https: /ttpps:///axxiv arxiv org/ abs / 2000. 00113.
Chanin, D., Wilken-Smith, J., Dulka, T., Bhatnagar, H.,, and Bloom, J. A is for absorption: Studying feature splitting and absorption in sparse autoencoders. arXiv preprint arXiv:2409. 14507, 2024.
Kissane, C., Krzyzanowski, R., Bloom, I., Conmy, A., and Nanda, N. Interpreting attention layer outputs with sparse autoencoders. arXiv preprint arXiv:240.117759 2024.
Cunningham, H., Ewart, A., Riggs, Sharkey, L. Sparse autoencoders find highly inter- pretable features in language models. arXiv preprint arXiv:2309.0866/ 086000 2023.
I.,, Huben, R., and
Kusupati, A., Bhatt, G., Rege, A., Wallingford, M., Sinha, A., Ramanujan, V., Howard-Snyder, W., Chen, K., Kakade, S., Jain, P., and Farhadi, A. Matryoshkarepresee tation learning, 2024. URL abs / 12005.11477
Eldan, R. and Li, Y. Tinystories: How small can language models be and still speak coherent english?, 2023. URL https: arxii.i.... www.v.org//ooooooooo..t...oooorr abs 230..0075555
Elhage, N., Hume, T., Olsson, C., Schiefer, N., Henighan, T., Kravec, S., Hatfield-Dodds Z., Lasenby, R., Drain, D., Chen, C., et al. Toy models of superposition. arXiv preprint arXiv:2209. 10652, 2022.
Leask, P., Bussmann, B., Pearce, M., Bloom, J., Tigges, C., Moubayed, N. A., Sharkey, L., and Nanda, N. Sparse autoencoders do not find canonical units of analysis. 2025. URL https: : wawpiv.vvvvo.////////aa orrg abs/ 250.. 04878.
Gao, L.,, Biderman, S., Black, S., Golding, L.,, Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
Lieberum, T., Rajamanoharan, S., Conmy, A., Smith, L., Sonnerat, N., Varma, V., Kramár, J., Dragan, A., Shah, R., and Nanda, N. Gemma scope: Open sparse autoencoders everywhere all at once on gemma 2, 2024. URL https: / / //aaxiv..ooo www./vi.org/bb////.. 320.//ab///2...5533 05147.
Gao. L., la Tour, T. D., Tillman, H., Goh, G., Troll, R., Radford, A., Sutskever, I., Leike, J., and Wu, J. Scal- ing and evaluating sparse autoencoders. arXiv preprint arXiv:2406.040993 2024.
Lin, J. Neuronpedia: Interactive reference and tooling for analyzing neural networks, 2023. URL https: / //www www.nneuroppdia.. org. Software available from neuronpedia.org.
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L.., Hendricks, I.. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 15556, 2022.
Lindsey, J., Templeton, A., Marcus, J., Conerly, T., Batson, J., and Olah, C. Sparse crosscoders for cross-layer fea- tures and model diffing. Transformer Circuits Thread, 2024.
Juang, C., Paulo, G.. Drori, J., and Belrose, N. Open source automated interpretability for sparse autoencoder features. EleutherAl Blog, July, 30, 2024.
Makelov, A., Lange, G., and Nanda, N. Towards principled evaluations of sparse autoencoders for interpretability and control. arXiv preprint arXivv2445.. 00336666 2024.
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.088 083611 2020.
Marks, S., Rager, C., Michaud, E. J., Belinkov, Y., Bau, D., and Mueller, A. Sparse feature circuits: Discovering and editing interpretable causal graphs in language models. 2024.
arXiv preprint arXiv:2403.
Karvonen, A., Rager, C., Lin, J., Tigges, C., Bloom, J., Chanin, D., Lau, Y.-T., Farrell, E., Conmy, A., Mc- Dougall, C., Ayonrinde, K., Wearden, M., Marks, S., and Nanda, N. Saebench: A comprehensive benchmark for sparse autoencoders, December 2024a. URL https: // www. neuronpedia. org/ waaae-benchaee info. Ac- cessed: 2025-01-20.
Olah, C., Cammarata, N., Schubert, I.., Goh, G., Petrov, M., and Carter, S. Zoom in: An introduction circuits. Distill, 5()ee00024-001, 2020.
Paulo, G., Mallen, A., Juang, C., and Belrose, N. Automati- cally interpreting millions of features in large language models. arXiv preprint arXiv:2410. 13928, 2024.
10

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Rajamanoharan, S., Conmy, A., Smith, L., Lieberum, T., Varma, V., Kramár, J., Shah, R., and Nanda, N. Improving dictionary learning with gated sparse autoencoders. arXiv preprint arXiv:2404.1600 2024a.
Rajamanoharan, S., Lieberum, T., Sonnerat, N., Conmy, A., Varma, V., Kramár, J., and Nanda, N. Jumping ahead: Improving reconstruction fidelity with jumprelu sparse autoencoders. arXiv preprint arXiv:2407. 14435, 2024b.
Templeton, A. Scaling monosemanticity: Extracting in- terpretable features from claude 3 sonnet. Anthropic, 2024.
Wattenberg, M. and Viégas, F. Relational composition in neural networks: survey and call to action. arXiv preprint arXiv:2407.14662, 2024.

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
A. SAE Training Details
use baseline SAE suite provided by SAE Bench (Karvonen et al., 2024a). contains SAEs trained on layer of Gemma-2-28 and layer 8 of Pythia-160M, Biderman al. (2023). The SAE Bench suite contains proposed SAE variants.
SAE Bench Variants
ReLU (Anthropic Interpretability Team, 2024)
TopK (Gao et al., 2024)
BatchTopK (Bussmann al., 2024)
Gated (Rajamanoharan et al., 2024a)
JumpReL.U (Rajamanoharan al., 2024b)
P-Annealing (Karvonen al., 2024c)
Table List of evaluated sparse autoeneoder architectures
Our Matryoshka SAEs were trained in directly comparable manner, including identical data ordering and hyperparameters.
B. Code Implementation
This appendix shows example code implementation to get the reconstruction and loss of a Matryoshka SAE. For the code implementation used in our experiments see https://github. com/ batch top_k.tppppp.....
blob/main hwtp://iciiiiiiiriioooiii
def forward (self, input_dates
Compute features and apply sparsity
features = relu (input_data data - self.bias self encode))
sparse_features - batch_top_k (features)
Same features standard BatchTopk SAB
Dictionary size cutoffs
dict_sized 12048, 6144, 14336, 30720, 65536]
losses
start with blas current_output = self clone
Incrementally add ructions from each group for range
start_idx = 0 if -- else dict_sizes [i-l] end idx = dict_sizes/l
Add contribution from this group of features froup_features sparse_features start_idx:end_idd) group_weights del.fdecoderrcieo (strrdddrrsttrtt_iieeeeeddd current_cutput current_cutput + group_eeatures group_weights
start_idx:
Store and compute loss
losses.append ( (curreentoottpt - input_data) x*2) .))
Sum all losses
loss sum (Losses)
return total_loss current_outpur Return summed loss and final reconstruction
Listing Example code implementation Matryoshka SAE
C. Toy Model Details
Here we provide implementation details for training the SAEs on our synthetic hierarchical feature dataset. Code can be found at https::.............. //g/ thub. com/ iomnabeshinaaaaaa mat rywshka-ssssees -aaee....aas.............. The models are implemented PyTorch
12

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
and use the following configuration:
Model Architecture. Both the Vanilla and Matryoshka SAEs use the following hyperparamaters:
Input/output dimension: 20
Number of latents: 20
ReL.U activation function
Adaptive sparsity control targeting the ground-truth average to of 1.2338
Optimizer: Adam with learning rate 3e-2, betas=(0.5, 0.9375)
Training steps: 40,000
Batch size: 200
Gradient clipping norm: 1.0
Matryoshka-Speciffe Components. Per batch, 10 prefix lengths are sampled from a truncated Pareto(0.55 distribution, with the full SAE always included. Specifically, each possible prefix length l (1...., ml), where m the total number ( where a > 0 controls how heavily the distribution
of latents, is assigned a probability proportional to P(f) 0.5 in experiments). This yields monotonically decreasing probability distribution over prefix lengths. We then normalize this distribution and sample prefix lengths without replacement, always including the full prefix length to ensure the complete SAE is trained.
favors shorter prefixes (we
parent feature and its three children
Independent features
Ground-truit Features
Vanilla Activations
Matryoshka Activations
Botch
Figure 11. Full feature absorption toy model activations Ground-truth feature activations alongside Matryoshka SAE activations and Vanila activations. Notice for the Vanilla SAE, parent latents (columns bracketed in blue) do not fire when their children (bracketed red) fire.
D. TinyStories examples
Figure 12 shows a set examples features learned by the Matryoshka SAE trained on the TinyStories model. The first feature activates on third feature activates on nouns that refer the location where story takes place.
tokens, the second feature on adjectives indicating that something is small, whereas the
13

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
(necen)ow day, littleljesal namel sue Eound light shoe. and it rit just right, She ais very happy to have chee --dal stacted walk azound with hez nee shot. She sue her and
(BEGIN) upen a time, there was clunsy dog naned soot. Spot iiked play with big. rnd equaee tay. woid nu pack with his frland, allittie giel nased Lucy. Locy would these oquaze, Spet
laaii. Spat was as dbedient deg, always listening to his eos. sun, Jumn chese birth. Today, Spet soe betterfly wanted
to
it.olpoot tried the butterfly, but it flle away.
(RCGEN)One day, littlelated nased Lay seet the park,
angel, but she toe sloe... called out, "lgge.. please snst for nel" but
son poetty angal claving, angel
anilee (iipy marted
(REGIN)Once
girl naned See. Sue was restloss gixl. She liked to play and zun day. day, ahe found Altiin used wanted thumb to gently take
(BEGIN)ONce
tzee, Under the
lived the hale.
these little hole. tiny was caned Yet.. day, vet bag.
sadd,
floee. aas big anvelope with peutty caters. Isalll
(neces)ow day, littlelgttal named sally found envelope
(BEGIN]Once upee time, there aae elleeetd dag nseed Spet. dee day, Spot found big tip on the ground. It was vety fsiends ware going celobrate...
(BESIN) one day, little gisl naned Lily went to thellann age his lng. Lily walked over hin and noticed that the Sancy. bright
band
"Roy, can I open this envelope " ofher non said, "Yes, Halli,
special
wanted
on TT..
(R6GSN)00e littieligeen nomed LiTy found bex. box. these mire masy little pleces. Lill did not know shat uasted to play the gane. but fiest. she vanted to wach the pleces.
(BeGIN)Once upon time, there was allleest bay naned Tin. The Leved spend with his won. They liked may
(BEGIN)ONN One day, cheestul giel named Lily and her and big socks. Lily Liked the musiss lot..inn the minnnn
wanted to play music for Tim.. is non sadd, TTn, let's is
listen
Liby telling
to childtenn
people
lady
(BCGEN)
gro strong. So,
snall. cheap tulh in het gazden. She santed to help the buib
little
water
(BEGIN)Once
Tiny loved to play with his ball. One day. he west the pack play with fziend, Big the They turna ball... ball!" said Tiny.
allmalll
Tiny.
(B0GS1)) Dnce upen time, these was little bey naned Tie. thelittone with his to find something yumsy for his angty
sunahire
day, while talking bind came to wisit. The
feeding the bulb
cat. The big and
(RG4EN]0nce
liail ramed Sue. lisse sestless gizll. the liked to play
(BEGIN)Once upen time, there aas big. hold man named Ton in the park. One day, it wan vesy cold. and the wind made the trees trueble. Spot was scared shaking trees... -oo Soot
(eee and Ban are tuins, They like to play with their toye
and nn all day. Onn day, she found tiny but stuck de
has plcture
Saea and Ben sant
thueb to gently take
read book. Thee ank book.
they can have
Figure 12, A set of example features learned by the Matryoshka SAE trained on the TinyStories model.
E. Tree Methodology
To generate the hierarchical visualizations in sparselatents.come view, we implemented a variant of Masked Cosine Similarity (MCS), building upon the metric introduced in "Towards Monosemanticity" (Bricken et al., 2023).
E.1. Modified Masked Cosine Similarity
The original MCS between two latents A and B is calculated as follows:
1. Compute the cosine similarity between activations of A and B, but only for tokens where latent A is active
2. Compute the cosine similarity between activations of A and B, but only for tokens where latent B is active
3. Take the larger of these two similarity values as the final MCS
Our approach modifies MCS in two ways:
1. We only consider the cosine similarity between A and B's activations on tokens where B is active, for a directed version of MCS.
2. We scale this directed similarity by the ratio max(B activations) /aatiionnn////AAAAAAa activations), which we hoped would down- weight relationships involving latents with minimal activation
We use a more conservative similarity threshold of 0.6 compared to the original implementation.
E.2, Tree Construction Process
The tree generation proceeds through these steps:
1. Starting with parent latent (e.g., S/1/12), we identify all latents in a wider SAE model (e.g., S/2) that exceed our directed MCS threshold
2. For each identified S/2 latent, we recursively apply the same methodology to find children in S/3
3. We continue this process to build the complete hierarchical structure
14

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
E.3. Handling Multi-Parent Relationships
The resulting structure often forms directed acyclic graph (DAG) rather than a strict tree, some latents deeper layers (e.g,, S/3) may have relationships exceeding the similarity threshold with multiple parent latents in the previous layer (e.g.. S/2).
To simplify visualization, we assign each multi-parent latent exclusively to the parent with which it exhibits the highest similarity score. While this approach necessarily obscures some of the complexity in the underlying network structure, enables clearer visualization and interpretation.
E.4. Interpretation
It's important to note that these tree visualizations should not be interpreted as comprehensive representations of the full model structure. Rather, they serve as targeted explorations highlighting sets of latents that potentially demonstrate feature absorption relationships when examined together. They provide window into how features may be organized and recomposed across different model widths.
F. Investigating the sub-SAEs
In addition to using the complete Matryoshka SAEs, we also investigate using only number of the sub-SAEs. For this analysis, we train a Matryoshka SAE with 36864 latents and 5 subgroups of (2304, 4608, 9216, 18432, and 36864). For fair comparison, we train five standard BatchTopK SAEs with dictionary sizes matching our nested sub-SAEs (2304, 4608, 9216, 18432, and 36864), each calibrated match the effective sparsity of the corresponding sub-SAr (average LO norms of 22, 25, 27, 29, and 32 respectively).
In Figure 13, we find distinct activation patterns across different latent groups, with early latents showing consistently higher activation rates due to their participation in multiple reconstruction objectives.
Latent Sparsity vs Latent Index
2304]
[2304, 4608]
[4608, 9216]
Latent Sparsity
10-2
[9216, 18432]
[18432, 36864)
10:5
5000
10000
15000
20000
25000
30000
35000
Latent Index
Figure 13. Latent frequency latent index. Distribution of latent activations across different groups in the Matryoshka SAE, Each color represents different nested dictionary size. Early latents (blue) show higher average activation rates, reflecting their use in multiple reconstruction objectives.
By using only a limited number of the sub-SAEs of the single Matryoshka SAE, we can construct dictionaries of different sizes. Figure 14 shows the reconstruction performance of the sub-SAEs compared to their matched BatchTopK SAEs. Although the BatK SAES explain a larger percentage of the variance of the input, the difference in degradation of downstream language model cross-entropy loss decreases when more sub-SAFs are used.
15

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
Performance Comparison: Mätryoshka BatchTopK SAES
Average LD
Variance Explainad
LM Loss Degradation
Logit KL Divergence
Parryasnaa
0.86
Hotryeshka
11
MatryoslNa
sarch
sanchtopk.
sanch
202.3ge LG Nomm
30-
Explaned
Logst Divergence
0.9
0.8
0.7
24
0.6-
Metryoshka Betut/Tip
-0.6
20200
20000
30000
10000
20000
30000
10000
2C003
30000
30000
20000
30000
kumber of interes used
Number of iatenss UsDd
number of lorents used
Number intenes used
Figure Reconstruction performance of sub-SAE. Although BatchTopK SAEs explain a larger percentage the variance of the input, the difference in degradation of downstream language model gross-entropy loss decreases when more sub-SAEs are used.
F.1. Hierarchies of Matryoshka Features
The nested training objective of Matryoshka SAEs is designed to encourage the formation of feature hierarchy. Early latents, belonging to smaller sub-SAEs, are pressured to capture broad, general concepts since they must reconstruct the input independently. Subsequent latents in larger sub-SAEs can then learn more specific features, potentially representing refinements or components of the earlier, more general features, without the same incentive to absorb or fragment these foundational concepts.
We can visualize these learned hierarchies by examining the co-occurance between features. Figure 15 illustrates such hierarchy. In this hierarchy, we observe clear parent-child relationship structure. The root latent (103) represents general "sports team names" concept that activates broadly any mention of sports teams. This general feature branches into more specific subcategories, such as rankings of sports teams and news headlines containing sports teams.
103
263
DIATL VeQUA nickmames
14183
20892
spors Mars
sperts names g. Tigers).
28135
Sports headlines with meeme rarnes.
Figure Hierarchy of features related sports names. The thickness of the connections indicates the fraction of inputs on which the child features co-occur with their parent fcatures.
16

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
F.2. Investigating Composition with Meta-SAEs
We train meta-SAEs on the decoder matrices of both architectures, Each meta-SAE has a dictionary size one-quarter that of its input SAE and uses average of active latents. The variance explained by these meta-SAEs serves proxy for shared information between latents.
Meta-SAEs explain substantially more variance in BatchTopK SAE decoder directions compared to Matryoshka SAE directions (Figure 16). For the largest dictionary size Meta-SAEs explain 55% of variance in BatchTopK decoder directions, whereas they only explain 42% of variance in Matryoshka SAE decoder directions.
This indicates that Matryoshka SAE latents are more disentangled, with less shared information between them. The effect strengthens with dictionary size: BatchTopK SAEs show increasing levels of shared structure while Matryoshka SAEs maintain relatively constant levels of disentanglement.
Meta-SAE Evaluation Results
W dec Fraction Variance Explained
0.52
0.50
0.48
0.46
Matryoshka
0.44
BatchTopk
5000
10000
15000
20000
25000
30000
35000
Dictionary
Figure 16. Meta-SAE Evaluation Comparison between Matryoshka SAEs and standard BatchTopK SAEs.
G. Evaluations with Board Game Models
Karvonen et al. (2024c) evaluated SAEs specialized models trained exclusively chess and Othello games, leveraging the well-defined ground truth features in these domains. Their evaluation created two key metrics:
Board reconstruction: The ability to reconstruct the game board from the SAE latents
Feature coverage: The alignment of SAE latents with predefined board game features
While these toy models offer controlled environment for testing, find that the results do not show significant differentia- tion between Matryoshka SAEs and other evaluated architectures (ReLU, P P-Anneal, Gated, and TopK) in terms of these metrics. In the low LO regime (LO < 150), Matryoshka performs comparably to other architectures on both metrics. Notably, the best overall performance achieved by each architecture tends to occur in this lower LO range. with peak scores being fairly comparable across architectures.
However, Matryoshka shows some degradation in performance higher L0 values. This behavior may be partially explained by the model architecture: Chess-GPP uses relatively small hidden dimension of 512, meaning that high LO values (2 150) represent significant fraction of the model's representational capacity. This differs substantially from our main experiments on Gemma-2B, where even our largest LO values represent much smaller fraction of the model's 's hidden dimension.
Given that our primary interest lies in understanding SAE behavior on large-scale language models that better reflect real-world applications, we focus our main analysis on results from Gemma-2-22 2-2B.. However, we include these toy model results for completeness and to facilitate comparison with prior work. Figure 17 shows the performance comparison on the two metrics from the chess model evaluation.
17

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
0.45
Citess Board Stale Reconsiruciin
0.40
0.35
0.6
0.30
0.25
0.4
Coverage
Natryeshka
Matryoshka
Gated SAE
0.2
Gated SAE
0.15
Standard w/ paansalinng
Standard w/ p-anncaling
0.10
TupK
TopK
0.0
100
350
150
200
350
Ln (Lower is sparser)
In (Lower sparser)
Figure Evaluation results the ChessGPT toy model showing (Icft) feature coverage and (right) board reconstruction accuracy across different architectures. While Matryoshka SAEs show comparable performance at low 10 values, they slightly underperform at higher compared to other architectures.
H. Ablations
To better understand the key components of Matryoshka SAEs and validate our design choices, we conduct series of ablation studies. These experiments systematically modify different aspects of the architecture while keeping other variables constant. All ablation studies were trained on 200M tokens from The Pile (Gao et al., 2020), with evaluations performed on examine both 16k and 65k dictionary sizes and find differences become more pronounced at
layer of Gemma-2-2B. larger sizes, so we focus on 65k our ablation studies.
H.1. Ablation: Loss Weighting
The standard Matryoshka SAE equally weights the reconstruction loss for each nested dictionary size. However, we can interpolate between Matryoshka and standard BatchTopK behavior by weighting the losses according dictionary size. The BatchTopK SAE can be thought of special case of Matryoshka with single group.
We implement Weighted Matryoshka variant that assigns loss weights proportional to the number of new latents in each group. For nested dictionaries using [25%, 50%, 100%] of total latents, this yields weights [0.25, 0.25, 0.5] compared to the standard [0.33, 0.33, 0.33]. This weighting scheme emphasizes accurate reconstruction of the full dictionary while maintaining some pressure for hierarchical feature learning.
With our ablation training budget of 200M tokens, the Weighted Matryoshka achieves downstream cross-entropy loss similar to BatchTopK SAEs, outperforming the equally-weighted variant. However, this comes at the cost of reduced feature quality the weighted variant performs worse on feature absorption and spurious correlation removal metrics, suggesting that equal weighting is important for maintaining the hierarchical benefits of Matryoshka SAFs. Given that the cross-entropy gap between standard Matryoshka and BatchTopK disappears with larger training budgets, larger LOs, and larger dictionary sizes (as shown Section 4.3), these results support using equal weighting to preserve feature quality without significant reconstruction trade-offs.
18

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
0.995
0.88
Loss Recovered
0.990
Autointerp Score
0.87
0.985
0.86
0.85
0.975
0.84
0.970
102
10a
Lo (Sparsity)
Lo (Sparsity)
Mean Absorption Fraclion Score
1.0
Sparse Probing Top 1 Test Accuracy
0.80
0.78
0.8
0.7
0.74
0.6
0.72
0.5
0.70
102
10a
Lo (Sparsity)
Lo (Sparsity)
0.30
0.40
0.25
SCR 1op 20 Metric
TPP Top 20 Metric
0.35
0.20
0.30
0.15
0.10
0.25
0.05
0.20
10a Lo (Sparsity)
102
Lo (Sparsity)
Batch TopK
Matryoshka Batch TopK Weighted Matryoshka
Figure 18. Evaluation Results comparing the Matryoshka, Weighted Matryoshka, and BatchTopK 65K width Gemma-2-2B SAEs.
19

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
H.2. Ablation: Stop Gradients
We investigate how gradient flow between different dictionary sizes affects the learning dynamics of Matryoshka SAEs. In the standard implementation, the reconstruction for each nested dictionary incorporates the reconstructions from smaller dictionaries. This allows gradients to flow through the entire chain of reconstructions when optimizing the loss for larger dictionaries, gradients can affect how smaller dictionaries learn their features. In our stop gradient variant, we detach each partial reconstruction before it's used in the next one, forcing each group of latents to learn independently of the others.
These changes effectively isolate the training each nested dictionary, preventing larger dictionaries from influencing how smaller ones learn features. Our evaluation shows mixed results: while the stop gradient variant achieves better feature absorption scores low LO values, it shows degraded performance on loss recovered, sparse probing, and Spurious Correlation Removal (Figure H.2). Given these degradations, did not pursue this variant further, though future work with more sophisticated metrics or qualitative analysis may reveal additional benefits to this approach.
20

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
0,995
0.88
0.990
0.87
Loss Recovered
Autointerp Score
0.985
0.86
0.980
0.85
0.975
0.84
0.970
0.83
0.965
10a
10a
Lo (Sparsity)
(Sparsity)
1.0
Mean Absorption Fraclion Score
Sparse Probing Top Test Accuracy
0.80
0.9
0.78
0.8
0.76
0.7
0.74
0.6
0.72
0.5
0.70
10a
10a
Lo (Sparsity)
(Sparsity)
0.30
0.40
0.25
SCR Top 20 Metric
TPP Top 20 Metric
0.35
0.20
0.30
0.15
0.10
0.25
0.05
0.20
10a LD (Sparsity)
102 LD (Sparsity)
Batch TopK
Matryoshka Batch TopK
Stop Gradients
Figure 19, Evaluation Results comparing the Matryoshka, Stop Gradient Matryoshka, and BatchTopK 65K width Gemma-2-2 SAEs.
21

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
H.3. Ablation: Number of Nested Dictionaries
We investigate how the number of nested dictionaries affects Matryoshka SAE performance. Our standard implementation uses five nested dictionaries, with each group representing the following fractions of the total latents: [1/32, 1/16, 1/8, 1/4, (1/2 1/32)]. We compare this against two variants:
A three-group variant with ratios [1/8, 1/4, (5/8))
A ten-group variant with ratios (1/16384)99,, 152, 241, 383, 607, 964, 1531, 2430, 3857, 6123), derived from logarithmically spaced values between 128 and 8192 and normalized to sum to the total dictionary size. We chose this logarithmic spacing to avoid the extremely small dictionary sizes that would result from extending our standard geometric progression to 10 groups.
The three-group variant improves on loss recovered compared to our standard five-group implementation but performs worse on feature absorption and SCR, suggesting it interpolates between Matryoshka and BatchTopK behavior. Our ten-group variant slightly improves feature absorption but performs significantly worse loss recovered and automated interpretability metrics. Based on these results, we find that while the exact number of groups is not critical, five groups provides a reasonable balance between the trade-offs present in these metrics (Figure H.3).
22

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
0.995
0.88
0.990
Loss Recovered
Autointerp Score
0.985
0.980
0.86
0.975
0.85
0.970
0,965
0.84
102
10a
Lo (Sparsity)
Lo (Sparsity)
Mear Absorption Fraction Score
1.0
Sparse Probing Top 1 Test Accuracy
0.80
0.78
0.8
0.76
0.7
0.74
0.6
0.72
0.5
0.70
102
102 L0 (Sparsity)
Lo (Sparsity)
0.30
0.40
0.25
SCR Top 20 Metric
TPP Top 20 Metric
0.35
0.20
0.30
0.15
0.10
0.25
0.05
0.20
10a Lo (Sparsity)
102 L0 (Sparsity)
Batch TopK
Matryoshka Batch TopK Ten Fixed Groups
Three Fixed Groups
Figure 20. Evaluation Results comparing Matryoshka variants with different numbers of nested dictionaries (3, and 10 groups) against BatchTopK SAE, All models use 65K width on Gemma-2-2B
23

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
H.4. Ablation: Randomly Sampled Dictionary Sizes
As an alternative fixed dictionary sizes, we explored dynamically sampling prefix lengths from a truncated Pareto distribution during training. Specifically, each possible prefix length & {I....,m,, m), where m is the total number of latents, is assigned probability proportional to P(e) x 1 ( ( where > 0 controls how heavily the distribution favors shorter prefixes (we use 0.5 in our experiments). This yields 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.
Random sampling could theoretically create more continuous feature hierarchy by exposing the model to diverse dictionary sizes throughout training. We evaluated both strategies 16k-width SAEs with three different configurations for the number of sampled prefixes (3, and 10 groups).
On our current evaluation metrics, random sampling showed similar performance to fixed dictionary sizes, with minor degradation in spurious correlation removal and sparse probing performance, as shown in Figure 21. Random sampling may also introduce additional complexity for distributed training across multiple GPUs. While our evaluation framework may not capture all relevant aspects feature quality, we opted for fixed dictionary sizes for our main experiment given these considerations.
24

---

Learning Multi-Level Features with Matryoshka Sparse Autoencoders
0.875
0.990
0.870
Loss Recovered
0.985
Autointerp Score
0.865
0.980
0.860
0.975
0.855
0.850
0.965
0.845
0.960
0.840
10
(Sparsity)
Lo (Sparsity)
1.00
Sparse Probing Top 1 Test Accuracy
0.78
Mean Absorption Fraction Score
0.95
0.76
0.90
0.74
0.85
0.80
0.72
0.70
10a
10a
Lo (Sparsity)
LO (Sparsity)
0.40
0.200
0.38
0.175
SCR Top 20 Metric
TPP Top 20 Metric
0.36
0.150
0.34
0.125
0.32
0.100
0.30
0.075
0.28
0.050
0.26
0.025
10a
10a
LO (Sparsity)
LO (Sparsity)
Batch TopK
Matryoshka Batch TopK
Ten Random Groups
Five Random Groups
Three Random Groups
Figure 21. Random vs Fixed Dictionary Sizes. Comparison of Matryoshka SAE variants using randomly sampled dictionary sizes against fixed dictionary sizes and BatchTopK baseline. We evaluate configurations with 3, 5, and 10 groups on 16k-width SAEs trained on Gemma-2-20 Random sampling shows similar slightly worse performance compared fixed dictionary sizes across most metrics.
25


Reviews:
[{"review_id": "mKJSLekOPh", "reviewer": "Reviewer_gKza", "summary": "This paper aims to improve concept learning in sparse autoencoders (SAEs), which are models that have recently become popular as a means to disentangle features from large deep models, particularly LLMs, into a sparse set of disentangled concepts. This work focuses on the problem with existing SAEs of choosing the optimal size of the latent dimension, where large sizes may help capture more concepts, but cause them to have irregular and inconsistent granularities. Specifically, it looks into the problems of feature absorption, feature splitting, and feature composition. It proposes an alternative training objective inspired by Matryoshka representation learning, where the SAE is forced to reconstruct the input feature independently using varying latent sizes, enforcing some latents to encode coarse concepts and others to encode more specialized, fine-grained concepts. Experimental evaluation is performed to show that the proposed Matryoshka SAEs are performant and interpretable, while also avoiding the aforementioned problems with existing SAEs.\n\n## Update after rebuttal\nThank you for your detailed response. I agree that this paper should be accepted.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nWe address your key questions below:\n\n1. 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.\n    \n2. 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.\n    \n3. 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.\n    \n4. 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.\n    \n5. 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:\n$$\nP(l) \\propto 1 - \\left( \\frac{l}{m} \\right)^{\\alpha}\n$$\nwhere $\\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.\n\n6. 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. \n    \n7. 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:\n\n\t- 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.\n\n\t- 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].\n\nThanks 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.\n\n[1] https://www.neuronpedia.org/sae-bench/info#feature-absorption\n\n[2] Chanin, David, et al. \"A is for absorption: Studying feature splitting and absorption in sparse autoencoders.\" arXiv preprint arXiv:2409.14507 (2024).", "claims_and_evidence": "The claims are generally supported by evidence. Some of the experimental evaluation appears to be limited to specific tasks (e.g. first letter recognition) as discussed in the Weaknesses section below, and it would help to expand the evaluation to more tasks.", "methods_and_evaluation_criteria": "Yes, they make sense.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "The experimental design and analyses appear to be sound. It would help to have a more diverse set of experiments however, as discussed in the Weaknesses section below.", "relation_to_broader_scientific_literature": "This work builds upon existing literature that develops SAEs as a mechanistic interpretability tool to disentangle concepts learnt by deep models such as LLMs. While most prior research on improving SAEs has focused on improving sparsity of SAEs (e.g. TopK SAE, Gao et al. 2024), this work looks at improving the interpretability of SAEs, and to avoid the recently reported problems of feature absorption, splitting, and composition.", "essential_references_not_discussed": "None that I am aware of.", "other_strengths_and_weaknesses": "## Strengths\n1. This work deals with an important problem—while many attempts have been made to Pareto-improve the reconstruction error and sparsity tradeoff, this work focuses on improving the interpretability of SAEs and addresses issues caused by scale, which is a valuable contribution.\n2. Compared to the baselines, Matryoshka SAE performance improves with scale, which is promising.\n3. The idea of learning concepts at different granularities makes intuitive sense and the proposed approach generally appears to be sound.\n4. The paper is well written and easy to follow.\n\n## Weaknesses\n1. The training objective of the Matryoshka SAE as per Equation 5 only enforces that latent subsets from $0:m$ for all $m\\in\\mathcal M$ must independently be able to reconstruct the features. In particular, given $m_1<m_2$, there is no constraint that any latent outside of the first $m_1$ latents but within the first $m_2$ latents must activate at all. As a result, why should the SAE learn any specialized concepts at all? Is it simply that learning such concepts helps performance? An analysis of this, with respect to SAE size and reconstruction error, would be interesting to have, both in the toy setup and with real data.\n2. In addition to the previous point, it might be interesting to see if it helps to enforce specific values of $K$ for different latent granularities, to control how many coarse and fine-grained concepts the model learns.\n3. Currently there doesn't seem to be any evaluation to check how parent and child concepts relate, except in the toy setup. It would be useful to see if similar parent-child trees could also be constructed in the more realistic setups used in Section 4.2 and 4.3.\n4. Generally, the evaluations in Section 4.2 and 4.3 are observation-based and limited to very specific tasks, such as first letter recognition. While this is promising, it would be helpful to also extend to at least a few more tasks, to assess the generality of the method.\n5. The \"truncated Pareto distribution\" described in L210 is unclear, and there is no cited reference. A clarification would be helpful.\n6. In the experiments in Section 4.2, it would help to also show qualitative results with Matroyshka SAEs. L284-291 claims it avoids feature absorption, but without a qualitative figure. Additionally, what happens if $K$ is varied in this experiment?\n7. The metrics used in Figure 6 need to be defined precisely, they are currently unclear.", "other_comments_or_suggestions": "None.", "questions_for_authors": "Please refer to the Weaknesses section above. Overall I believe this work provides a valuable contribution and should be accepted. However, it would be helpful to have clarifications on the concerns raised in the rebuttal.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "W1CpKTpnwx", "reviewer": "Reviewer_HNXv", "summary": "This paper presents Matryoshka SAEs, inspired by Matryoshka representation learning, that learns a nested series of SAEs simultaneously to address issues such as feature splitting and feature absorption. Comparisons with multiple well-established baseline SAEs demonstrate Matryoshka SAEs’ superior quality to overcome these problems as well as exhibiting better concept isolation performance. Though their reconstruction performance is worse than baselines, Matryoshka SAEs show better automatic interpretability metrics and better scalability.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n\n**Regarding the toy model demonstration:**\n\n1. **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.\n\n2. **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.\n\n3. **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.\n\n**Regarding the first-letter experiment:**\nTo 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].\n\n**Regarding sparse probing performance:**\nThe 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.\n\n**Regarding BatchTopK batch size:**\nWe 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.\n\nWe hope these explanations address your questions. We would appreciate your consideration if these clarifications strengthen your support of the paper.\n\n[1] Park, Kiho, et al. \"The geometry of categorical and hierarchical concepts in large language models.\" arXiv preprint arXiv:2406.01506 (2024).\n\n[2] https://www.neuronpedia.org/sae-bench/info#feature-absorption", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes. The Matryoshka structure is well-suited to address the feature splitting/absorption problems, and the evaluation metrics used as well as the choice of baseline SAEs all follow the community standard.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes. All designs are valid, although I have the following concerns related to Section 4.1, Toy Model Demonstration of Feature Absorption.\n- Here, different from real-world scenarios, the authors are not training sparse or overcomplete autoencoders; instead, the number of latents is set equal to the ground-truth. According to the superposition hypothesis (e.g., [1]), an SAE is trained to approximate the high-dimensional hypothetical disentangled model using low-dimensional features from the observed model, which translates into d < L in this setup. Will the Matryoshka SAE maintain its superior performance under this condition?\n- As detailed in Appendix E, there are two important Matryoshka-specific components. First is an “adaptive sparsity control targeting the ground-truth average l_0 of 1.2338”. Could the author be more specific on what this means? How is this ground-truth information incorporated into the training process? The second component is the “l_1 sparsity penalty on normalized activations”. Does this imply that the vanilla SAE lacks an l_1 penalty? If so, then it would be beneficial to include ablations on the vanilla SAEs that also incorporate these biases.\n- In this setup, all ground-truth features are orthogonal to each other; however, in a more realistic setting, one would expect (at least) the parent feature and a child feature to naturally have a higher cosine similarity score compared to inter-group cosine similarity. I have this concern because I’ve spotted the correlation between the parent feature and the children features in the Matryoshka SAE as depicted in Figure 3, which is unexpected given the goal of avoiding absorption. Will this effect be amplified in a non-orthogonal setup?\n\n[1] Toy Models of Superposition. Section link: https://transformer-circuits.pub/2022/toy_model/index.html#motivation-superposition", "relation_to_broader_scientific_literature": "The proposed idea of incorporating Matryoshka representation learning into SAEs helps address several crucial problems in the current SAE community, which improves interpretability of the family of sparse autoencoders.", "essential_references_not_discussed": "I have not found essential missing references.", "other_strengths_and_weaknesses": "The paper is clearly written. Incorporating the Matryoshka learning idea into SAE is novel. It shows the potential to address problems in large-scale SAEs trained on LLMs.", "other_comments_or_suggestions": "I would suggest the authors mention the ablation studies in Appendix B (for example, use one sentence) into the main text, especially regarding loss weighting. It is noteworthy that average weighting works the best in this context.", "questions_for_authors": "I have the following questions for authors.\n- In the first-letter experiment, you first discover a direction via probing, and then “find the corresponding latents and measure when these SAE latents fail to activate on tokens starting with that letter and get absorbed into specific token-aligned latents.” How is this finding process performed? Is it through cosine similarity comparison? If this is the case, have you checked the token distribution with a specific letter, to avoid potential issues that lead to identifying a local feature/token feature with “holes” due to an imbalanced distribution (for example, corpus starts with L dominated by “Lily”)?\n- In the sparse probing experiment, the performance drops at higher sparsity, both in Figure 7 and Figure 10. This looks like a warning sign that when increasing the dictionary size, the Matryoshka SAE loses the ability (which it has when scaling is not performed) to some extent to isolate the target concept. Do you have an explanation or discussion on this phenomenon?\n- I wonder how the batch size setup in BatchTopK would affect the result. I suspect you chose 2048 to maintain consistency to the baselines, but if a stability result on this hyperparameter can be provided it would be beneficial.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ua1He3usK4", "reviewer": "Reviewer_phSQ", "summary": "The authors suggest a novel training objective for sparse autoencoders to address the issues of feature splitting, feature absorbtion and feature composition.\nThey test this idea on a toy, synthetic dataset explicitly designed to demonstrate improvements, then a 4-layer transformer-based language model traied on TinyStories, then Gemma-2-2B checkpoints (open-source language model from Google).\nThey demonstrate strong improvements through a variety of experiments, assuming well-tuned baselines.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nFurthermore, 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.\n\nThank you again for your strong support and enthusiastic review!", "claims_and_evidence": "The authors claim that despite worse reconstruction performance, Matryoshka SAEs better achieve the original goals of SAEs.\nThese claims are well-supported by their experiments in Section 4.3.", "methods_and_evaluation_criteria": "The methods (synthetic data, TinyStories, and Gemma-2-2B) and evalution criteria all make sense.\nIdeally there would be human evaluation as well, but the authors mention this as a limitation and have a public SAE latent viewer for TinyStories.\nI am very satisfied with the evaluation.\n\nIt would be interesting to see the effects of steering with the different latent subgroups.\nBut I'm happy for that to be future work.", "theoretical_claims": "There appear to be no theoretical claims.", "experimental_designs_or_analyses": "The experimental design is excellent.\nBaselines appear to be well-tuned.\nI did not find any issues in the soundness or validity of the experimental designs.", "relation_to_broader_scientific_literature": "This work fits in with other novel SAE architecture/objectives like TopK, JumpReLU, etc.\nIt explains its position within the broader literature well.", "essential_references_not_discussed": "There are no essential references not discussed.", "other_strengths_and_weaknesses": "The simplicity of the idea is phenomenal.\nMost SAE modifications are in the activation layer.\nIt's exciting to see novelty in the training objective, and I would love to see how this objective composes with other SOTA activation layers like JumpReLU.", "other_comments_or_suggestions": "* There's a fair amount of whitespace. Can you use that whitespace to make bigger (taller) figures?", "questions_for_authors": "* Can you explain why the PCA baseline has the highest variance explained and the lowest CE loss (figure 5)? Is it because PCA doesn't have a sparsity constraint?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yzMoxYWV7X", "reviewer": "Reviewer_3MV4", "summary": "This paper introduces a simple but novel approach for training SAEs with a nested structure in the feature space. As a consequence of this training, the authors present results suggesting that Matryoshka SAEs are more adept at overcoming feature splitting and feature absorption issues currently facing SAEs. The nested feature space concept of Matryoshka SAE is inspired by a previous work which introduces Matryoshka representations.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\nRegarding 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.\n\nThank you again for your positive assessment and valuable question. We hope our response reinforces your support of our paper.", "claims_and_evidence": "The claims made in the submission seem well supported by the experiments.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are sound.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Overall, the experimental design and analysis is sound.", "relation_to_broader_scientific_literature": "This work connects well with the existing literature. Feature splitting and feature absorption are known limitations of SAEs; this work presents a method which overcomes these issues while drawing inspiration from previous works, namely the SAE and Matryoshka representation learning literature.", "essential_references_not_discussed": "None", "other_strengths_and_weaknesses": "Strengths:\n- Paper is well written and has clear language\n- The set of experiments presented seems relevant and extensive. This includes presenting results for which their method is not necessarily the best (e.g., reconstruction and game board experiments)", "other_comments_or_suggestions": "None", "questions_for_authors": "Overall, the paper looks good, but I do have the following inquiry about the experiments: The extent of the large-scale experiments is on Gemma-2B. I understand that restrictions on compute may prohibit running experiments on additional models of similar or larger size, but evaluation on a single, relatively small language model is a limitation itself. I'm wondering if the authors have any insight on how the results may change as a function of model scale, architecture, or pretraining data?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes Matryoshka Sparse Autoencoders (Matryoshka SAEs), a new training method for sparse autoencoders designed to address interpretability issues such as feature splitting, feature absorption, and feature composition. The method is inspired by Matryoshka representation learning and introduces a nested structure, where the SAE is trained to reconstruct input features at multiple latent sizes. This encourages coarse features to be captured in smaller latent subsets and finer-grained features in larger ones. The approach is evaluated through experiments on synthetic data, a small transformer trained on TinyStories, and the Gemma-2-2B language model.

### Strengths

- Addresses an important problem in SAEs by targeting interpretability issues at scale rather than only focusing on sparsity vs. reconstruction trade-offs  
- Novel and intuitive objective that enforces multi-scale feature learning through nested reconstruction constraints  
- Demonstrates improved interpretability and concept isolation compared to strong baselines such as TopK SAE and JumpReLU SAE  
- Evaluations are performed on both synthetic and realistic language modeling tasks, including experiments on a 2B-parameter model  
- Reviewers highlighted the paper as well-written, with clear experimental design and extensive supplementary material  
- The approach is simple yet powerful, opening possibilities for future integration with advanced activation functions or architectures  
- Publicly available visualization tool (latent viewer) enhances transparency and usability of the method  

### Weaknesses

- Reconstruction quality is worse than some baselines, raising concerns about performance trade-offs  
- Evaluation tasks in certain sections (e.g., first-letter recognition) are narrow and may not reflect broader applicability  
- The toy model experiments use idealized settings (e.g., orthogonal ground truth, known latent count) that may not generalize well  
- Lack of evaluation of hierarchical (parent-child) relationships in real-world datasets beyond the toy model  
- Some methodological descriptions (e.g., truncated Pareto distribution, sparsity control) are unclear or underexplained  
- Unclear how well the method scales to larger language models beyond Gemma-2B, due to limited compute-based evaluation  
- Some metrics (e.g., in Figure 6) are not clearly defined, and more qualitative examples could improve clarity  

Most concerns have been addressed by the authors during the rebuttal period. All reviewers agree that this paper tackles an important problem and would be a great addition to the ICML program. The AC agrees.
Decision: Accept (poster)