Identifying Functionally Important Features with End-to-End Sparse Dictionary Learning

Authors: dan  braun
jordan  taylor
nicholas  goldowsky-dill
lee  sharkey
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 11:35:33

Abstract: Identifying the features learned by neural networks is a core challenge in mechanistic interpretability. Sparse autoencoders (SAEs), which learn a sparse, overcomplete dictionary that reconstructs a network's internal activations, have been used to identify these features. However, SAEs may learn more about the structure of the datatset than the computational structure of the network. There is therefore only indirect reason to believe that the directions found in these dictionaries are functionally important to the network. We propose end-to-end (e2e) sparse dictionary learning, a method for training SAEs that ensures the features learned are functionally important by minimizing the KL divergence between the output distributions of the original model and the model with SAE activations inserted. Compared to standard SAEs, e2e SAEs offer a Pareto improvement: They explain more network performance, require fewer total features, and require fewer simultaneously active features per datapoint, all with no cost to interpretability. We explore geometric and qualitative differences between e2e SAE features and standard SAE features. E2e dictionary learning brings us closer to methods that can explain network behavior concisely and accurately. We release our library for training e2e SAEs and reproducing our analysis at
https://github.com/ApolloResearch/e2e_sae.

Main paper: 
 introduction: 
 Sparse Autoencoders (SAEs) are a popular method in mechanistic interpretability [Sharkey et al., 2022 , Cunningham et al., 2023 , Bricken et al., 2023] . They have been proposed as a solution to the problem of superposition, the phenomenon by which networks represent more 'features' than they have neurons. 'Features' are directions in neural activation space that are considered to be the basic units of computation in neural networks. SAE dictionary elements (or 'SAE features') are thought to approximate the features used by the network. SAEs are typically trained to reconstruct the activations of an individual layer of a neural network using a sparsely activating, overcomplete set of dictionary elements (directions). It has been shown that this procedure identifies ground truth features in toy models [Sharkey et al., 2022] .However, current SAEs focus on the wrong goal: They are trained to minimize mean squared reconstruction error (MSE) of activations (in addition to minimizing their sparsity penalty). The issue is that the importance of a feature as measured by its effect on MSE may not strongly correlate with how important the feature is for explaining the network's performance. This would not be a problem if the network's activations used a small, finite set of ground truth features -the SAE would simply identify those features, and thus optimizing MSE would have led the SAE to learn the functionally Bottom: Pareto curves for three different types of GPT2-small layer 6 SAEs as the sparsity coefficient is varied. E2e-SAEs require fewer features per datapoint (i.e. have a lower L0) and fewer features over the entire dataset (i.e. have a low number of alive dictionary elements). GPT2-small has a CE loss of 3.139 over our evaluation set.important features. In practice, however, Bricken et al. [2023] observed the phenomenon of feature splitting, where increasing dictionary size while increasing sparsity allows SAEs to split a feature into multiple, more specific features, representing smaller and smaller portions of the dataset. In the limit of large dictionary size, it would be possible to represent each individual datapoint as its own dictionary element. Since minimizing MSE does not explicitly prioritize learning features based on how important they are for explaining the network's performance, an SAE may waste much of its fixed capacity on learning less important features. This is perhaps responsible for the observation that, when measuring the causal effects of circuits made from SAE features on network performance, a significant amount is mediated by the reconstruction residual errors (i.e. everything not explained by the SAE) and not mediated by SAE features [Marks et al., 2024] .Given these issues, it is therefore natural to ask how we can identify the functionally important features used by the network. We say a feature is functionally important if it is important for explaining the network's behavior on the training distribution. If we prioritize learning functionally important features, we should be able to maintain strong performance with fewer features used by the SAE per datapoint as well as fewer overall features.To optimize SAEs for these properties, we introduce a new training method. We still train SAEs using a sparsity penalty on the feature activations (to reduce the number of features used on each datapoint), but we no longer optimize activation reconstruction. Instead, we replace the original activations with the SAE output (Figure 1 ) and optimize the KL divergence between the original output logits and the output logits when passing the SAE output through the rest of the network, thus training the SAE end-to-end (e2e). We use SAE e2e to denote an SAE trained with KL divergence and a sparsity penalty.By contrast, we use SAE local to denote our baseline SAEs, trained only to reconstruct the activations at the current layer with a sparsity penalty.One risk with this method is that it may be possible for the outputs of SAE e2e to take a different computational pathway through subsequent layers of the network (compared with the original activations) while nevertheless producing a similar output distribution. For example, it might learn a new feature that exploits a particular transformation in a downstream layer that is unused by the regular network or that is used for other purposes. To reduce this likelihood, we also add terms to the loss for the reconstruction error between the original model and the model with the SAE at downstream layers in the network (Figure 1 ). We use SAE e2e+ds to denote SAEs trained with KL divergence, a sparsity penalty, and downstream reconstruction loss. We use e2e SAEs to refer to the family of methods introduced in this work, including both SAE e2e and SAE e2e+ds .Previous work has used the performance explained -measured by cross-entropy loss difference when replacing the original activations with SAE outputs -as a measure of SAE quality [Cunningham et al., 2023 , Bricken et al., 2023 , Bloom, 2024 ]. It's reasonable to ask whether our approach runs afoul of Goodhart's law ("When a measure becomes a target, it ceases to be a good measure"). We contend that mechanistic interpretability should prefer explanations of networks (and the components of those explanations, such as features) that explain more network performance over other explanations. Therefore, optimizing directly for quantitative proxies of performance explained (such as CE loss difference, KL divergence, and downstream reconstruction error) is preferred.We train each SAE type on language models (GPT2-small [Radford et al., 2019] and Tinystories-1M [Eldan and Li, 2023] ), and present three key findings:1. For the same level of performance explained, SAE local requires activating more than twice as many features per datapoint compared to SAE e2e+ds and SAE e2e (Section 3.1).2. SAE e2e+ds performs equally well as SAE e2e in terms of the number of features activated per datapoint (Section 3.1), yet its activations take pathways through the network that are much more similar to SAE local (Sections 3.2, 3.3).3. SAE local requires more features in total over the dataset to explain the same amount of network performance compared with SAE e2e and SAE e2e+ds (Section 3.1).These findings suggest that e2e SAEs are more efficient in capturing the essential features that contribute to the network's performance. Moreover, our automated-interpretability and qualitative analyses reveal that SAE e2e+ds features are at least as interpretable as SAE local features, demonstrating that the improvements in efficiency do not come at the cost of interpretability (Section 3.4). These gains nevertheless come at the cost of longer wall-clock time to train (Appendix H).As a supplementary investigation, we tested e2e SAEs on a set of subject-verb agreement tasks from Finlayson et al. [2021] . These tests had inconclusive results (Appendix I), indicating that further work is needed to identify which downstream tasks would benefit most from e2e SAEs.In addition to this article, we also provide: A library for training all SAE types presented in this article (https://github.com/ApolloResearch/e2e_sae); a Weights and Biases [Biewald, 2020] report that links to training metrics for all runs (https://api.wandb.ai/links/sparsify/ evnqx8t6); a Neuronpedia [Lin and Bloom, 2023] page for interacting with the features in a subset of SAEs (those presented in Tables 2 and 3 ) (https://www.neuronpedia.org/gpt2sm-apollojt) as well as a repository for downloading these SAEs directly (https://huggingface.co/ apollo-research/e2e-saes-gpt2). 
