Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2 (2024)

Abstract: Sparse autoencoders (SAEs) are an unsupervised method for learning a sparse decomposition of a neural network’s latent representations into seemingly interpretable features.Despite recent excitement about their potential, research applications outside of industry are limited by the high cost of training a comprehensive suite of SAEs.In this work, we introduce Gemma Scope, an open suite of JumpReLU SAEs trained on all layers and sub-layers of Gemma 2 2B and 9B and select layers of Gemma 2 27B base models.We primarily train SAEs on the Gemma 2 pre-trained models, but additionally release SAEs trained on instruction-tuned Gemma 2 9B for comparison.We evaluate the quality of each SAE on standard metrics and release these results.We hope that by releasing these SAE weights, we can help make more ambitious safety and interpretability research easier for the community. Weights and a tutorial can be found at https://huggingface.co/google/gemma-scope and an interactive demo can be found at https://neuronpedia.org/gemma-scope .

## introduction there are several lines of evidence that suggest that a significant fraction of the internal activations of language models are sparse, linear combination of vectors, each corresponding to meaningful features @xcite @xcite @xcite . but by default, it is difficult to identify which vectors are meaningful, or which meaningful vectors are present. sparse autoencoders are a promising unsupervised approach to do this, and have been shown to often find causally relevant, interpretable directions @xcite @xcite @xcite . if this approach succeeds it could help unlock many of the hoped for applications of interpretability @xcite @xcite , such as detecting and fixing hallucinations, being able to reliably explain and debug unexpected model behaviour and preventing deception or manipulation from autonomous ai agents. however, sparse autoencoders are still an immature technique, and there are many open problems to be resolved @xcite before these downstream uses can be unlocked -especially validating or red-teaming saes as an approach, learning how to measure their performance, learning how to train saes at scale efficiently and well, and exploring how saes can be productively applied to real-world tasks. as a result, there is an urgent need for further research, both in industry and in the broader community. however, unlike previous interpretability techniques like steering vectors @xcite or probing @xcite , sparse autoencoders can be highly expensive and difficult to train, limiting the ambition of interpretability research. though there has been a lot of excellent work with sparse autoencoders on smaller models @xcite @xcite , the works that use saes on more modern models have normally focused on residual stream saes at a single layer @xcite @xcite . in addition, many of these @xcite have been trained on proprietary models which makes it more challenging for the community at large to build on this work. to address this we have trained and released the weights of gemma scope: a comprehensive, open suite of jumprelu saes @xcite on every layer and sublayer of gemma 2 2b and 9b (gemma team, 2024a), 1 as well as select layers of the larger 27b model in this series. we release these weights under a permissive cc-by-4.0 license 2 on huggingface to enable and accelerate research by other members of the research community. gemma scope was a significant engineering challenge to train. it contains more than 400 sparse autoencoders in the main release 3 , with more than 30 million learned features in total (though many features likely overlap), trained on 4-16b tokens of text each. we used over 20% of the training compute of @xcite , saved about 20 pebibytes (pib) of activations to disk, and produced hundreds of billions of sparse autoencoder parameters in total. this was made more challenging by our decision to make a comprehensive suite of saes, on every layer and sublayer. we believe that a comprehensive suite is essential for enabling more ambitious applications of interpretability, such as circuit analysis @xcite @xcite , essentially scaling up @xcite to larger models, which may be necessary to answer mysteries about larger models like what happens during chain of thought or in-context learning. in section 2 we provide background on saes in general and jumprelu saes in particular. section 3 contains details of our training procedure, hyperparameters and computational infrastructure. we run extensive evaluations on the trained saes in section 4. ## evaluation in this section we evaluate the trained saes from various different angles. we note however that as of now there is no consensus on what constitutes a reliable metric for the quality of a sparse autoencoder or its learned latents and that this is an ongoing area of research and debate @xcite @xcite . unless otherwise noted all evaluations are on sequences from the same distribution as the sae training data, i.e. the pretraining distribution of gemma 1 (gemma team, 2024b). ## discussion and future work in this report we have introduced gemma scope, a comprehensive suite of sparse autoencoders (saes) on all layers and sublayers of gemma 2 2b and 9b pt. we have described the engineering challenges involved in this project and how we approached them. in order to shed light on the quality of the gemma scope saes, we have provided results of various evaluation experiments. while we have extensively evaluated these saes, their real test is how much they can enable and accelerate downstream interpretability research. to further underscore this point, we provide a broad range of open research questions related to saes which we think are enabled or aided by gemma scope in appendix a and which we would be excited to see pursued by the interpretability research community. training such a comprehensive suite of sparse autoencoders requires a significant upfront cost in compute and energy (section 3) and thus also has a certain carbon footprint. it is our hope that by paying this cost once, we can avoid the broader research community having to train their own models again and again. we think gemma scope will enable research into language model internals for years to come, even if and when the state of the art of sae training improves in the future, and so we are optimistic that the cost of training gemma scope can be amortized.