How Matryoshka Sparse AutoEncoders Recover Feature Hierarchies That Vanilla SAEs Lose

A walkthrough of the core findings and guided replication of the concepts from the original research on "Multi-level features discovery with Matryoshka Sparse AutoEncoders".

Sparse AutoEncoders (SAEs) are a cornerstone of mechanistic interpretability, but they struggle with scalability. As we increase the dictionary size to capture more features, we often encounter "feature splitting" and "feature absorption," where general concepts are lost or broken into fragmented, less interpretable components. Matryoshka Sparse Autoencoders (MSAEs) solve this by training nested dictionaries simultaneously. This forces the model to organize features hierarchically, where smaller dictionaries capture more general concepts and larger dictionaries capture the specific details, all within the same latent space.

1.Why Sparse AutoEncoders?

Sparse AutoEncoders is a technique originally used in Signal Processing. It has evolved, from dictionary learning methods, and has recently proven to be effective in helping us understand the features neural networks learn from the high-dimensional activations of LLMs.

In the early days of interpretability, researchers discovered that vision models often learn more features than they have neurons. This phenomenon is known as polysemanticity or superposition.This made mechanistic interpretability work very difficult and limited its progress for a few years. However, recent work, notably by Anthropic, has popularized the use of SAEs to disentangle these features. This allows us to extract monosemantic features and trace the circuits of concepts learned by LLMs.

Still, this technique is evolving. They are interesting and promising but they are currently suffering from some limitations that require better scaling solutions for it to be used in practical interpretability work.

2. Limitations of Classic SAEs

In the way SAEs are trained on model's activations, there is a fundamental problem. The more we increase the dictionary size, that is, the total number of features the SAE can learn, the more "sparsity" encourages the model to either split general features into fragmented sub-features or absorb them into broader, less useful clusters.

In this regard, the advised way of patching this problem is to sweep through many different dictionary sizes to find the best. But obviously, this is computationally expensive and, clearly, a "try and fail" approach that would not scale to larger, and more modern models.

The paper identifies three important problems with current SAEs: feature splitting, feature absorption and feature composition.

Feature Splitting is when a high-level feature, which could be important for interpretability, gets broken down into more specific features. As a consequence, the more general representation gets missing. In the original paper, the authors took the example of a sentence of punctuation. Like when a more general feature like "punctuation mark" gets splitted into {Question Mark, Comma, Period}, and we lose the more general feature "punctuation mark" in the cost of the 3 more specific features, which by the way are related to it.

Feature absorption is when that part of a general concept gets splitted and leaves a hole into that more general feature.

Feature Composition is when we have different independent features, and the sparsity objective encourages the model to learn latent that capture combination, even if the features are not forcefully related, ,instead of representing features independently, which would be more useful.

3. What are Matryoshka SAEs?

Matryoshka SAEs were introduced to solve this. They are inspired by "Matryoshka Representation Learning" (MRL), named after the Russian nesting dolls.

MRL is a representation learning method where information is encoded at different levels of abstraction within the same embedding vector. This main advantage allows the representation to be adoptable to computational constraints of downstreams tasks.

MRL is a representation learning method where information is encoded at different levels of abstraction within the same embedding vector. This main advantage allows the representation to be adoptable to computational constraints of downstreams tasks.

They come up with an approach where the training is done using multiple dictionaries of increasing size simultaneously. This forces the smaller dictionaries to reconstruct the inputs at their best extent without using larger dictionaries. On this, it organizes features hierarchically, so

Larger features learn more specific concepts without merging larger ones into them.

4. Architecture and Training Objective

At the core, MSAEs introduces a "nested" training objective.

Same as the classic SAEs, the encoder maps the hidden state into a latent representation , and the decoder reconstructs the input .

Decoder1. Sigma here is an activation function that introduces sparsity in by enforcing non-negativity, typically the ReLU function.2. The decoder uses a linear transformation to map the latents back to the original input, which produces an approximation reconstruction of the input .

Matryoshka SAEs extend this classic encoder decoder architecture by training multiple nested auto-encoders of increasing sizes, and introduces a nested objective function for the training.

To enforce hierarchy of features, they introduced the idea of nested dictionaries. In their methods, we are given a maximum dictionary size . Then from , we define a sequence of nested dictionaries of sizes = such that

In the training process, each represents the dictionary of a sub-SAE that should reconstruct the input the best it could only using latents at most.

In practice, in their experiments, is sampled from a batch. This leads to the novel nested training objective as follows.

In this formula, the term forces the model to reconstruct the input using only the first m latents, ensuring the hierarchy is strictly maintained.

5. Practical Experimentation

The paper used synthetic data, on a 4-layer toy tiny-stories model, with a tree-structured model to validate the statistical dependencies. In their validation experiments, they used a synthetic model where they constructed a tree structure of binary features.

In practice, in the experimental pipeline, they used a 4-layer model trained on tiny stories dataset, to validate the statistical dependencies. In their validation experiments, they used a synthetic model where they constructed a tree structure of binary features.

Every feature is a node, and the root feature is always active.

Children's features are sampled only if their parents are active, which creates the hierarchical statistical dependency.

For active features, their direction vector gets added to a final running sum to produce a final d-dim input that is fed into the autocender.

Here, the first latents must learn to reconstruct well alone.The vanilla gets one loss signal, but here it gets several nested loss signals at different prefix length. This forces the most important features to be encoded first.

The Ks are sampled from a Preto distribution. This means at each step, the mode is trained on a small prefix.

6. Results and Replication