training end-to-end saes: 
 Our experiments train SAEs using three kinds of loss function (Figure 1 ), which we evaluate according to several metrics (Section 2.5):1. L local trains SAEs to reconstruct activations at a particular layer (Section 2.2); 2. L e2e trains SAEs to learn functionally important features (Section 2.3);3. L e2e+downstream trains SAEs to learn functionally important features that optimize for faithfulness to the activations of the original network at subsequent layers (Section 2.4).Suppose we have a feedforward neural network (such as a decoder-only Transformer [Radford et al., 2018] ) with L layers and vectors of hidden activations a (l) :a (0) (x) = x a (l) (x) = f (l) (a (l-1) (x)), for l = 1, . . . , L -1 y = softmax f (L) (a (L-1) (x)) .We use SAEs that consist of an encoder network (an affine transformation followed by a ReLU activation function) and a dictionary of unit norm features, represented as a matrix D, with associated bias vector b d . The encoder takes as input network activations from a particular layer l. The architecture we use is:Enc a (l) (x) = ReLU W e a (l) (x) + b e SAE a (l) (x) = D ⊤ Enc a (l) (x) + b d ,where the dictionary D and encoder weights W e are both (N_dict_elements × d_hidden) matrices, b e is a N_dict_elements-dimensional vector, while b d and a (l) (x) are d_hidden-dimensional vectors.The standard, baseline method for training SAEs is SAE local training, where the output of the SAE is trained to reconstruct its input using a mean squared error loss with a sparsity penalty on the encoder activations (here an L1 loss):L local = L reconstruction + L sparsity = ||a (l) (x) -SAE local (a (l) (x))|| 2 2 + ϕ||Enc(a (l) (x))|| 1 . ϕ = λ dim(a (l)) is a sparsity coefficient λ scaled by the size of the input to the SAE (see Appendix D for details on hyperparameters).For SAE e2e , we do not train the SAE to reconstruct activations. Instead, we replace the model activations with the output of the SAE and pass them forward through the rest of the network:â(l) (x) = SAE e2e (a (l) (x)) â(k) (x) = f (k) (â (l) (x)) for k = l, . . . , L -1 ŷ = softmax f (L) (â (L-1) (x))We train the SAE by penalizing the KL divergence between the logits produced by the model with the SAE activations and the original model:L e2e = L KL + L sparsity = KL(ŷ, y) + ϕ||Enc(a (l) (x))|| 1Importantly, we freeze the parameters of the model, so that only the SAE is trained. This contrasts with Tamkin et al. [2023] , who train the model parameters in addition to training a 'codebook' (which is similar to a dictionary).A reasonable concern with the L e2e is that the model with the SAE inserted may compute the output using an importantly different pathway through the network, even though we've frozen the original model's parameters and trained the SAE to replicate the original model's output distribution. To counteract this possibility, we also compare an additional loss: The end-to-end with downstream reconstruction training loss (L e2e+downstream ) additionally minimizes the mean squared error between the activations of the new model at downstream layers and the activations of the original model:L e2e+downstream = L KL + L sparsity + L downstream = KL(ŷ, y) + ϕ||Enc(a (l) )|| 1 + β l L -l L-1 k=l+1 ||â (k) (x) -a (k) (x)|| 2 2 (1)where β l is a hyperparameter that controls the downstream reconstruction loss term (Appendix D).L e2e+downstream thus has the desirable properties of 1) incentivizing the SAE outputs to lead to similar computations in downstream layers in the model and 2) allowing the SAE to "clear out" some of the non-functional features by not training on a reconstruction error at the layer with the SAE. Note, however, the inclusion of the intermediate reconstruction terms means that L e2e+downstream may encourage the SAE to learn features that are less functionally important.We record several key metrics for each trained SAE:1. Cross-entropy loss increase between the original model and the model with SAE: We measure the increase in cross-entropy (CE) loss caused by using activations from the inserted SAE rather than the original model activations on an evaluation set. We sometimes refer to this as 'amount of performance explained', where a low CE loss increase means more performance explained. All other things being equal, a better SAE recovers more of the original model's performance. 2. L 0 : How many SAE features activate on average for each datapoint. All other things being equal, a better SAE needs fewer features to explain the performance of the model on a given datapoint. 
