This work was done as part of MATS 7.0.

MATS provides a generous compute stipend, and towards the end of the program we found we had some unspent compute. To let this not go to waste, we trained batch topk Matryoshka SAEs on all residual stream layers of Gemma-2-2b, Gemma-2-9b, and Gemma-3-1b, and are now releasing them publicly. The hyperparams for these SAEs were not aggressively optimized, but they should hopefully be decent. Below we describe our rationale for how these SAEs were trained and why, and the stats for each SAE. Key decisions:

• We use more narrow inner widths than in the original Matryoshka SAEs work, and increase each width by a larger amount. We do this to make it easier to study the highest-frequency features of the model.
• We include standard and snap loss variants for Gemma-2-2b. Snap loss and the rationale behind it is described in our Feature Hedging post. There is probably not much practical difference between the snap and standard versions of the SAEs.
• We do not stop gradients between Matryoshka layers. We find in toy models that hedging and absorption pull the encoder in opposite directions, and this helps moderate the severity of feature hedging in Matryoshka SAE inner layers.

I don't care about any of that, just give me the SAEs!

You can load all of the SAEs using SAELens via the following releases:

• Gemma-2-2b snap loss matryoshka: gemma-2-2b-res-snap-matryoshka-dc
• Gemma-2-2b standard matryoshka: gemma-2-2b-res-matryoshka-dc
• Gemma-2-9b standard matryoshka: gemma-2-9b-res-matryoshka-dc
• Gemma-3-1b standard matroshka: gemma-3-1b-res-matryoshka-dc

For each release, the SAE ID is just the corresponding Transformer Lens post residual stream hook point, e.g. blocks.5.hook_resid_post for the layer 5 residual stream SAE.

Each SAE can be loaded in SAELens as follows:

from sae_lens import SAE

sae = SAE.from_pretrained("<release>", "<sae_id>")[0]

For instance, to load the layer 5 snap variant SAE for gemma-2-2b, this would look like the following:

sae = SAE.from_pretrained("gemma-2-2b-res-snap-matryoshka-dc", "blocks.5.hook_resid_post")[0]

Neuronpedia

Neuronpedia has generously hosted some of these SAEs, with more coming soon. Check them out at: https://www.neuronpedia.org/res-matryoshka-dc.

Matryoshka SAEs should be much better than standard SAEs at finding general, high-frequency concepts like parts of speech. In standard SAEs, latents tracking these concepts will get shot to pieces by feature absorption as they co-occur with so many other concepts. As Matryoshka SAEs should be much more resilient to absorption, we thus expect to find more meaningful high-density latents in Matryoshka SAEs (although these latents may be messed up by feature hedging instead). For instance, here's a high-density first-layer latent from layer 12 of Gemma-2-2b, which appears to (very noisily) perform a grammatical function similar to Treebank's IN (Preposition or subordinating conjunction) part of speech. 

Higher frequency concepts should be concentrated in earlier latent indices. The highest frequency concepts should be in latents 0-127, then the next highest frequency should be in latents 128-511, etc...

Training info

All Matryoshka SAEs in this release are trained on 750M tokens from the Pile using a modified version of SAELens. The SAEs are all 32k width with the following Matryoshka levels: 128, 512, 2048, 8192, and 32768. We included two layers (128, 512) that are much narrower than the model residual stream to make it easier to study what the first  features are that the SAE learns. These are all batch top-k SAEs, following the original Matryoshka SAEs work. We largely did not optimize hyperparams for these SAEs, so it's likely possible to squeeze out more performance from the SAE with optimized choices of learning rate and more training tokens, but hopefully these SAEs should be decent.

Snap loss

One of the notable components of this release is the addition of snap loss variants of all SAEs for Gemma-2-2b. Snap loss is described in our post on Feature Hedging, and involves switching the reconstruction loss of the SAE from MSE to L2 mid-way through training. Practically, we don't see much difference in SAEs trained on LLMs using snap loss, but are releasing these regardless in case others are curious to investigate the effect of snap loss, as we have the SAE trained anyway. If you notice a meaningful difference in practice between the snap loss and standard variants of these SAEs, please let us know!

Balancing absorption and hedging

Intuitively, it might seem like we'd want the inner layers of Matryoshka SAEs to be insulated from gradients from outer layers. Outer layers will pull the inner latents towards absorption, which defeats the purpose of a Matryoshka SAE! However, in toy models, hedging and absorption have opposite effects on the SAE encoder, so allowing some absorption pressure can help counteract the hedging of the SAE and improve performance. We notice that the dictionary_learning implementation of matryoshka SAEs also does not stop gradients between layers, and likely this is because stopping gradients causes hedging to mess up the SAE more severly.

For a further investigation of balancing hedging and absorption in Matryoshka SAEs, check out this colab.

We suspect that it may be possible to intentionally balance hedging with absorption in a more optimal way, and we plan to investigate this in future work. 

SAEs and stats

Below we list all the SAEs trained along with some core stats.

Gemma-2-2b

We trained both snap and standard variants of SAEs for Gemma-2-2b. These SAEs have the relase ID gemma-2-2b-res-snap-matryoshka-dc for snap-loss variant, and gemma-2-2b-res-matryoshka-dc for the standard variant.

Snap Matryoshka SAEs

Standard Matryoshka SAEs

Gemma-2-9b

These SAEs have the release ID gemma-2-9b-res-matryoshka-dc.

Gemma-3-1b

These SAEs have the release ID gemma-3-1b-res-matryoshka-dc.