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

Authors: tom  lieberum
senthooran  rajamanoharan
arthur  conmy
lewis  smith
nicolas  sonnerat
vikrant  varma
jános  kramár
anca  dragan
rohin  shah
neel  nanda
Venue: EMNLP_2024_Workshop_BlackBoxNLP
Published: 2024-08-16 12:20:22

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 \url{https://huggingface.co/google/gemma-scope} and an interactive demo can be found at \url{https://neuronpedia.org/gemma-scope}.

Main paper: 
 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 (Elhage et al., 2022; Gurnee et al., 2023; Olah et al., 2020; Park et al., 2023; Nanda et al., 2023a; Mikolov et al., 2013) . 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 (Bricken et al., 2023; Cunningham et al., 2023; Templeton et al., 2024; Gao et al., 2024; Marks et al., 2024) . If this approach succeeds it could help unlock many of the hoped for applications of interpretability (Nanda, 2022; Olah, 2021; Hubinger, 2022) , 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 (Templeton et al., 2024 ) 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 (Turner et al., 2024; Li et al., 2023) or probing (Belinkov, 2022) , 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 (Bricken et al., 2023; Cunningham et al., 2023; Dunefsky et al., 2024; Marks et al., 2024) , the works that use SAEs on more modern models have normally focused on residual stream SAEs at a single layer (Templeton et al., 2024; Gao et al., 2024; Engels et al., 2024) . In addition, many of these (Templeton et al., 2024; Gao et al., 2024) 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 (Rajamanoharan et al., 2024b) 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 GPT-3 (Brown et al., 2020) , 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 (Conmy et al., 2023; Wang et al., 2022; Hanna et al., 2023) , essentially scaling up Marks et al. (2024) 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. 
sparse autoencoders: 
 Given activations x ∈ R n from a language model, a sparse autoencoder (SAE) decomposes and reconstructs the activations using a pair of encoder and decoder functions (f , x) defined by: f (x) := σ (W enc x + b enc ) ,(1) x(f ) := W dec f + b dec .(2)These functions are trained to map x(f (x)) back to x, making them an autoencoder. Thus, f (x) ∈ R M is a set of linear weights that specify how to combine the M ≫ n columns of W dec to reproduce x.The columns of W dec , which we denote by d i for i = 1 . . . M , represent the dictionary of directions into which the SAE decomposes x. We will refer to to these learned directions as latents to disambiguate between learnt 'features' and the conceptual features which are hypothesized to comprise the language model's representation vectors. 4The decomposition f (x) is made non-negative and sparse through the choice of activation function σ and appropriate regularization, such that f (x) typically has much fewer than n non-zero entries. Initial work (Cunningham et al., 2023; Bricken et al., 2023) used a ReLU activation function to enforce non-negativity, and an L1 penalty on the decomposition f (x) to encourage sparsity. TopK SAEs (Gao et al., 2024) enforce sparsity by zeroing all but the top K entries of f (x), whereas the JumpReLU SAEs (Rajamanoharan et al., 2024b) enforce sparsity by zeroing out all entries of f (x) below a positive threshold. Both TopK and JumpReLU SAEs allow for greater separation between the tasks of determining which latents are active, and estimating their magnitudes.In this work we focus on JumpReLU SAEs as they have been shown to be a slight Pareto improvement over other approaches, and allow for a variable number of active latents at different tokens (unlike TopK SAEs).The JumpReLU activation is a shifted Heaviside step function as a gating mechanism together with a conventional ReLU:σ(z) = JumpReLU θ (z) := z ⊙ H(z -θ). (3)Here, θ > 0 is the JumpReLU's vector-valued learnable threshold parameter, ⊙ denotes elementwise multiplication, and H is the Heaviside step function, which is 1 if its input is positive and 0 otherwise. Intuitively, the JumpReLU leaves the pre-activations unchanged above the threshold, but sets them to zero below the threshold, with a different learned threshold per latent.As loss function we use a squared error reconstruction loss, and directly regularize the number of active (non-zero) latents using the L0 penalty:EQUATIONwhere λ is the sparsity penalty coefficient. Since the L0 penalty and JumpReLU activation function are piecewise constant with respect to threshold parameters θ, we use straight-through estimators (STEs) to train θ, using the approach described in Rajamanoharan et al. (2024b) . This introduces an additional hyperparameter, the kernel density estimator bandwidth ε, which controls the quality of the gradient estimates used to train the threshold parameters θ.5 3 Training detailsacross layers, e.g. are there persistent features that one can "match up" across adjacent layers? 3. Better understanding the phenomenon of "feature splitting" (Bricken et al., 2023) where high-level features in a small SAE break apart into several finer-grained features in a wider SAE. 4. We know that SAEs introduce error, and completely miss out on some features that are captured by wider SAEs (Templeton et al., 2024; Bussmann et al., 2024) . Can we quantify and easily measure "how much" they miss and how much this matters in practice? 5. How are circuits connecting up superposed features represented in the weights? How do models deal with the interference between features (Nanda et al., 2023b) ?Using SAEs to improve performance on realworld tasks (compared to fair baselines) 1. Detecting or fixing jailbreaks. 2. Helping find new jailbreaks/red-teaming models (Ziegler et al., 2022) . 3. Comparing steering vectors (Turner et al., 2024) to SAE feature steering (Conmy and Nanda, 2024) or clamping (Templeton et al., 2024) . 4. Can SAEs be used to improve interpretability techniques, like steering vectors, such as by removing irrelevant features (Conmy and Nanda, 2024) ?Red-teaming SAEs 1. Do SAEs really find the "true" concepts in a model? 2. How robust are claims about the interpretability of SAE features (Huang et al., 2023) ? 3. Can we find computable, quantitative measures that are a useful proxy for how "interpretable" humans think a feature vector is (Bills et al., 2023) ? 4. Can we find the "dark matter" of truly nonlinear features?11 5. Do SAEs learn spurious compositions of independent features to improve sparsity as has been shown to happen in toy models (Anders et al., 2024) , and can we fix this?Scalable circuit analysis: What interesting circuits can we find in these models?1. What's the learned algorithm for addition (Stolfo et al., 2023) in Gemma 2 2B? 2. How can we practically extend the SAE feature circuit finding algorithm in Marks et al. (2024) to larger models? 3. Can we use SAE-like techniques such as MLP transcoders (Dunefsky et al., 2024) to find input independent, weights-based circuits?Using SAEs as a tool to answer existing questions in interpretability 1. What does finetuning do to a model's internals (Jain et al., 2024) ? 2. What is actually going on when a model uses chain of thought? 3. Is in-context learning true learning, or just promoting existing circuits (Hendel et al., 2023; Todd et al., 2024) ? 4. Can we find any "macroscopic structure" in language models, e.g. families of features that work together to perform specialised roles, like organs in biological organisms?12 5. Does attention head superposition (Jermyn et al., 2023) As described in Section 3, during training, we normalize LM activations and subtract b dec from them before passing them to the encoder. However, after training, we reparameterize the Gemma Scope SAEs so that neither of these steps are required during inference. Let x raw be the raw LM activations that we rescale by a scalar constant C, i.e. x := x raw /C, such that E ∥x∥ 2 2 = 1. Then, as parameterized during training, the SAE forward pass is defined byEQUATIONEQUATIONIt is straightforward to show that by defining the following rescaled and shifted parameters:b ′ enc := C b enc -C W enc b dec (7) b ′ dec := C b dec (8) θ ′ := C θ (9)we can simplify the SAE forward pass (operating on the raw activations x raw ) as follows:EQUATION)EQUATIONMLP SAEs are trained on the output of MLPs, but we can also replace the whole MLP with a transcoder (Dunefsky et al., 2024) for easier circuit analysis. Transcoders are not autoencoders: while SAEs are trained to reconstruct their input, transcoders are trained to approximate the output of MLP layers from the input of the MLP layer. We train one suite of transcoders on Gemma 2B PT, and release these at https://huggingface.co/ google/gemma-scope-2b-pt-transcoders.Evaluation We find that transcoders cause a greater increase in loss to the base model relative to the MLP output SAEs (Fig. 6 ), at a fixed sparsity (L0). This reverses the trend from GPT-2 Small found by Dunefsky et al. (2024) . This could be due to a number of factors, such as: 1. (2024) also do not use resampling (Bricken et al., 2023) or an architecture which prevents dead features like more recent SAE research (Rajamanoharan et al., 2024a; Conerly et al., 2024; Gao et al., 2024) , which means their results are in a fairly different setting to other SAE research.We fold the pre-MLP RMS norm gain parameters (Zhang and Sennrich (2019) , Section 3) into the MLP input matrices, as described in (Gurnee et al. (2024) , Appendix A.1) and then train the transcoder on input activations just after the pre-MLP RMSNorm, to reconstruct the MLP sublayer's output as the target activations. To make it easier for Gemma Scope users to apply this change, in Fig. 7 we provide TransformerLens code for loading Gemma 2 2B with this weight folding applied. Fig. 7 also includes an explanation of why only a subset of the weight folding techniques described in Appendix A.1 of Gurnee et al. (2024) can be applied to Gemma 2, due to its architecture.We train transcoders identically to MLP SAEs except for the following two differences: 1. We do not initialize the encoder kernel W enc to the transpose of the decoder kernel W dec ; 2. We do not use a pre-encoder bias, i.e. we do not subtract b dec from the input to the transcoder (although we still add b dec at the transcoder output). These two training changes were motivated by the fact that, unlike SAEs, the input and outputs spaces for transcoders are not identical. To spell out how we apply normalization: we divide the input and target activations by the root mean square of the input activations. Since the input activations all have norm √ d model due to RMSNorm, this means we divide input and output activations by √ d model .D.1 Sparsity-fidelity tradeoff Fig. 11 illustrates the trade off between fidelity as measured by fraction of variance unexplained (FVU) against sparsity for layer 12 Gemma 2 2B and layer 20 Gemma 2 9B SAEs. Fig. 12 shows the sparsity-fidelity trade off for the 131K-width residual stream SAEs trained on Gemma 2 27B after layers 10, 22 and 34 that we include as part of this release.Fig. 15 and Fig. 16 show fidelity versus sparsity curves for more layers (approximately evenly spaced) and all sites of Gemma 2 2B and Gemma 2 9B, demonstrating consistent and smoothly variance performance throughout these models.Methodology Prior research has shown that language models tend to have lower loss on later token positions (Olsson et al., 2022) . It is thus natural to ask how an SAE's performance changes over the length of a sequence. Similar to Section 4.1, we track reconstruction loss and delta loss for various sparsity settings, however this time we do not aggregate over the sequence position. Again, we mask out special tokens.Result Fig. 8 shows how reconstruction loss varies by position for 131K-width SAEs trained on the middle-layer of Gemma 2 9B. Reconstruction loss increases rapidly from close to zero over the first few tokens. The loss monotonically increases by position for attention SAEs, although it is essentially flat after 100 tokens. For MLP SAEs, the loss peaks at around the tenth token before gradually declining slightly. We speculate that this is because attention is most useful when tracking long-range dependencies in text, which matters most when there is significant prior context to draw from, while MLP layers do a lot of local processing, like detecting n-grams (Gurnee et al., 2023) , that does not need much context. Like attention SAEs, residual stream SAEs' loss monotonically increases, plateauing more gradually. Curves for other models, layers, widths and sparsity coefficients were found to be qualitatively similar. Fig. 13 shows how delta LM loss varies by sequence position. The high level of noise in the delta loss measurements makes it difficult to robustly measure the effect of position, however there are signs that this too is slightly lower for the first few tokens, particularly for residual stream SAEs.Methodology We perform the sparsity-fidelity evaluation from Section 4.1 on different validation subsets of The Pile (Gao et al., 2020) , to gauge whether SAEs struggle with a particular type of data. 13Results In Fig. 9 we show delta loss by subset. Of the studied subsets, SAEs perform best on Deep-Mind mathematics (Saxton et al., 2019) . Possibly this is due to the especially formulaic nature of the data. SAEs perform worst on Europarl (Koehn, 2005) , a multilingual dataset. We conjecture that this is due to the Gemma 1 pretraining data, which was used to train the SAEs, containing predominantly English text.We train all SAEs in 32-bit floating point precision. It is common to make model inference less memory and compute intensive by reducing the precision at inference time. This is particularly important for applications like circuit analysis, where users may wish to splice several SAEs into a language model simultaneously. Fig. 10 compares fidelity-versussparsity curves computed using float32 SAE and LM weights versus the same curves computed using bfloat16 SAE and LM weights, suggesting there is negligible impact switching bfloat16 inference.Methodology Conventionally, sparsity of SAE latent activations is measured as the L0 norm of the latent activations. Olah et al. (2024) suggest to train SAEs to have low L1 activation of attributionweighted latent activations, taking into account that some latents may be more important than others. We repurpose their loss function as a metric for our SAEs, which were trained penalising activation sparsity as normal. As in Rajamanoharan et al. (2024b) , we define the attribution-weighted latent activation vector y asEQUATIONwhere we choose the mean-centered logit of the correct next token as the loss function L.We then normalize the magnitudes of the entries of y to obtain a probability distribution p ≡ p(y). We can measure how far this distribution diverges from a uniform distribution u over active latents via the KL divergence with 1 ∥y∥ 0 ≤ r L0 ≤ 1. Note that since e S can be interpreted as the effective number of active elements, r L0 is the ratio of the effective number of active latents (after re-weighting) to the total number of active latents, which we call the 'Uniformity of Active Latent Importance'.EQUATIONIn Fig. 14 we show r L0 on middle layer SAEs. In line with Rajamanoharan et al. (2024b) , we find that the attributed effect becomes more diffuse as more latents are active. This effect is most pronounced for residual stream SAEs, and seems to be independent of language model size and number of SAE latents.In this sub-appendix, we provide further evaluations of SAEs on the activations of IT models, continuing Section 4.4.As mentioned in Section 4.4, we find in Fig. 19 that PT SAEs achieve reasonable FVU on rollouts, but the gap between PT and IT SAEs is larger than in the change in loss in the main text (Fig. 5 ).In Fig. 17 we evaluate the FVU on the user prompt and model prefix (not the rollout). In Fig. 18 we evaluate the change in loss (delta loss) on the user prompts, and surprisingly find that splicing in the base model SAE can reduce the loss in expectation in some cases. Our explanation for this result is that post-training does not train models to predict user queries (only predict high-preference model rollouts) and therefore the model is not incentivised to have good predictive loss by default on the user prompt.While we do not train IT SAEs on Gemma 2 2B, we find that the base SAEs transfer well as measured by FVU in Fig. 20 .Finally, we do not find evidence that rescaling IT activations to have same norm in expectation to the pretraining activations is beneficial (Fig. The trend for individual SAEs in this plot is that their L0 decreases but the Pareto frontier is very slightly worse. This is consistent with prior observations that SAEs are surprisingly adaptable to different L0s (Smith, 2024; Gao et al., 2024) . We also release one suite of transcoders(Dunefsky et al.(2024); Appendix C), a 'feature-splitting' suite of SAEs with multiple widths trained on the same site (Section 4.2), and some SAEs trained on the Gemma 2 9B IT model(Kissane et al. (2024b); Section 4.4).2 Note that the Gemma 2 models are released under a different, custom license.3 For each model, layer and site we in fact release multiple SAEs with differing levels of sparsity; taking this into account, we release the weights of over 2,000 SAEs in total.This is different terminology from earlier work(Bricken et al., 2023; Rajamanoharan et al., 2024a,b), where feature is normally used interchangeably for both SAE latents and the language models featuresA large value of ε results in biased but low variance estimates, leading to SAEs with good sparsity but sub-optimal fidelity, whereas a low value of ε results in high variance estimates that cause the threshold to fail to train at all, resulting in SAEs that fail to be sparse. We find through hyperparameter sweeps across multiple layers and sites that ε = 0.001 provides a good trade-off (when SAE inputs are normalized to have an unit mean squared norm) and use this to train the SAEs released as part of GemmaScope. 6 This is similar in spirit toConerly et al. (2024), who normalize the dataset to have mean norm of √ dmodel.The residual stream is the bottleneck by which the previous layers communicate with all later layers. Any given MLP or attention layer adds to the residual stream, and is typically only a small fraction of the residual, so even a large error in the layer is a small error to the residual stream and so to theNote the 1M-width SAEs included in Fig.2do not form part of this suite as they were trained using a different range of values for the sparsity coefficient λ.See e.g. https://huggingface.co/google/ gemma-2-2b-it for the user and model prefixes.https://huggingface.co/google/ gemma-scope-9b-it-resWe distinguish truly non-linear features from low-rank subspaces of linear features as found inEngels et al. (2024).We know this happens in image models(Voss et al., 2021) but have not seen much evidence in language models. But superposition is incentivized for features that do not co-occur(Gurnee et al., 2023), so specialized macroscopic structure may be a prime candidate to have in superposition. Now we have SAEs, can we find and recover it?Note that this is a different dataset to the dataset used to train the Gemma Scope SAEs. 
