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

TL;DR

Gemma Scope computionally democratizes interpretability research by open-sourcing a large, JumpReLU SAE suite trained on Gemma 2 across multiple scales. The work thoroughly characterizes sparsity-fidelity trade-offs, sequence-position effects, width-related phenomena, and cross-model transfer to IT variants, while detailing the substantial infrastructure required for training at this scale. By releasing thousands of latents, evaluation results, and an interactive demo, it provides a valuable resource for safety and interpretability research, including potential circuit analysis and red-teaming applications. The findings highlight both the utility and the caveats of SAEs for understanding and improving real-world LM behavior, and point to concrete open problems to guide future work.

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://www.neuronpedia.org/gemma-scope

Figures (23)

• Figure 1: Locations of sparse autoencoders inside a transformer block of Gemma 2. Note that Gemma 2 has RMS Norm at the start and end of each attention and MLP block.
• Figure 2: Sparsity-fidelity trade-off for layer 12 Gemma 2 2B and layer 20 Gemma 2 9B SAEs. An ideal SAE should have low delta loss and low L0, i.e. correspond to a point towards the bottom-left corner of each plot. For an analogous plot using FVU as the measure of fidelity see \ref{['fig:pareto_main_fvu']}.
• Figure 3: Reconstruction loss by sequence position for Gemma 2 9B middle-layer 131K-width SAEs with $\lambda=10^{-3}$.
• Figure 4: Delta loss versus sparsity curves for a series of SAEs of differing width (keeping $\lambda$ and other hyperparameters constant), trained on the residual stream after layer 20 of Gemma 2 9B.
• Figure 5: Frequency histogram of SAEs trained on Gemma 2 9B, layer 20, post MLP residual with sparsity coefficient $\lambda=6\times10^{-4}$. (These correspond to the SAEs with $\text{L0}\approx 50$ in \ref{['fig:delta-loss-width-ladder']}.)