number of alive dictionary elements:: 
 The number of features in training that have not 'died' (which we define to mean that they have not activated over a set of 500k tokens of data). All other things being equal, a better SAE needs a smaller number of alive features to explain the performance of model over the dataset.We also record the reconstruction loss at downstream layers. This is the mean squared error between the activations of the original model and the model with the SAE at all layers following the insertion of the SAE (i.e. downstream layers). If reconstruction loss at downstream layers is low, then the activations take a similar pathway through the network as in the original model. This minimizes the risk that the SAEs are learning features that take different computational pathways through the downstream layers compared to the original model. Finally, following Bills et al. [2023] , we perform automated-interpretability scoring and qualitative analysis on a subset of the SAEs, to verify that improved quantitative metrics does not sacrifice the interpretability of the learned features.We show results for experiments performed on GPT2-small's residual stream before attention layer 6.3 Results for layers 2, 6, and 10 of GPT2-small and some runs on a model trained on the TinyStories dataset [Eldan and Li, 2023 ] can be found in Appendices A.1 and A.2, respectively. They are qualitatively similar to those presented in the main text. For our GPT2-small experiments, we train SAEs with each type of loss function on 400k samples of context size 1024 from the Open Web Text dataset [Gokaslan and Cohen, 2019 ] over a range of sparsity coefficients λ. Our dictionary is fixed at 60 times the size of the residual stream (i.e. 60 × 768 = 46080 initial dictionary elements). Hyperparameters, along with sweeps over dictionary size and number of training examples, are shown in Appendices D and E, respectively.3.1 End-to-end SAEs are a Pareto improvement over local SAEsWe compare the trained SAEs according to CE loss increase, L 0 , and number of alive dictionary elements. The learning rates for each SAE type were selected to be Pareto-optimal according to their L 0 vs CE loss increase curves. 4 Each experiment uses a range of sparsity coefficients λ. In Figure 1 , we see that both SAE e2e and SAE e2e+ds achieve better CE loss increase for a given L 0 or for a given number of alive dictionary elements. This means they need fewer features to explain the same amount of network performance for a given datapoint or for the dataset as a whole, respectively. For similar results at other layers see Appendix A.1.This difference is large: For a given L 0 , both SAE e2e and SAE e2e+ds have a CE loss increase that is less than 45% of the CE loss increase of SAE local .5 SAE local must therefore be learning features that are not maximally important for explaining network performance.This improved performance comes at the expense of increased compute costs (2-3.5 times longer runtime, see Appendix H). We test to see if additional compute improves our SAE local baseline in Appendix E. We find neither increasing dictionary size from 60 * 768 to 100 * 768 nor increasing training samples from 400k to 800k noticeably improves the Pareto frontier, implying that our e2e SAEs maintain their advantage even when compared against SAE local dictionaries trained with more compute.For comparability, our subsequent analyses focus on 3 particular SAEs that have approximately equivalent CE loss increases (Table 1 ). 2 ). This indicates that the activations following the insertion of SAE e2e take a different path through the network than in the original model, and therefore potentially permit the model to achieve its performance using different computations from the original model. This possibility motivated the training of SAE e2e+ds s.In later layers, the reconstruction errors of SAE local and SAE e2e+ds are extremely similar (Figure 2 ). SAE e2e+ds therefore has the desirable properties of both learning features that explain approximately as much network performance as SAE e2e (Figure 1 ) while having reconstruction errors that are much closer to SAE local . There remains a difference in reconstruction at layer 6 between SAE e2e+ds and SAE local . This is not surprising given that SAE e2e+ds is not trained with a reconstruction loss at this layer. In Appendix B, we examine how much of this difference is explained by feature scaling. In Appendix G.3, we find a specific example of a direction with low functional importance that is faithfully represented in SAE local but not SAE e2e+ds .3.3 Differences in feature geometries between SAE types 3.3.1 End-to-end SAEs have more orthogonal features than SAE local Bricken et al. [2023] observed 'feature splitting', where a locally trained SAEs learns a cluster of features which represent similar categories of inputs and have dictionary elements pointing in similar directions. A key question is to what extent these subtle distinctions are functionally important for the Figure 2 : Reconstruction mean squared error (MSE) at later layers for our set of GPT2-small layer 6 SAEs with similar CE loss increases (Table 1 ). SAElocal is trained to minimize MSE at layer 6, SAEe2e was trained to match the output probability distribution, SAEe2e+ds was trained to match the output probability distribution and minimize MSE in all downstream layers.network's predictions, or if they are only helpful for reconstructing functionally unimportant patterns in the data.We have already seen that SAE e2e and SAE e2e+ds learn smaller dictionaries compared with SAE local for a given level of performance explained (Figure 1 ). In this section, we explore if this is due to less feature splitting. We measure the cosine similarities between each SAE dictionary feature and nextclosest feature in the same dictionary. While this does not account for potential semantic differences between directions with high cosine similarities, it serves as a useful proxy for feature splitting, since split features tend to be highly similar directions [Bricken et al., 2023] . We find that SAE local has features that are more tightly clustered, suggesting higher feature splitting (Figure 3a ). Compared to SAE e2e+ds the mean cosine similarity is 0.04 higher (bootstrapped 95% CI [0.037-0.043]); compared to SAE e2e the difference is 0.166 (95% CI [0.163 -0.168]). We measure this for all runs in our Pareto frontiers and find that this difference is not explained by SAE local having more alive dictionary elements than e2e SAEs (Appendix A.5).We find that SAE local s trained with one seed learn similar features as SAE local s trained with a different seed (Figure 3b ). The same is true for two SAE e2e+ds s. However, features learned by SAE e2e are quite different for different seeds. This suggests there are many different sets of SAE e2e features that achieve the same output distribution, despite taking different paths through the network.The cosine similarity plots between SAE e2e and SAE local (Figure 3c Using the automated-interpretability library [Lin, 2024] (an adaptation of Bills et al. [2023] ), we generate automated explanations of our SAE features by prompting gpt-4-turbo-2024-04-09 [OpenAI et al., 2024] with five max-activating examples for each feature, before generating "interpretability scores" by tasking gpt-3.5-turbo to use that explanation to predict the SAE feature's true activations on a random sample of 20 max-activating examples. For each SAE we generate automated-interpretabilty scores for a random sample of features (n = 198 to 201 per SAE). We then measure the difference between average interpretability scores. This interpretability score is an Figure 3 : Geometric comparisons for our set of GPT2-small layer 6 SAEs with similar CE loss increases (Table 1 ). For each dictionary element, we find the max cosine similarity between itself and all other dictionary elements.In 3a we compare to others directions in the same SAE, in 3b to directions in an SAE of the same type trained with a different random seed, in 3c to directions in the SAElocal with similar CE loss increase.imperfect metric of interpretability, but it serves as an unbiased verification and is therefore useful for ensuring that we are not trading better training losses for significantly less interpretable features.For pairs of SAEs with similar L 0 (listed in Table 3 ), we find no difference between the average interpretability scores of SAE local and SAE e2e+ds . If we repeat the analysis for pairs with similar CE loss increases, we find the SAE e2e+ds features to be more interpretable than SAE local features in Layers 2 (p = 0.0053) and 6 (p = 0.0005) but no significant difference in layer 10. For additional automated-interpretability analysis, see Appendix A.7.We also provide some qualitative, human-generated interpretations of some groups of features for different SAE types in Appendix G. Features from the SAEs in Table 2 and Table 3 can be viewed interactively at https://www.neuronpedia.org/gpt2sm-apollojt.4 Related work 
using sparse autoencoders and sparse coding in mechanistic interpretability: 
 When Elhage et al. [2022] identified superposition as a key bottleneck to progress in mechanistic interpretability, the field found a promising scalable solution in SAEs [Sharkey et al., 2022] . SAEs have since been used to interpret language models [Cunningham et al., 2023 , Bricken et al., 2023 , Bloom, 2024] and have been used to improve performance of classifiers on downstream tasks [Marks et al., 2024] . Earlier work by Yun et al. [2021] concatenated together the residual stream of a language model and used sparse coding to identify an undercomplete set of sparse 'factors' that spanned multiple layers. This echoes even earlier work that applied sparse coding to word embeddings and found sparse linear structure [Faruqui et al., 2015 , Subramanian et al., 2017 , Arora et al., 2018] . Similar to our work is Tamkin et al. [2023] , who trained sparse feature codebooks, which are similar to SAEs, and trained them end-to-end. However, to achieve adequate performance, they needed to train the model parameters alongside the sparse codebooks. Here, we only trained the SAEs and left the interpreted model unchanged.Although useful for mechanistic interpretability, current SAE approaches have several shortcomings. One issue is the functional importance of features, which we have aimed to address here. Some work has noted problems with SAEs, including Anders and Bloom [2024] , who found that SAE features trained on a language model with a given context length failed to generalize to activations collected from activations in longer contexts. Other work has addressed 'feature suppression' [Wright and Sharkey, 2024] , also known as 'shrinkage' [Jermyn et al., 2024] , where SAE feature activations systematically undershoot the 'true' activation value because of the sparsity penalty. While Wright and Sharkey [2024] approached this problem using finetuning after SAE training, Jermyn et al. [2024] and Riggs and Brinkmann [2024] explored alternative sparsity penalties during training that aimed to reduce feature suppression (with mixed success). Farrell [2024] , taking an approach similar to Jermyn et al. [2024] , has explored different sparsity penalties, though here not to address shrinkage, but instead to optimize for other metrics of SAE quality. Rajamanoharan et al. [2024] introduce Gated SAEs, an architectural variation for the encoder which both addresses shrinkage and improves on the Pareto frontier of L 0 vs CE loss increase.One of the main challenges in using SAEs for mechanistic interpretability is that there is no known 'ground truth' against which to benchmark the features learned by SAEs. Prior to our work, several metrics have been used, including: Comparison with ground truth features in toy data; activation reconstruction loss; L 1 loss; number of alive dictionary elements; similarity of SAE features across different seeds and dictionary sizes [Sharkey et al., 2022] ; L 0 ; KL divergence (between the output distributions of the original model and the model with SAE activations) upon causal interventions on the SAE features [Cunningham et al., 2023] ; reconstructed negative log likelihood of the model with SAE activations inserted [Cunningham et al., 2023 , Bricken et al., 2023] ; feature interpretability Cunningham et al. [2023] (as measured by automatic interpretability methods [Bills et al., 2023] ); and task-specific comparisons [Makelov et al., 2024] . In our work, we use (1) L 0 , (2) number of alive dictionary elements, (3) the average KL divergence between the output distribution of the original model and the model with SAE activations, and (4) the reconstruction error of activations in layers that follow the layer where we replace the original model's activations with the SAE activations.In our work, we optimize for functional importance directly, but previous work measured functional importance post hoc using different approaches. Cunningham et al. [2023] used activation patching [Vig et al., 2020] , a form of causal mediation analysis, where they intervened on feature activations and found the output distribution was more sensitive (had higher KL divergence with the original model's distribution) in the direction of SAE features than other directions, such as PCA directions.With the same motivation, Marks et al. [2024] use a similar, approximate, but more efficient, method of causal mediation analysis [Nanda, 2022 , Sundararajan et al., 2017] . Unlike our work, these works use the measures of functional importance to construct circuits of sparse features. Bricken et al.[2023] used logit attribution, measuring the effect the feature has on the output logits. 
conclusion: 
 In this work, we introduce end-to-end dictionary learning as a method for training SAEs to identify functionally important features in neural networks. By optimizing SAEs to minimize the KL divergence between the output distributions of the original model and the model with SAE activations inserted, we demonstrate that e2e SAEs learn features that better explain network performance compared to the standard locally trained SAEs.Our experiments on GPT2-small and Tinystories-1M reveal several key findings. First, for a given level of performance explained, e2e SAEs require activating significantly fewer features per datapoint and fewer total features over the entire dataset. Second, SAE e2e+ds , which has additional loss terms for the reconstruction errors at downstream layers in the model, achieves a similar performance explained to SAE e2e while maintaining activations that follow similar pathways through later layers compared to the original model. Third, the improved efficiency of e2e SAEs does not come at the cost of interpretability, as measured by automated-interpretability scores and qualitative analysis.These results suggest that standard, locally trained SAEs are capturing information about dataset structure that is not maximally useful for explaining the algorithm implemented by the network. By directly optimizing for functional importance, e2e SAEs offer a more targeted approach to identifying the essential features that contribute to a network's performance. 
