Source-linked AI summary
Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2
Tom Lieberum, Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Nicolas Sonnerat, Vikrant Varma, János Kramár, Anca Dragan, Rohin Shah, Neel Nanda
TL;DR
Interpretability research needs broader access to sparse autoencoders, whose training cost and unresolved evaluation questions limit community use. Gemma Scope releases a comprehensive suite of JumpReLU SAEs across Gemma 2 models and evaluates them with standard metrics. The release also covers instruction-tuned comparison models and provides open weights for further research.
Problem
High training costs, limited coverage of modern models, and unresolved questions about SAE validation constrain broader community research.
Method
The authors train and release JumpReLU SAEs across Gemma 2 2B, 9B, and selected 27B layers, with additional 9B instruction-tuned models.
Results
The release evaluates SAE quality using delta LM loss and fraction of variance unexplained, while reporting comparisons across sites, widths, architectures, and model variants.
Takeaways & Limitations
Open Gemma Scope weights and comprehensive coverage are intended to accelerate community safety and interpretability research.
Takeaways & Limitations
The usefulness of feature-splitting from wider SAEs remains unclear for downstream tasks, and Gemma 2 uses a custom license.
Abstract
from arXiv · showhide
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
1. Introduction
Sparse autoencoders could expose meaningful features in language-model activations, but their cost and immaturity limit broader interpretability research. Gemma Scope addresses this gap by releasing a comprehensive open suite of JumpReLU SAEs across Gemma 2 models.
- Sparse autoencoders offer an unsupervised way to identify interpretable, potentially causally relevant directions in language-model activations.
- SAE research remains constrained by expensive training, incomplete coverage of modern models, and reliance on proprietary systems.
- Gemma Scope releases JumpReLU SAEs for every layer and sublayer of Gemma 2 2B and 9B, plus selected layers of Gemma 2 27B.
- The main release contains more than 400 SAEs and over 30 million learned features, trained on 4-16B tokens per SAE.
- The authors argue that comprehensive layer and sublayer coverage is essential for ambitious interpretability applications such as circuit analysis.
- The release includes SAEs trained on Gemma 2 9B instruction-tuned models, alongside additional transcoders and feature-splitting suites.
2. Preliminaries
Sparse autoencoders encode language-model activations into sparse, nonnegative latent coefficients and reconstruct them from an overcomplete dictionary. Gemma Scope focuses on JumpReLU gating, which uses learned thresholds to control latent activation sparsity.
- Sparse autoencoders: An SAE maps an activation vector x into a higher-dimensional latent representation and reconstructs x through encoder and decoder functions.
- Sparse autoencoders: Decoder columns form a dictionary of latent directions, while encoder weights determine how those directions combine to reproduce the input.
- Sparse autoencoders: SAE latents are nonnegative and sparse because the activation function and regularization suppress most coefficients.
- JumpReLU SAEs: JumpReLU is used because it offers a slight Pareto improvement and permits different numbers of active latents across tokens.
- JumpReLU SAEs: JumpReLU leaves pre-activations unchanged above each learned positive threshold and sets those below threshold to zero.
- JumpReLU SAEs: Training uses squared-error reconstruction with an L0 penalty, and straight-through estimators optimize the piecewise-constant threshold parameters.
3. Training details
Gemma Scope trains SAEs on multiple locations throughout Gemma 2 layers using standardized data processing, normalization, and optimization choices. The training infrastructure combines large activation storage with shared, distributed data loading to support extensive SAE sweeps.
- Data and locations: Training activations come from Gemma pretraining text, exclude BOS, EOS, and padding tokens, and are shuffled in buckets of about 10^6 activations.The dataset also includes a separate suite trained on instruction-tuned Gemma 2 9B data.
- Optimization: ε = 0.001 is selected as a trade-off between sparse but biased thresholds and unstable thresholds that fail to produce sparse SAEs.The choice is based on hyperparameter sweeps across multiple layers and sites with normalized inputs.
- Normalization: Activation vectors are normalized to unit mean squared norm during training, then SAE parameters are rescaled so inference requires no input normalization.This normalization makes hyperparameters such as λ and ε more transferable across layers and sites.
- Optimization: The release varies SAE width and training duration: 16.4K-latent SAEs use 4B tokens, 1M-width SAEs use 16B tokens, and other SAEs use 8B tokens.Multiple sparsity levels are released for each model, width, site, and layer.
- Data and locations: SAEs are trained at three locations per layer: attention head outputs, post-RMSNorm MLP outputs, and the post-MLP residual stream.Attention outputs are concatenated across heads before learning a joint SAE.
- Data pipeline: Shared buffers, distributed disk reads, and dynamic fetching reduce data-loading bottlenecks when training multiple SAEs on the same site and layer.A single Gemma 2 9B SAE requires more than 1 GiB/s of disk-read speed to avoid being bottlenecked by data loading.
4. Evaluation
Gemma Scope evaluations examine sparsity–fidelity trade-offs, sequence-position effects, width, interpretability, transfer to instruction-tuned models, and inference precision. Results show site-dependent loss sensitivity, improved fidelity from wider SAEs, comparable interpretability across architectures, near-transfer from base-model SAEs, and negligible bfloat16 impact.
- Evaluation setup: SAE evaluation uses mean L0 for sparsity, delta LM loss as the primary fidelity metric, and FVU as a secondary reconstruction metric.Delta LM loss measures the cross-entropy increase caused by splicing an SAE into the model; FVU normalizes reconstruction loss against predicting the dataset mean.
- Sparsity–fidelity trade-off: Residual-stream SAEs have consistently higher delta loss than MLP and attention SAEs, although FVU is roughly comparable across sites.The authors attribute this to the residual stream’s role as the communication bottleneck for later layers.
- Sequence position: Reconstruction loss rises rapidly over the first tokens; attention and residual-stream losses then increase or plateau, while MLP loss peaks near token ten and declines slightly.These patterns were qualitatively similar across other models, layers, widths, and sparsity coefficients.
- SAE width: Wider SAEs learn more latent directions and provide better reconstruction fidelity at a fixed sparsity level, but feature-splitting may limit breadth gains.The paper notes that wider SAEs can split narrow-SAE latents into multiple specialized latents, and the usefulness of this remains unclear.
- Transfer to instruction-tuned models: Using pretraining-model SAEs on instruction-tuned activations causes cross-entropy increases almost as small as using SAEs trained directly on instruction-tuned activations.The corresponding FVU results are somewhat faithful but less strong than the loss comparison.
- Inference precision: bfloat16 SAE and language-model weights have negligible impact on fidelity-versus-sparsity curves relative to float32 weights.This comparison targets lower-memory and lower-compute inference for applications that may splice several SAEs simultaneously.
5. Open problems that Gemma Scope may help tackle
Gemma Scope is intended as a foundation for broader safety and interpretability research. The authors identify open problems spanning feature structure, feature splitting, safety applications, interpretability validation, circuit analysis, model algorithms, fine-tuning, and quantization.
- Research scope: The release is intended to help safety and interpretability communities advance understanding of interpretability and its use in making models safer.The authors present the open problems as starting points where Gemma Scope may help research progress.
- Feature structure: Open questions include discovering relationships among SAE features and matching persistent residual-stream features across adjacent layers.These questions target the structure and cross-layer organization of learned features.
- Feature splitting: The authors highlight feature splitting, in which high-level features in small SAEs break into finer-grained features in wider SAEs.They also ask how much wider SAEs recover features missed by narrower ones and how this matters in practice.
- Safety applications: Potential safety applications include detecting or fixing jailbreaks, red-teaming models, and comparing SAE feature steering with steering vectors or clamping.The list also asks whether irrelevant features can be removed to improve interpretability techniques.
- Interpretability: Open interpretability questions ask whether SAEs recover true concepts, how robust latent-interpretability claims are, and whether computable proxies can match human judgments.The authors also point to the unresolved possibility of genuinely nonlinear features.
- Circuits and model behavior: Further problems concern feature compositions, circuit representations, learned algorithms, fine-tuning, chain of thought, in-context learning, and macroscopic feature organization.The proposed questions extend from model internals to practical circuit-finding methods for larger models.
- Scaling and efficiency: The authors ask how SAEs can capture circular features and cross-layer superposition, and how much quantization is possible without significant inference or training degradation.These questions address representational structure and deployment efficiency.
A. Standardizing SAE parameters for inference
Gemma Scope normalizes language-model activations during SAE training, then reparameterizes the trained models so inference can operate directly on raw activations.
- During training, LM activations are normalized and bdec is subtracted before encoding.
- After training, Gemma Scope reparameterizes the SAEs so inference requires neither activation normalization nor subtraction of bdec.
- Raw activations xraw are rescaled by a scalar C to form x := xraw/C with unit expected mean squared norm.
- The training-time SAE forward pass is transformed by rescaling and shifting parameters.
B. Transcoders
The paper releases Gemma 2B transcoders for circuit analysis and evaluates them against MLP-output SAEs. Transcoders increase base-model loss more at fixed sparsity, with several possible explanations discussed.
- Transcoders approximate MLP outputs from MLP inputs, whereas SAEs reconstruct the activations they receive.
- One suite of transcoders is trained on Gemma 2B pre-trained models and released for research.
- At fixed sparsity (L0), transcoders cause a greater increase in base-model loss than MLP-output SAEs.
- The observed trend reverses the result reported for GPT-2 Small by Dunefsky et al. (2024).
- Possible explanations include limited scaling to larger or modern architectures, a larger JumpReLU benefit for SAEs, implementation errors, and differing training details.
- Transcoder training differs from MLP SAE training in encoder initialization, pre-encoder bias handling, and normalization of input and target activations.
C.1. Sparsity-fidelity tradeoff
The paper reports fidelity-versus-sparsity evaluations across representative Gemma 2 SAE configurations and broader model coverage, including Gemma 2 27B release layers.
- Fig. 13 compares fraction of variance unexplained (FVU) with sparsity for layer 12 Gemma 2 2B and layer 20 Gemma 2 9B SAEs.
- Fig. 14 shows the sparsity-fidelity tradeoff for 131K-width residual-stream SAEs from Gemma 2 27B layers 10, 22, and 34.
- Figs. 17 and 18 show fidelity-versus-sparsity curves for approximately evenly spaced layers and all sites of Gemma 2 2B and 9B.
- These broader curves demonstrate consistent and smoothly varying performance throughout the evaluated Gemma 2 2B and 9B models.
C.2. Impact of sequence position
The paper examines how delta loss varies with sequence position.
- Fig. 15 shows how delta loss varies by position.
C.3. Uniformity of active latent importance
This section introduces Uniformity of Active Latent Importance, which measures how evenly attribution-weighted importance is distributed among active latents. On middle-layer SAEs, attributed effects become more diffuse as more latents are active, especially for residual-stream SAEs.
- Methodology: The metric uses the mean-centered logit of the correct next token as the loss function for attribution weighting.
- Uniformity of Active Latent Importance is the ratio of the effective number of active latents after re-weighting to the total number of active latents.It is derived from the entropy of a probability distribution over attribution-weighted active latents.
- Results: The attributed effect becomes more diffuse as more latents are active in middle-layer SAEs.
- Results: This diffusion is most pronounced for residual stream SAEs and appears independent of language-model size and the number of SAE latents.
C.4. Additional Gemma 2 IT evaluation results
The additional evaluations test how Gemma Scope SAEs trained on pretraining or instruction-tuned activations reconstruct instruction-tuned model activations. Pretraining SAEs transfer reasonably in several settings, but instruction-tuned SAEs perform better on rollouts and activation rescaling is not beneficial.
- Rollouts: Pretraining SAEs achieve reasonable FVU on instruction-tuned rollouts, but their gap from instruction-tuned SAEs is larger for loss change.
- User prompts: On user prompts, splicing in base-model SAEs can reduce loss in expectation in some cases.The authors attribute this to post-training optimizing prediction of high-preference model rollouts rather than user queries.
- Cross-model transfer: Base SAEs transfer well to Gemma 2 2B instruction-tuned activations when measured by FVU.
- Activation scaling: Rescaling instruction-tuned activations to match pretraining norms shows no evidence of benefit.For individual SAEs, L0 decreases while the Pareto frontier becomes very slightly worse.
- Evaluation scope: The evaluations compare FVU and change in loss across model sizes, layers, sequence positions, SAE widths, and training sources.