data: 
 We train SAEs on the activations of Gemma 2 models generated using text data from the same distribution as the pretraining text data for Gemma 1 (Gemma Team, 2024b) , except for the one suite of SAEs trained on the instruction-tuned (IT) model (Section 4.4). We generate activations on sequences of length 1024.For a given sequence we only collect activations from tokens which are neither BOS, EOS, nor padding. After activations have been generated, they are shuffled in buckets of about 106 activations. We speculate that a perfect shuffle would not significantly improve results, but this was not systematically checked. We would welcome further investigation into this topic in future work.During training, activation vectors are normalized by a fixed scalar to have unit mean squared norm. 6 This allows more reliable transfer of hyperparameters (in particular the sparsity coefficient λ and bandwidth ε) between layers and sites, as the raw activation norms can vary over multiple orders of magnitude, changing the scale of the reconstruction loss in Eq. ( 4). Once training is complete, we rescale the trained SAE parameters so that no Location We train SAEs on three locations per layer, as indicated by Fig. 1 . We train on the attention head outputs before the final linear transformation W O and RMSNorm has been applied (Kissane et al., 2024a) , on the MLP outputs after the RMSNorm has been applied and on the post MLP residual stream. For the attention output SAEs, we concatenate the outputs of the individual attention heads and learn a joint SAE for the full set of heads. We zero-index the layers, so layer 0 refers to the first transformer block after the embedding layer. In Appendix C we define transcoders (Dunefsky et al., 2024) and train one suite of these.Optimization We use the same bandwidth ε = 0.001 and learning rate η = 7 × 10 -5 across all training runs. We use a cosine learning rate warmup from 0.1η to η over the first 1,000 training steps. We train with the Adam optimizer (Kingma and Ba, 2017) with (β 1 , β 2 ) = (0, 0.999), ϵ = 10 -8 and a batch size of 4,096. We use a linear warmup for the sparsity coefficient from 0 to λ over the first 10,000 training steps.During training, we parameterise the SAE using a pre-encoder bias (Bricken et al., 2023) , subtracting b dec from activations before the encoder. However, after training is complete, we fold this bias into the encoder parameters, so that no preencoder bias needs to be applied during inference. See Appendix B for details.Throughout training, we restrict the columns of W dec to have unit norm by renormalizing after every update. We also project out the part of the gradients parallel to these columns before computing the Adam update, as described in Bricken et al. (2023) .Initialization We initialize the JumpReLU threshold as the vector θ = {0.001} M . We initialize W dec using He-uniform (He et al., 2015) initialization and rescale each latent vector to be unit norm. W enc is initalized as the transpose of W dec , but they are not tied afterwards (Conerly et al., 2024; Gao et al., 2024) . The biases b dec and b enc are initialized to zero vectors.Topology We train most of our SAEs using TPUv3 in a 4x2 configuration. Some SAEs, especially the most wide ones, were trained using TPUv5p in either a 2x2x1 or 2x2x4 configuration.Sharding We train SAEs with 16.4K latents with maximum amount of data parallelism, while using maximal amounts of tensor parallelism using Megatron sharding (Shoeybi et al., 2020) for all other configurations. We find that as one goes to small SAEs and correspondingly small update step time, the time spent on host-to-device (H2D) transfers outgrows the time spent on the update step, favoring data sharding. For larger SAEs on the other hand, larger batch sizes enable higher arithmetic intensity by reducing transfers between HBM and VMEM of the TPU. Furthermore, the specific archi-tecture of SAEs means that when using Megatron sharding, device-to-device (D2D) communication is minimal, while data parallelism causes a costly all-reduce of the full gradients. Thus we recommend choosing the smallest degree of data sharding such that the H2D transfer takes slightly less time than the update step.As an example, with proper step time optimization this enables us to process one batch for a 131Kwidth SAE in 45ms on 8 TPUv3 chips, i.e. a model FLOP utilization (MFU) of about 50.8%.Disk storage We store all collected activations on hard drives as raw bytes in shards of 10-20GiB. We use 32-bit precision in all our experiments. This means that storing 8B worth of activations for a single site and layer takes about 100TiB for Gemma 2 9B, or about 17PiB for all sites and layers of both Gemma 2 2B and 9B. The total amount is somewhat higher still, as we train some SAEs for 16B tokens and also train some SAEs on Gemma 2 27B, as well as having a generous buffer of additional tokens. While this is a significant amount of disk space, it is still cheaper than regenerating the data every time one wishes to train an SAE on it. Concretely, in our calculations we find that storing activations for 10-100 days is typically at least an order of magnitude cheaper than regenerating them one additional time. The exact numbers depend on the model used and the specifics of the infrastructure, but we expect this relationship to hold true in general. If there is a hard limit on the amount of disk space available, however, or if fast disk I/O can not be provided (see next paragraph), then this will favor on-the-fly generation instead. This would also be the case if the exact hyperparameter combinations were known in advance. In practice, we find it advantageous for research iteration speed to be able to sweep sparsity independently from other hyperparameters and to retrain SAEs at will. EQUATIONTable 1 : Overview of the SAEs that were trained for which sites and layers. For each model, width, site and layer, we release multiple SAEs with differing levels of sparsity (L0).All+: We also train one suite of transcoders on the MLP sublayers on Gemma 2.6B PT (Appendix C).tuning.To overcome this bottleneck we implement a shared server system, enabling us to amortize disk reads for a single site and layer combination:  
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 (Gao et al., 2024; Karvonen et al., 2024; Makelov et al., 2024) .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).Methodology For a fixed dictionary size, we trained SAEs of varying levels of sparsity by sweeping the sparsity coefficient λ. We then plot curves showing the level of reconstruction fidelity attainable at a given level of sparsity.Metrics We use the mean L0-norm of latent activations, E x ∥f (x)∥ 0 , as a measure of sparsity. To measure reconstruction fidelity, we use two metrics:• Our primary metric is delta LM loss, the increase in the cross-entropy loss experienced by the LM when we splice the SAE into the LM's forward pass. • As a secondary metric, we also use fraction of variance unexplained (FVU) -also called the normalized loss (Gao et al., 2024) -as a measure of reconstruction fidelity. This is the mean reconstruction loss L reconstruct of a SAE normalized by the reconstruction loss obtained by always predicting the dataset mean. Note that FVU is purely a measure of the SAE's ability to reconstruction the input activations, not taking into account the causal effect of any error on the downstream loss. All metrics were computed on 2,048 sequences of length 1,024, after masking out BOS, EOS, and padding tokens when aggregating the results.Results Fig. 2 compares the sparsity-fidelity trade-off for SAEs in the middle of each Gemma model. For the full results see Appendix D. Delta loss is consistently higher for residual stream SAEs compared to MLP and attention SAEs, whereas FVU (Fig. 11 ) is roughly comparable across sites. We conjecture this is because even small errors in reconstructing the residual stream can have a significant impact on LM loss, whereas relatively large errors in reconstructing a single MLP or attention sub-layer's output have a limited impact on the LM loss.7 Holding all else equal, wide SAEs learn more latent directions and provide better reconstruction fidelity at a given level of sparsity than narrow SAEs. Intuitively, this suggests that we should use the widest SAEs practicable for downstream tasks. However, there are also signs that this intuition may come with caveats. The phenomenon of 'featuresplitting' (Bricken et al., 2023) -where latents in a narrow SAE seem to split into multiple specialized latents within wider SAEs -is one sign that wide SAEs do not always use their extra capacity to learn a greater breadth of features (Bussmann et al., 2024) . It is plausible that the sparsity penalty used to train SAEs encourages wide SAEs to learn frequent compositions of existing features instead of or in addition to learning new features (Anders et al., 2024) . If this is the case, it is currently unclear whether this is good or bad for the usefulness of SAEs on downstream tasks.In order to facilitate research into how SAEs' properties vary with width, and in particular how SAEs with different widths trained on the same data relate to one another, we train and release a 'feature-splitting' suite of mid-network residual stream SAEs for Gemma 2 2B and 9B PT with matching sparsity coefficients and widths between 2 14 ≈ 16.4K and 2 19 ≈ 524K in steps of powers of rest of the network's processing. On the other hand, a large error to the residual stream itself will significantly harm loss. For the same reason, mean ablating the residual stream has far higher impact on the loss than mean ablating a single layer. two. 8 The SAEs are trained with different sparsity settings after layers 12 and 20 of Gemma 2 2B and 9B respectively.Sparsity-fidelity trade-off Similar to Section 4.1, Fig. 3 compares fidelity-versus-sparsity curves for SAEs of differing width in this ladder.Latent firing frequency Fig. 4 shows frequency histograms for λ = 6 × 10 -4 SAEs in the same ladder of widths from 2 14 to 2 19 latents. To compute these, we calculate the firing frequency of each latent over 20,000 sequences of length 1,024, masking out special tokens. The mode and most of the mass shifts towards lower frequencies with increased number of latents. However there remains a cluster of ultra-high frequency latents, which has also been observed for TopK SAEs but not for Gated SAEs (Cunningham and Conerly, 2024; Gao et al., 2024; Rajamanoharan et al., 2024b) .The interpretability of latents for JumpReLU SAEs and other architectures was investigated in Rajamanoharan et al. (2024b) , finding little difference between various SAE architectures. Since we also use JumpReLU SAEs, we refer to section 5.3 of that work for a detailed discussion of the methodology and results.Additional IT SAE training Prior research has shown that SAEs trained on base model activations also faithfully reconstruct the activations of IT models derived from these base models (Kissane et al., 2024b) . We find further evidence for these results by comparing the Gemma Scope SAEs with several SAEs we train on the activations from Gemma 2B 9B IT. Specifically, we train these IT SAEs by taking the same pretraining documents used for all other SAEs (Section 3.1) and prepend them with Gemma's IT prefix for the user's query, and append Gemma's IT prefix for the model's response. 9 We then train each SAE to reconstruct activations at all token positions besides the user prefix (since these tokens have much larger norm (Kissane et al., 2024b) , and are the same for every document). We also release the weights for these SAEs to enable further research into the differences between training SAEs on base and IT models.10 Methodology To evaluate the SAEs trained on the IT model's activations, we generate 1,024 rollouts of the Gemma 2 9B IT model on a random sample of the SFT data used to train Gemini 1.0 Ultra (GoogleDeepmind, 2024) , with temperature 1.0.We then use SAEs trained on the residual stream of the base model and the IT model to reconstruct these activations, and measured the FVU.In Fig. 5 we show that using PT model SAEs results in increases in cross-entropy loss almost as small as the increase from the SAEs directly trained on the IT model's activation. We show further evaluations such as on Gemma 2 2B, measuring FVU rather than loss, and using activations from the user query (not just the rollout) in Appendix D.6. In Fig. 19 we find that the FVU for the PT SAEs is somewhat faithful, but does not paint as strong a picture as Fig. 5 . A speculative explanation for this is that finetuning consists of 'reweighting' old features from the base model, in addition to learning some new, chat-specific features that do not have as big an impact on next-token prediction. This would mean the FVU looks worse than the increase in loss since the FVU would be impacted by low impact chat features, but change in loss would not be.Future work could look into finetuning these SAEs on chat interactions if even lower reconstruction error is desired (Kissane et al., 2024b) , or evaluating on multi-turn and targeted rollouts.In Appendix D, we present additional evaluation results covering• Sparsity-Fidelity trade-off for more SAEs • Studying SAE performance as a function of token position.• SAE performance on different subsets of The Pile (Gao et al., 2020) , showing stronger performance on Deepmind Mathematics (Saxton et al., 2019) and weaker performance on Europarl (Koehn, 2005) .• Impact of low precision inference, showing little performance regression from using bfloat16.• Uniformity of active latent importance, which is a measure for how diffuse the downstream effect of a single SAE latent is, introduced by Rajamanoharan et al. (2024b).• Additional evaluation results for SAEs trained on the activations of Gemma 2 IT models. 
related work: 
 Open Weights Sparse Autoencoders There have been several open weights SAE contributions by the research community. However, all releases we are aware of have focused on smaller and older language models or have not released a comprehensive set of autoencoder weights. Marks and Mueller (2023) trained SAEs on the MLP outputs of all layers of Pythia-70M (Biderman et al., 2023) . Braun et al. (2024) trained different variants of SAEs on GPT-2 small (Radford et al., 2019) and Tinystories-1M (Eldan and Li, 2023) on the residual stream activations in select layers. Belrose (2024) In contrast to the above work, Gemma Scope is the first release of SAE weights which contains SAEs for all layers and sublayers of a recently released, performant 2B and 9B language model. 
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.Our main goal in releasing Gemma Scope is to help the broader safety and interpretability communities advance our understanding of interpretability, and how it can be used to make models safer. As a starting point, we provide a list of open problems we would be particularly excited to see progress on, where we believe Gemma Scope may be able to help. Where possible we cite work that may be a helpful starting point, even if it is not tackling exactly the same question.Deepening our understanding of SAEs 1. Exploring the structure and relationships between SAE features, as suggested in Wattenberg and Viégas (2024). 
 
 Appendix:  