impact statement: 
 This article proposes an improvement to methods used in mechanistic interpretability. Mechanistic interpretability, and interpretability broadly, promises to let us understand the inner workings of neural networks. This may be useful for debugging and improving issues with neural networks. For instance, it may enable the evaluation of a model's fairness or bias. Interpretability may relatedly be useful for improving the trust-worthiness of AI systems, potentially enabling AI's use in certain high stakes settings, such as healthcare, finance, and justice. However, increasing the trust-worthiness of AI systems may be dual use in that may also enable its use in settings such as military applications. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: We provide statistical significance testing and report confidence intervals across our geometrics comparisons and auto-interpretability. We also provide some results showing the robustness to seeds (Appendix F), though we do not run statistical significance testing on pareto differences between our methods over various seeds due to computational feasibility.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Details on compute resources given at the end of Appendix D and in Appendix H.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We conform to all aspects of the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: See our impact statement (Section 6).Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We do not release data or models that have a high risk of misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The sources and licenses of all dataset, model and core libraries used are given in Appendix D.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Pareto curves for SAEs at other layers We also tested our methods on Tinystories-1M, a 1M parameter model trained on short, simple stories [Eldan and Li, 2023] . Figure 5 shows our key results generalising to the residual stream halfway through the model (before the 5 th of 8 layers).Note that most of our Tinystories-1M runs were for SAE local and SAE e2e , and we did not perform several of the analyses that we performed for GPT2-small elsewhere in this report. But the clear improvement in L 0 and alive_dict_elements vs CE loss increase was apparent for SAE e2e vs SAE local . More results can be found at https://api.wandb.ai/links/sparsify/yk5etolk. Future work would test that these results hold on more models of different sizes and architectures, as well as on SAEs trained not just on the residual stream. A.3 Comparison of runs with similar L 0 , CE loss increase, or number of alive dictionary elements Figure 6 shows that layers 2 and 10 also have the property that SAE e2e+ds has a very similar reconstruction loss to SAE local at downstreams layers, and SAE e2e has a much higher reconstruction loss.In Section 3.3.1 we showed that at layer 6, SAE local is less orthogonal than SAE e2e and SAE e2e+ds , indicating a higher level of feature splitting. In Figure 7 we extend the analysis to runs on other layers.In almost all cases we find that SAE e2e contains the most orthogonal dictionaries, followed by SAE e2e+ds and then SAE local . Perhaps surprisingly, as the number of alive dictionary elements decrease for each SAE type, we see an increase in the mean of the within-SAE similarities, indicating less feature splitting. One hypothesis for this result is that the the orthogonality of the dictionary depends much more on the output performance (as measured by CE loss difference) or sparsity (as measured by L 0 ) of the model with the SAE than on the number of alive dictionary elements, though further analysis is needed.A.6 Cross-type similarity at other layers In Section 3.3.2 we show that downstream and local SAEs have more similar decoder directions than e2e and local SAEs. In Figure 8 we show this is true for layers 2, 6, and 10.Figure 7 : Mean over all SAE dictionary elements of the cosine similarity to the next-closest element in the same dictionary. Plotted against L0, CE loss increase, and number of alive dictionary elements for all SAE types on runs with a variety of sparsity coefficients for GPT2-small Figure 8 : For runs with similar CE loss increase in layers 2, 6, 10, for each SAEe2e and SAEe2e+ds dictionary direction, we take the max cosine similarity over all SAElocal directions.In Section 3.4 we claim that when comparing auto-interpretability scores we find no difference between pairs of similar L 0 , but do find SAE e2e+ds is more interpretable than SAE local in layers 2 and 6. These results are presented in more detail in Figure 9 and Table 4 . 3 ) and similar CE loss increases (see Table 2 ). Error bars are a bootstraped 95% confidence interval for the true mean auto-interpretability scores. Measured on approximately 200(±2) randomly selected features per dictionary.Table 4 : Estimates of the difference between the mean auto-interpretability scores for SAEe2e+ds and SAElocal (Figure 9 ). A positive difference indicates SAEe2e+ds is more interpretable. For each comparison we use bootstrapping to compute a 95% confidence interval and a two-tailed p-value that the means are equal. We saw in Appendix A.4 that our e2e-trained SAEs are much worse at reconstructing the exact activation compared to locally-trained SAEs. We performed some initial analysis of why this is.A common problem with SAEs is "feature-supression", where the SAE output has considerably smaller norm than the input [Wright and Sharkey, 2024, Rajamanoharan et al., 2024] . We observe this as well, as shown in Figure 10 for an SAE e2e+ds in layer 6. Note the cluster of activations with original norm around 3000; these are the activations at position 0. We can measure suppression with the metric:L 2 Ratio = E x∈D ||â(x)|| ||a(x)||which is presented in Table 5 for all of the similar CE loss increase SAEs in Table 2 . Generally, SAE e2e has the most feature-suppression. This is as layer-norm is applied to the residual stream before the activations are used, which can allow the network to re-normalize the downscaled activations and keep similar outputs. The downscaled activations will still disrupt the normal ratio between the residual stream before the SAE is applied and the outputs of future layers that are added to the residual stream.Both SAE e2e and SAE e2e+ds do significantly worse at reconstructing the directions of the original activations than SAE local (Figure 11 ). Note, however, that we are comparing runs with similar CE loss increases. SAE local is the only one of the three that is trained directly on reconstructing these activations, and achieves this reconstruction with significantly higher average L 0 .Overall, SAE e2e+ds and SAE e2e reconstruct the activation direction in the current layer similarly well, with SAE e2e+ds doing better at layer 6 and but worse at layer 10.Figure 11 : Distribution of cosine similarities between the original and reconstructed activations, for our SAEs with similar CE loss increases (Table 2 ). We measure 100 sequences of length 1024.How much of the reconstruction error seen earlier (Section 3.2) is due to feature shrinkage? One way to investigate this is to normalize the activations of the SAE output before comparing them to the original activation. 6 In Figure 12 , we compare the explained variance for the reconstructed activations of each type of SAE in layer 6, both with and without normalizing the activations first. Normalizing the activations greatly improves the explained variance of our e2e SAEs. Despite this, the overall story and relative shapes of the curves are similar.Figure 12 : Explained variance between activations from the model with and without the SAE inserted. Measured at all later layers for our set of SAEs with similar CE loss increase in layer 6 (Table 1 ). In (b) we apply Layer Normalization to the activations before comparison.One of our goals is to reduce the total number of features needed over a dataset (i.e. the alive dictionary elements), thereby reducing the computational overhead of any method that makes use of these features. We showed in Figure 4 that SAE e2e and SAE e2e+ds consistently use fewer dictionary elements for the same amount of performance when compared with SAE local . We also see that SAE e2e+ds uses fewer elements than SAE e2e for layers 2 and 6 but not layer 10.Notice in Table 2 , however, that the number of alive dictionary elements is negatively correlated with the norm of the gradients during training. This begs the question: If we increase the learning rate, is it possible to maintain performance in L 0 vs CE loss increase while also decreasing the number of alive dictionary elements?Figure 13 : Varying the learning rate for SAElocal on layers 2, 6 and 10. All other parameters are the same as the local runs listed in the similar CE loss increase Table 2 .In Figures 13, 14, 15, we show the effect that varying the learning rate has on performance for SAE local , SAE e2e , and SAE e2e+ds , respectively. In all cases, we see that learning rates higher than our 2 .default of 0.0005 require fewer dictionary elements for the same level of performance on CE loss increase. We also see that these runs with higher learning rates (up to a limit) can have a better L 0 vs CE loss increase frontier at high sparsity levels and is similar or worse at low sparsity levels.This effect appears to be more pronounced for SAE e2e and SAE e2e+ds than SAE local , indicating that e2e SAEs may require even fewer alive dictionary elements compared to SAE local s than what is presented in the figures in the main text.While not shown in these figures, a downside of using learning rates larger than 0.0005 is that it can cause the L 0 metric to steadily increase during training after an initial period of decreasing. This occurred for all of our SAE types, and was especially apparent in later layers. Due to this instability, we persisted with a learning rate of 0.0005 for our main experiments. We expect that training tweaks such as using a sparsity schedule could help remedy this issue and allow for using higher learning rates. Our architectural and training design choices were selected with the goal of maximizing L 0 vs CE loss increase Pareto frontier of SAE local . We then used the same design choices for SAE e2e and SAE e2e+ds . Much of our design choice iteration took place on the smaller Tinystories-1m due to time and cost constraints.Our SAE encoder and decoder both have a regular, trainable bias, and use Kaiming initialization. To form our dictionary elements, we transform our decoder to have unit norm on every forward pass. We do not employ any resampling techniques [Bricken et al., 2023] as it is unclear how these methods affect the types of features that are found, especially when aiming to find functional features with e2e training. We clip the gradients norms of our parameters to a fixed value (10 for GPT2-small). This only affects the very large grad norms at the start of training and the occasional spike later in training.We do not have strong evidence that this is worthwhile to do on GPT2-small, and it does comes at a computational cost.We train for 400k samples of context size 1024 on Open Web Text with an effective batch size of 16.We use a learning rate of 5e -4, with a warmup of 20k samples, a cosine schedule decaying to 10% of the max learning rate, and the Adam optimizer [Kingma and Ba, 2017] with default hyperparameters.For SAE e2e+ds , we multiply our KL loss term by a value of 0.5 in our implementation. Note that if we instead fixed this value to 1 and varied the other loss coefficients, we would also need to vary other coefficients such as learning rate and effective batch size accordingly, which may have been difficult. This said, fixing this parameter to 1 and having fewer overall hyperparemeters may be a better option going forward, as it turns out to be difficult to tune the other coefficients in this setting anyway. We set the total_coeff (i.e., the coefficient that multiplies the downstream reconstruction MSE, denoted β in Equation 1) to 2.5 for layers 2 and 6, and to 0.05 for layer 10. Note from Equation 1 that this coefficient gets split evenly among all downstream layers. It's likely that a different weighting of these parameters is more desirable, but we did not explore this for this report.It's worth noting that we did not iterate heavily on loss function design for SAE e2e+ds , so it's likely that other configurations have better performance (e.g. having different downstream reconstruction loss coefficients depending on the layer, and/or including the reconstruction loss at the layer containing the SAE).Note that in our loss formulation (Section 2), we divide our sparsity coefficient λ by the size of the residual stream dim(a (l) (x)). This is done in an attempt to make our sparsity coefficient robust to changes in model size. The idea is that the L 1 score for an optimal SAE will be a function of the size of the residual stream. However, we did not explore this relationship in detail and expect that other functions of residual stream size (and perhaps dictionary size) are more suitable for scaling the sparsity coefficient.For GPT2-small we stream the dataset https://huggingface.co/datasets/ apollo-research/Skylion007-openwebtext-tokenizer-gpt2 which is a tokenized version of OpenWebText ( [Gokaslan and Cohen, 2019] ) (released under the license CC0-1.0). The tokenization process is the same as was used in GPT2 training, with a 'BOS' token between documents.We evaluate our models on 500 samples of the Open Web Text dataset (a different seed to that used for training). We consider a dictionary element alive if it activates at all on 500k training tokens.Note that information from all of our runs are accessible in this Weights and Biases ([Biewald, 2020] ) report, including the weights, configs and numerous metrics tracked throughout training. The SAEs from these runs can be loaded and further analysed in our library https://github.com/ ApolloResearch/e2e_sae/.We used NVIDIA A100 GPUs with 80GB VRAM (although the GPU was saturated when using smaller batch sizes that used 40GB VRAM or less).Our library imports from the TransformerLens library [Nanda and Bloom, 2022] (released under the MIT License), which is used to download models via HuggingFace's Transformers library [Wolf et al., 2020] (released under the Apache License 2.0). GPT2-small is released under the MIT license.The Tinystories-1M model is released under the Apache License 2.0 and it's accompanying dataset is released under CDLA-Sharing-1.0.In Figure 16 we show the effect of varying the initial dictionary size for our layer 6 similar CE loss increase SAEs in Table 2 . For all SAE types, we see L 0 vs CE loss increase improve with diminishing returns as the dictionary size is scaled up, capping out at a dictionary size of roughly 60. This comes at the cost of having more alive dictionary elements with increasing dictionary size.It's worth mentioning that, after preliminary investigation on Tinystories-1M, it's possible to reduce the dictionary ratio to 5 times the residual stream and still achieve a good L 0 vs CE loss increase tradeoff, as well as reducing the number of alive dictionary elements. See this Weights and Biases report for details https://wandb.ai/sparsify/tinystories-1m-ratio/ reports/Scaling-dict-size-tinystories-blocks-4-layerwise--Vmlldzo3MzMzOTcw.Figure 16 : Sweep over the SAE dictionary size for layer 6 (where 'ratio' is the size of the initial dictionary divided by the residual stream size of 768). All other parameters are the same as in the similar CE loss increase runs in Table 2 .In Figure 17 we analyse the effect of varying the number of training samples for each SAE type on layer 6 of our similar CE loss increase SAEs. For SAE local , training for 50k samples is clearly insufficient. The difference between training on 200k, 400k, and 800k samples is quite minimal for both L 0 vs CE loss increase and alive_dict_elements vs CE loss increase.For SAE e2e , we see improvements to L 0 vs CE loss increase when increasing from 50k to 800k samples but with diminishing returns. In contrast to SAE local , we see a steady improvement in alive_dict_elements vs CE loss increase as we increase the number of samples. Note that training SAE e2e or SAE e2e+ds for 800k samples takes approximately 23 hours on a single A100.For SAE e2e+ds , the L 0 vs CE loss increase and alive_dict_elements vs CE loss increase improves up until 400k samples where performance maxes out.Figure 17 : Sweep over number of samples trained on layer 6. All other parameters are the same as in the similar CE loss increase runs in Table 2 .We show in Figure 18 that, for a variety of sparsity coefficients and layers, our training runs are robust to the random seed. Note that the seed is responsible for both SAE weight initialization as well as the dataset samples used in training and evaluation. To explore the qualitative differences between the features learned by SAE local and SAE e2e+ds , we first visualize the SAE features using UMAP [McInnes et al., 2018] (Figures 19, 20) .Although there is substantial overlap between the features from both types of SAE in the plot, there are some distinct regions that are dense with SAE e2e+ds features but void of SAE local , and vice versa. We look at the features in these regions along with features in other identified regions of interest such as small mixed clusters in layer 6 of GPT2-small in more detail. We label the regions of interest from A to G in Figure 19 , and provide human-generated overview of these features below. Features from this UMAP plot can be explored interactively at https://www.neuronpedia.org/gpt2sm-apollojt.For each region, we also share links to lists of features in that region which go to an interactive dashboards on Neuronpedia.Figure 19 : UMAP plot of SAEe2e+ds and SAElocal features for layer 6 on runs with similar CE loss increase in GPT2-small.Many of these features appear to be late-context positional features, or miscellaneous tokens that only activate in particularly late context positions. It may be the case that SAE e2e+ds has fewer positional features than local, as indicated by the 18 SAE e2e+ds vs 91 SAE local local features in this region (and similar in surrounding reasons). This said, we have not ruled out whether positional features for SAE e2e+ds are found elsewhere in the UMAP plot.This region mostly contains features which activate on 
 tokens, in addition to some newline and double newline. These are tokens that mark the beginning of a new context. Seemingly SAE e2e+ds contains many more distinct features for 
 than SAE local .Region C potentially suggests more feature splitting happening in SAE local than SAE e2e+ds . For SAE e2e+ds , each feature activates most strongly on tokens "by" or "from" in a broad range of contexts.For SAE local , each feature activates most strongly in fine-grained contexts, such as "goes by" vs "led by" vs ". By" vs "stop by" vs "<media>, by author" vs "despised by" vs "overtaken by" vs "issued by" vs "step-by-step / case-by-case / frame-by-frame" vs "Posted by" vs "Directed by" vs "killed by" vs "by".These features all activate on "at" in various contexts. As in Region C, the SAE e2e+ds features appear less fine-grained. Examples of SAE local features not present in SAE e2e+ds : "Announced at" or "presented at" or "revealed at" feature (https://www.neuronpedia.org/gpt2-small/6-res_ scl-ajt/40197). "At" in technical contexts (https://www.neuronpedia.org/gpt2-small/ 6-res_scl-ajt/34541).All features appear to boost starting words which would come after a paragraph or a full stop to start a new idea, such as "Finally", "Moreover", "Similarly", "Furthermore", "Regardless", "However" and so on. They seem to be differentiated by perhaps activating in different contexts.For example https://www.neuronpedia.org/gpt2-small/6-res_scl-ajt/4284 activates on full stops and newlines in technical contexts so it can predict things like "Additionally", "However", and "Specifically". On the other hand, https://www.neuronpedia.org/gpt2-small/6-res_ scl-ajt/13519 activates on full stops in baking recipes so it can predict things like "Then", "Afterwards", "Alternatively", "Depending" and so on.These seem mostly similar to Region E. It's not clear what distinguishes the regions looking at the feature dashboards alone.The features in both SAEs seem to activate on fairly specific different words or phrases. There is no obvious distinguishing features. It's possible that SAE e2e+ds features tend to activate more specifically and on fewer tokens than the corresponding SAE local features. An example of this can be seen when comparing https://www.neuronpedia.org/gpt2-small/6-res_ scefr-ajt/13910 (a SAE e2e+ds feature), with https://www.neuronpedia.org/gpt2-small/ 6-res_scl-ajt/45568 (a SAE local feature).In Figure 20 , we show UMAP plots for layers 2 and 10. We interpret a single region from layer 10 in the next section. While some individual features in this region are interpretable, there is no obvious uniting theme semantically. There is, however, a geometric connection. In particular, these are features that point away from the 0th PCA direction in the original model's activations (Figure 21 ). The 0th PCA direction is nearly exactly the direction of the outlier activations at position 0 (see also Appendix B). Activations in this direction are tri-modal, with large outliers at position 0 and smaller outliers at end-of-text tokens (Figure 22 ).We can measure how well an SAE preserves a particular direction by measuring the correlation between the input and output components in that direction. Our SAE e2e+ds faithfully reconstructs the activations in this direction at position 0 (r = 0.996), but not at other positions (r = 0.262). This is a particularly poor reconstruction compared to SAE local or other PCA directions (Figure 23a ). SAE e2e+ds 's poor reconstruction of the activations in this direction implies that the differences may not be functionally relevant. We can measure this by resample ablating the activation in this direction at all non-zero positions. This means we perform the following intervention in a forward hook:a(x) i ← a(x) i -P a(x) i + P a(x ′ ) jWhere a(x) i is the activation at position i > 0, a(x ′ ) j is the resampled activation for a different input x ′ and position j > 0, and P is a projection matrix onto the 0th PCA direction.After performing this ablation, the kl-divergence from the original activations is only 0.01. This difference is smaller than repeating the experiment for any other direction in the first 30 PCA directions (Figure 23b ).This means that the exact value of this component of the activation (at positions > 0) is mostly functionally irrelevant for the model. SAE local still captures the direction faithfully, as it is purely trained to minimize MSE. While SAE e2e+ds fails to preserve this direction accurately, this seems to allow it to have a cleaner dictionary, avoiding SAE local 's cluster of features that point partially away from this direction.The training time for each type of SAE in GPT2-small is shown in Table 6 . We see that e2e SAEs are 2-3.5x slower than SAE local . Note that one can reduce training time with little performance cost by training on fewer that 400k samples (Figure 17 ) and/or using an initial dictionary ratio of less than 60x the residual stream size (Figure 16 ). Using locally trained SAEs as initialization for e2e SAEs or training multiple SAEs at different layers concurrently are also possible solutions. Our main evaluation metrics presented in Section 3 measure the functional importance of the features learned by the SAEs on the next-token language modeling task used to train the model. We also experimented with evaluating the SAEs on a downstream task: how faithfully the dictionaries represent the information the model uses to perform subject-verb agreement. This task is directly inspired by the analysis in Marks et al. [2024] .We use datasets from Finlayson et al. [2021] with 4 variations of a subject-verb agreement task:• Simple: The parent/s is/are• Across participle phrase (PP): The secretary/secretaries near the cars has/have• Within relative clause (RC): The athlete that the manager/managers likes/like• Across RC: The athlete/athletes that the managers like do/doesFor each template, we use 1000 datapoints with different subjects and verbs. For each input, we can compute the logit difference that the model assigns to the correct and incorrect forms of the verb.Following Marks et al. [2024] we compute the faithfulness of this logit difference when intervening on the network's activations. Let m represent the mean logit difference between the correct and incorrect verb forms across the dataset. Let M be the original model and M be the model under some intervention. We measure the faithfulness of the intervention as m( M )-m(∅) m(M )-m(∅) where ∅ represents ablating the entire residual stream.A faithfulness of 0% thus means M performs no better than random, while a faithfulness of 100% means the intervention does not change performance. Faithfulness numbers greater than 100% mean the model is, on average, more confident in the correct verb with the intervention than without.We first test the faithfulness of the models with the SAEs inserted (Table 7 ). An SAE which preserves functionally relevant features in the activations would have faithfulness close to 100% on all tasks the model is trained to perform. While all SAEs preserve most of the logit-difference, there is significant variation across SAE types, layers, and tasks. The local SAEs in layer 10 have the worst faithfulness, although SAE e2e+ds in layer 6 also has poor faithfulness across participial phrases.We are also interested in if end-to-end training helps concentrate the functional relevance into a few specific SAE features. We thus ranked SAE features in Layer 6 Similar L 0 SAEs by their indirect effect, measured by the change in faithfulness score when mean ablating that feature. 7 We then mean ablate all SAE features except the k with largest indirect effect (Figure 24 ). We see SAE e2e needs comparatively few features to achieve high faithfulness scores, but SAE local and SAE e2e+ds have roughly similar curves.Despite certain (SAE type, layer) combinations showing superior performance, there are no clear patterns between SAE types across tasks. These results indicate that e2e SAEs do not provide an obvious benefit on the selected downstream tasks. While e2e SAEs demonstrate benefits for the language modeling task on the full OpenWebText distribution, further work would be needed to find specific tasks and sub-distributions in which they provide the most benefit. Or, perhaps these sort of task-specific results are very noisy and it would be necessary to aggregate across many tasks and templates to differentiate between SAE training methodologies.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification:Repository attached to the submission (https://github.com/ ApolloResearch/e2e_sae) Guidelines:• The answer NA means that paper does not include experiments requiring code. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material.We use zero-indexed layer numbers throughout this articleWe show in Appendix C that it is possible to reduce the number of alive dictionary elements for any SAE type by increasing the learning rate. This has minimal cost according to L0 vs CE loss increase Pareto-optimality up to some limit.Measured using linear interpolation over a range of L0 ∈ (50, 300). This range was chosen based on two criteria: (1) L0 should be significantly smaller than the residual stream size for the SAE to be effective (we conservatively chose 300 compared to the residual stream size of 768), and (2) the CE loss should not start to increase dramatically, which occurs at approximately L0 = 50 (Figure1).To "normalize" we apply center along the embedding dimension and scale the resulting vector to have unit norm. This is equivalent to Layer Normalization with no affine transformation. We use this as Layer Normalization is applied to the residual-stream activations before they are used by the network. (a) Unmodified activations (b) Normalized activationsWhen mean ablating an SAE feature we take it's mean position-specific value. 


Decision: **Summary of the paper:**
The paper proposes a new way to train sparse auto-encoders (SAE) to be used for interpreting pre-trained deep networks. Instead of directly minimizing the discrepancy of the feature out of SAE with that of an intermediate layer of the original network, it minimizes the discrepancy at the output layer, i.e., probability distribution over the next token in the case of LLMs considered in this work. It obtains the output for SAE by passing the SAE’s encoded feature through the rest of the layers of the original network. It also uses a sparsity inducing regularization. It further adds a reconstruction loss to some later layers.

**Summary of the reviews:**
The reviewers found the proposed method well-justified, the reasoning of the paper convincing, the limitations well-covered, the results mostly suggesting improved performance and the geometric study of the learnt features informative .

On the other hand, they found evaluating metrics not easy to interpret, not entirely convincing, or sometimes incrementally improved, the mitigation of inconsistent pathways still possibly problematic when moving the reconstruction to a later layer, a lack of ablation study when only the later reconstruction is used (without KLD on output). 

**Summary of the rebuttal and discussions:**
The authors provided new experiments on a different evaluation metric than CE for faithfulness. 
This along with the textual clarifications on the evaluation metrics rectified some of the concerns of the reviewers. However, some reviewers remained somewhat concerned about the metrics. Despite that, all reviewers are leaning towards acceptance.

**Consolidation report:**
As reviewers have pointed out a crucial strength of the paper is the outright transparency on the limitation of their approach and importantly the limitations of the experiments and evaluations in supporting their basic claims and, in fact, the paper often discusses alternative hypotheses. This mitigates most of the issues raised above, which are especially valid regarding the metrics. This strength in tandem with the fact that the paper has taken clear steps forward for SAE interpretation of LLMs, make a good paper.

**Recommendation:**
The AC agrees with the general tendency of all reviewers and recommends acceptance. The discussions (including the concerns raised by reviewers) and additional points and results provided by the authors during the rebuttal would be quite enlightening for future readers and thus the authors are encouraged to incorporate them in the camera ready version of the paper. It would be good to have a section discussing the metrics used and their limitations as was extensively discussed with the reviewers.

Reviews: review: paper_summary: paper_summary: The paper proposes new framework to improve the interpretability of large language models. Inspired by the work proposed in Anthropic's "", the proposed work replaces the original reconstruction loss in the sparse auto-encoder (SAE) with KL-div loss, and additionally add further constraint to minimize the error between reconstructed downstream activations and original activations in the LLM. Empirical results demonstrates the adavantage of the proposed new SAE_{e2e} architecture series.
strength_weakness: strength_weakness: W1: The authors may further explain the intuition between the evaluation metrics. The current version of the evaluations is not explicitly clear to me why the proposed method is advantageous in comparison to the baselines. 

W2: It is now clear how the KL divergence is implemented to replace the reconstruction loss. KL is only practically useful when the distributions are available. However, in this problem, it is unclear how the distribution loss between the two deterministic feature outputs are computed.  The motivation of the replacement between KL divergence and the L2 reconstruction is also not very clear to me, and these two losses should be ultimately equivalent and optimal at the same parameters global optima, although the loss changes the optimization landscapes.
questions: questions: I appreciate it if the authors could please see the above weakness for my questions.
limitations: limitations: Yes the authors have adequately discussed the limitations of the work.
paper_summary: paper_summary: This paper proposes a way to train sparse autoencoders (SAEs) that encourages them to learn features that are causally relevant to the model's output. This is done by replacing the usual SAE reconstruction objective - an L2 penalty between original activations and their reconstructions - by the KL divergence between the model's output distribution and the distribution when activations are replaced by their reconstructions. Additional loss terms encourage the reconstructions to lead to downstream layer activations similar to the original downstream layer activations. The sparsity penalty on feature activations is kept as in the "vanilla" SAE.

The paper finds that this approach is a clear Pareto improvement over vanilla
SAEs w.r.t. the tradeoff between sparsity of the feature activations (L0) and
loss recovered when using reconstructions. An automatic interpretability pipeline is used to compare the features learned to vanilla SAE features. A statistically significant advantage for e2e+ds SAEs is found.

The paper also attempts to avoid a potential failure mode of the approach, whereby the SAE may learn to exploit different pathways through the model in order to drive down the KL divergence to the original output distribution, because this is easier from an optimization point of view (as the reconstructions are no longer optimized to match the original activations at the layer where the SAE is applied). This is why one of the SAE variants considered includes an L2 penalty for downstream layer activations too. Results here show that using only the KL penalty in the loss produces very different downstream activations compared to adding these additional terms, suggesting that the
strength_weakness: strength_weakness: - given that the KL divergence to the output distribution is incorporated in the
loss function, it is not that surprising that the methods in this paper have
better values of the loss recovered metric (to their credit, the authors
acknowledge this). To truly conclude the superiority of the suggested method for
surfacing *individual* causally important features, it would be extremely helpful to have some external (to the KL metric)
evaluation. 
- To some extent the paper shows such evaluations. For example, in Appendix G.3 it is shown that vanilla SAEs represent well a feature that is not causally important, and the e2e SAEs in turn represent it poorly. However, what is *really* required to establish superiority is the opposite: to exhibit a causally important feature not represented by vanilla SAEs, but represented by e2e ones.
- the key problem when not using the same-layer activations as a reconstruction target for the SAE is (as the authors describe) that the optimization may prefer to learn features that achieve a good KL divergence value eventually, but do so through "pathological" pathways in the model. To fix this, the authors encourage closeness with activations starting from the next layer up. However, such "pathological" pathways may exist in a single layer (this was shown e.g. here https://arxiv.org/abs/2311.17030). So it is unclear whether the problem has been fully overcome or rather restricted to a more narrow part of the model. Furthermore, even if reconstructions of downstream layers are similar to the true activations of these layers, this does not establish that the individual features themselves are not "pathological". Again, some additional, *per-feature* evaluation is needed to make the (strong) claims of the paper more believable.
- out of the three kinds of SAEs considered - vanilla, e2e, e2e+ds - there seems to be a missing one implied by the others: an SAE that encourages faithful reconstructions only of the *next* layer activations, and does not involve the KL divergence loss. Would such an SAE lead to similar improvements? If so, this would be stronger evidence, as the KL divergence w.r.t. the final output distribution won't be a part of the loss.
questions: questions: - did you try the "next layer only" SAE variant described in the weaknesses?
- how can we get a more fine-grained picture of how individual features change between the vanilla SAE and the e2e ones?
limitations: limitations: - I think the main limitation that I would have loved to see addressed more in the paper is that the evaluations are somewhat indirect w.r.t. the main claim of the paper. Even if in some average sense the e2e SAEs' reconstructions lead to better KL divergence, we still don't know how this plays out on the level of individual features.
paper_summary: paper_summary: The authors present their observations on modified sparse autoencoder (SAE) learning. In the proposed method SAE is trained to reconstruct original model weights (features). SAE is optimized with the KL-divergence loss between the model output and the output at the same location in the original model. An additional variant is training with the reconstruction loss computed after each layer between the original model and SAE output. The sparsity loss term is added to increase the number of weights that will converge to zero. The authors claim that their method explains the model features in a better way with fewer dictionary elements.
strength_weakness: strength_weakness: - The modification is very simple and raises the question of whether additional variants could work too, e.g. local SAE with MSE end-to-end loss (in addition to the layer reconstruction MSE)
- The numeric results in Figure 1 are very incremental: e.g. reduction of 0.05 in CE loss increase L0=100 is ~1.6% of the evaluation CE loss
- The interpretability is not well defined, in addition, the claim that fewer features means more interpretability should be proved by empirical evidence, white the authors showed that interpretability was not changed or changed significantly (appendix A7).  If I get it correctly, the main goal of the proposed method is interpretability but a few experiments were done to evaluate it. Also, qualitative results are not provided (anonymous url)
-  Std values are missing (for different random initialization seeds) in Figure 1
- The writing should be improved, Sec. 3.4 should be more self-contained
- It is not clear which method e2e or e2e-ds  is the best performer overall?
questions: questions: - Can you please explain the diagram in Figure 1 why some blocks are of different shapes?
- It looks like that the Sec. 3.4 result is trivial, if we don’t encourage the SAE to reconstruct after wach activation, that will be the result.
- What is the number of parameters you train in SAE in total for all layers? 
- Again, if the main benefit of the method is interpretability, I do not understand how this work addresses this property and what is its main contribution in improving interpretability. 
- Are $W_e$ and $D$ the matrices defined for each layer? If they are - an appropriate indexing should be applied.
limitations: limitations: not discussed
paper_summary: paper_summary: This paper proposes to train Sparse Autoencoders (SAEs) with an additional loss function term: The KL divergence between the original output distribution and the output distribution obtained when using model with the inserted SAEs. This additional loss pushes the SAEs to focus on functionally important features
The authors find that this setup requires fewer total features and fewer simultaneously active features per data point.
strength_weakness: strength_weakness: - I think the major weakness is that the contribution is not very groundbreaking.
- There are no new findings wrt interpretability, and the e2e SAEs require more compute to be trained
- while the authors discuss several findings (robustness to random seeds, reconstruction loss, training time) the reader is left waiting a bit for a conclusion or take home message
- the comparison to local SAEs is sometimes a bit unclear:
	- the authors state that "locally trained SAEs are capturing information about dataset structure that is not maximally useful for explaining the algorithm implemented by the network." but do not show that their e2e SAEs are significantly more interpretable or helpful at explaining the algorithm implemented by the network, instead the e2e SAEs are merely "at least as interpretable as SAElocal"
	- the features found by local SAEs are not the functionally important ones , but SAElocal and SAEe2e+ds features are somewhat similar so local SAEs can be used as initializations (line 233)?


- notation could be improved 
	- N_dict_elements and d_hidden seem to belong more into pseudo code so it looks a bit odd to mix it with more traditional mathy notation
	- $\lambda$ is hidden in the $\phi$ parameter but then you show values for lambda in Table 1 and looking back at i.e. Eq. (1) it's not immediately clear what $\lambda$ is and why $\phi$ is not mentioned. I would suggest either putting $\lambda$ directly in the equations or citing values for $\phi$
questions: questions: line 50: but we no longer optimize activation reconstruction. -> but you do use activation reconstruction for $SAE_{e2e+ds}$, don't you?

Maybe I just missed i but where in the network do you plug in the SAEs? the residual stream/decoder layer output?

You say that taking different computational pathways through subsequent layers of the network might be a problem, but do you have any evidence for this happening? is this just section 3.3.2?

While reading, I was wondering if taking different computational pathways through subsequent layers of the network is still a problem if only one SAE is trained at a time. There should be little incentive for computing additional information since the original network post SAE insertion would not have capacity to handle these anyway, so it would only matter for the output directly.


small things:
- line 68: not sure if it's necessary to mention goodhart's law since the target is not cross-entropy loss difference but interpretability so I would be more worried about this if you start optimizing interpretability measures 
- line 5: datatset -> dataset
- line 44: a feature is functional important -> a feature is functionally important
- line 113: ")" missing in equation
limitations: limitations: Authors briefly discuss the lack of a ground truth to evaluate SAEs.

I think discussion of evaluation metrics could be a bit more divided (maybe into metrics to assess interpretability and metrics to assess how close to the original network a network with inserted SAEs is).
paper_summary: paper_summary: This work describes two new type of sparse autoencoders (SAEs) that are trained with two new loss functions, L_{e2e} and L_{e2e + ds}. L_{e2e} loss penalizes the KL divergence of the original model vs. when the SAE is inserted, as opposed to existing SAE training techniques which instead optimize MSE with the original activations. L_{e2e + downstream} additionally seeks to minimize the MSE of downstream layers after the SAE is inserted vs. the original model. Both types of new SAE have a better L0 vs. CE tradeoff than naive SAEs, and e2e + downstream additionally has similar downstream reconstruction loss. Interpretability is no worse for the new SAEs. Finally, e2e+ds has similar geometric properties to naive SAEs, including discovering similar features.
strength_weakness: strength_weakness: - The main evaluation metrics that show improvement (CE loss for both new SAEs and later layer reconstruction error for e2e + downstream) are the same ones that are being optimized. This is noted by the authors when they mention Goodhart's law, but it would greatly improve the paper if there were additional metrics for downstream SAE quality that were tested (e.g. some of the methods described in section 4.4). 
- Feature shrinkage (shown in Table 5) seems much worse with the new SAE training methods. Moreover, even with e2e + ds SAEs, later layer reconstructions in Figure 12 are much worse than local. This seems potentially a problem for e.g. finding circuits or understanding interactions between multiple SAEs via feature similarity. This is another reason I think additional downstream metrics would be useful.
questions: questions: - Building off the above weaknesses, could the authors add a measure of downstream SAE functional importance that is not CE loss? 
- Why doesn't the original SAE learn the "left block" of features that SAE e2e + ds does? It would be great to see if e.g. these have a small or negative impact on the variance explained.
- Did the authors try training an SAE with KL divergence + the local reconstruction error (i.e. MSE of layer X reconstruction for an SAE trained on layer X)? It is a little concerning to me that the layer X variance explained is so strongly negative because it means the SAE's reconstruction is very far off distribution.
- Can you quantify how much farther the <e2e ds to local> difference is vs. the inter seed distance? It is hard to know how to interpret Appendix G without this context, as inter-seed SAEs have different features learned too.
- The ~0.15 CE loss e2e+downstream SAE looks closer to the other chosen SAEs than the 0.125 CE loss one, does it change things to run with the ~0.15 one?

Small points:
- Line 41 states "a significant amount of [causal effects of some features is] mediated by the reconstruction residual errors", but I believe the cited work shows that this is true for circuits as a whole, not features individually.
- Lines 133-136 are confusing. Why is it okay to "clear out" features from layer X, but not layers X + 1...? In general, a more clear discussion of the tradeoff between learning functionally important features and maintaining model behavior would help I think.

I would be happy to raise my score if some of these questions/concerns were addressed.
limitations: limitations: I believe the authors should more explicitly discuss the Goodhart's law problem I've discussed above and ideally the increased feature suppression as well. Otherwise, I think the limitations and impacts are well addressed.

Comments: title: Thanks for your rebuttal
Valuable, if negative, results
comment: Thanks for the rebuttal and the further clarifications on my questions. Given the explanations on the motivations behind using KL to replace L2, I decide to increase my score. I also encourage the authors to include the discussions on these differentiating factors between L2 and KL into the final version of their paper.
Thank you for the explanations. I think that the paper is still borderline and should be improved. I raise my score appropriately.
> In Figure 1 you present CE loss increase, not the error, so I don't understand you claim of "We would not call a halving of the error incremental.".

Apologies if our use of the terms "error" and "CE loss increase" caused confusion in the rebuttal, we in fact did mean to use them interchangeably. To expand: The optimal CE Loss Increase (the metric used in Figure 1) is 0. Achieving this would mean that we could splice in the SAE, run the SAE activations through the rest of the network, and achieve exactly the same output as the original model. This would thus give us a complete representation of the features that matter for the output of the network.

When we splice in a "vanilla" local SAE and then pass the SAE activations through the rest of the network, the CE loss of that network is worse than that of the original network. The yellow line in Figure 1 (left) shows how much worse it is. E.g. at an L0 of 100, it gives a CE Loss Increase of 0.11 compared to the original model.

When we splice in an e2e SAE (either pure e2e or e2e+ds), we only get a CE loss increase of 0.05. So the amount of CE Loss Increase has ~halved from 0.11 to 0.05. As mentioned in our rebuttal, we do not think this is a small difference when considering that models containing many times more parameters only slightly decrease CE loss as per model scaling laws. Or perhaps more starkly, to get the same amount of CE Loss as our e2e models with L0=100 (CE Loss Increase of 0.05), a local SAE needs to activate ~400 features on average for each input (i.e. 4 times larger description length!).

> I meant Sec. 3.2 not 3.4

Ah, in this case we do agree that it is at least unsurprising that pure e2e performs much worse than local SAEs on this metric. Although we do find it valuable to know exactly how close our e2e+ds and local are on downstream layers. We use a reconstruction loss at downstream layers in the e2e+ds case, but not the local case, so it wasn't clear a priori how this plot would turn out for the optimal hyperparameters in each setting. It could have turned out that a smaller downstream reconstruction coefficient was optimal for e2e+ds SAEs, meaning that the downstream reconstruction MSE would be worse and we would be more concerned about the SAE taking different computational pathways through the network.

We hope this helps clarify your remaining concerns.
Thank you for the detailed and thoughtful rebuttal, as well as the overall response.

I appreciate the additional evaluations of your proposed method. 
My interpretation of the additional experiments is that they provide at best mixed evidence
for the superiority of the proposed SAE variants. That being said, I appreciate the honesty
of the authors, and I believe these results will be valuable (and perhaps surprising) to the
interpretability community, so I am in favor of disseminating these results widely. 

This is why I maintain my recommendation to accept this work (though I will not be changing
my score at present).
Thank you for the rebuttal.

In Figure 1 you present CE loss increase, not the error, so I don't understand you claim of "We would not call a halving of the error incremental.".

I meant Sec. 3.2 not 3.4
Thanks for emphasizing again that your main contribution lies in reducing the description length. I do agree that this is a valid improvement over existing SAEs.
I'm sorry for being too unclear wrt my comment about evaluation metrics. You do discuss different metrics in different sections. I guess I was looking more for a high level picture. Sth like you want to maximize faithfulness between the network with and without inserted SAEs (to make sure you interpret the original network) and interpretability at the same time and how the different metrics serve these goals.
I do think it is hard to assess how much the reduction in description length helps with the goal of interpretability (maybe showing that pure reconstruction loss SAEs features can be misleading or sth could be helpful), so I did not update my score.
Thank you for responding to my comments! Most of my concerns are addressed, but I remain concerned that the only real metric that shows improvement is the one that is being optimized (the new experiments are a pretty limited evaluation and don't show a compelling improvement). Some ideas might be using board games like in the recent work from [Karvonen et. al.](https://arxiv.org/abs/2408.00113), showing more faithful circuits over a wide dataset like the ones studied in [Marks et. al.](https://arxiv.org/abs/2403.19647), or something else. I believe without more compelling downstream metrics, I cannot raise my score, so I will keep it as is. Thank you again!