Replicating Matryoshka Sparse Autoencoders

A guided replication of "Learning Multi-Level Features with Matryoshka Sparse Autoencoders" (Bussmann et. al (2025))(arXiv:2503.17547)

Read the full write-up, the story behind this code, on LessWrong...Feature Hierarchies Without the Absorption: How Matryoshka SAEs recover hierarchical features

The paper's central claim is that when features are hierarchically organized, a vanilla SAE absorbs and splits them, while a Matryoshka SAE, with its nested prefix dictionaries in objective function, recovers the hierarchy cleanly.

The paper's official repo is noanabeshima/matryoshka-saes. The LLM-scale pipeline here was rebuilt from the paper on top of the shared sae.py.

Layout

Path	What
sae.py	The Matryoshka SAE implementation (n_prefixes=1 recovers a vanilla SAE)
toy_model.py	Synthetic hierarchical feature generator
tree.json, tree.schema.json	The feature hierarchy definition and schema
train_toy.py	Headless toy-model training and evaluation driver
modal_tinystories.py	Modal app: cache activations, then train SAEs on a TinyStories transformer
heatmap.py	Plotly heatmap helper for the figures
figures/	Output heatmaps (.png / interactive .html) + toy checkpoint

Results

In short, on the toy model the Matryoshka SAE cleanly recovers 20/20 ground-truth features versus 11/20 for vanilla; on TinyStories the reconstruction/sparsity trade-off holds at matched sparsity, with both SAEs healthy.