Decision: Reviewers appreciated the contribution of a massive suite of SAEs trained at large scale on SOTA models, as well as the technical details. One reviewer raised a concern about environmental costs in case the particular type of SAEs becomes obsolete. It would be good to discuss this in the camera ready.

Reviews: review: paper_summary: paper_summary: This technical report describes the training procedures used to train a comprehensive suite of JumpReLU SAEs on Gemma 2 (Instruct). Specifically, the authors trained multiple suites: a comprehensive suite of SAEs on every layer and sublayer of Gemma 9B and 2B, a selection of layers of Gemma 27B, a suite of transcoders, a suite of SAEs with multiple widths, and a selection of layers on Gemma 2 9B Instruct (over 2,000 SAEs in total). Finally, the authors evaluate and compare the quality of their SAEs on standard metrics, and propose a set of research questions this suite of SAEs could help to address.
strength_weakness: strength_weakness: NA
questions: questions: - The reference to Nora Belrose's SAEs is titled "Untitled" (p. 9).
paper_summary: paper_summary: This paper aims to accelerate interpretability research using sparse autoencoders (SAE), which suffer from high training cost and a lack of openly released SAEs for modern models. To support future research, authors release SAEs of varying sparsity for each (sub)layer of Gemma 2 9B and 2B, as well as selected layers of Gemma 2 27B. 

The paper goes into detail about training details and engineering challenges addressed during this project, and has some evaluation results, although the main success criterion is to what extent future interpretability research will be able to benefit from the SAEs released.
strength_weakness: strength_weakness: - massive scale is not only a strength, but brings environmental costs. I hope the resources put into this don't go to waste, for example because the particular flavour of SAE used becomes obsolete. Authors say that they used "over 20% of the training compute of GPT-3", but applications of SAEs are much more narrow than the LLMs themselves.
questions: questions: -