Efficient Dictionary Learning with Switch Sparse Autoencoders

Response to weaknesses, part 2

This paper only uses one Language Model. GPT-2 Small has a specfic architecture which is unlike modern LLMs in several ways (only uses GeLU not gated MLPs, uses absolute positional embeddings, ...) which may mean the results do not generalize. (But this is unlikely since most work has found so far that SAE approaches that work on model 1 also work on model 2). The benefits of MoE SAEs are only at very large scale, and GPT-2 Small 800M token training runs are not large scale. When scaling ML techniques it is often found that small-scale hacks perform worse in large setups. This is still a plausible issue given the results in paper.

Thank you for pointing these concerns out! We have added the following table in the appendix, which shows results for training GPT-2 TopK and Switch SAEs at the same width and sparsity on additional layers and sites, as well as Gemma-2B SAEs on a single layer, and have added a reference to it in section 4.2.2:

The * means that these SAEs have not yet trained to convergence, but should converge in the next couple of days. Once they are done we plan to update the table in this comment and add a reply.

Overall, we find that our results replicate on a larger model, Gemma 2 2B, and additional GPT-2 residual layers, but do not replicate as well on earlier attention and MLP layers of GPT-2. Currently, residual stream SAEs are the most widely used by the community, so we believe this result does not strongly negatively affect our contribution, but we certainly are interested in future work investigating this disparity. We have added a discussion of this limitation to our conclusion.

We further expect mixture of experts based techniques to work well at scale for two reasons:

• We see increased speedup with larger SAEs in Figure 1

• Mixture of expert models have been shown to work very well at large scales in language models, see e.g. Mixtral of Experts [1]

Response to questions

How can I calculate the FLOPs for a Switch SAE with given num experts, width, input dim and training datapoints? I did not see this in the paper, and therefore sentences like "As we scale up the number of experts and represent more features, performance continues to increase while keeping computational costs and memory costs (from storing the pre-activations) roughly constant" are very confusing.

Thank you very much for bringing up this concern. We have added Appendix A.3 to explain the FLOP calculations and have added a reference to Appendix A.3 in section 4.2.1. We apologize for not including this initially and greatly appreciate your attention to detail.

Are the authors able to open source training code and/or weights when this is deanonymized? Since there have been known cases in the SAE literature where promising work has not replicated [8, 9], open weights and training code will enable downstream validation of this technique to be much stronger.

Yes, we will open source our training code and weights when our work is de-anonymized. We are also planning on integrating Switch SAEs within existing tools such as SAE Lens (https://jbloomaus.github.io/SAELens/).

Do you have any experiments on different models or different sites?

Yes, see above, thank you!

If our comments resolve your concerns, we would appreciate it if you would consider raising your score!

[1] https://mistral.ai/news/mixtral-of-experts/

Thank you for your extremely thoughtful comments. We greatly appreciate you taking the time to think deeply about our work.

Response to weaknesses, part 1

There is no application of the Switch SAE in the paper. Recently, SAEs have been applied to steering [10] and circuit finding [11] as well as many other applications in AI safety and capabilities. These applications rely on a sparse set of features explaining model behavior. But if there are N sets of sparse features, different sets for each of the experts, then these applications may not want to use Switch SAEs despite performance benefits. This is a deeper problem than merely "there may be duplicated features in different experts (as measured by cosine sim / t-SNE)". For example, on a given dataset of interest, such as a dataset used to generate an SFC graph [11] or the set of prompts related to e.g. the Golden Gate Bridge [10], there may be different ways the different experts reconstruct similar prompts. As a toy example, on a set of prompts where the "Queen" concept is always present, on one prompt Expert 1 may be chosen, and encode V(Queen) = Feature(Woman) + Feature(Royalty). But on another prompt Expert 2 may be chosen and have a dedicated feature Feature(Queen). Let me know if I need to make this explanation clearer.

This is an excellent and thoughtful point that we had not considered, thank you! We have added a discussion of it to the limitations section of our conclusion. Interestingly, however, we believe that our existing experiments show that it may not actually be a large issue. Specifically, our auto-interp evaluation asks an LLM to predict when a feature fires depending on input text. If your concern was happening frequently, then detection scores should be much lower, since sometimes the “same” feature would fire on similar types of text in expert 1 versus expert 2 in different contexts, and we would not be able to predict that “cross-expert feature” successfully. However, detection scores are only slightly lower for the Switch SAE. Thus, we believe that reconstruction MSE, loss recovered and auto-interpretability score are together a good proxy for downstream tasks, even in the Switch SAE setting.

This paper bases their approach off of the Switch MoE architecture [1], but this is a weak MoE variant. My understanding is that the switch MoE architecture is weaker than different MoE architectures. For example [2] benchmarks Switch MoEs against GShard [3] and their own DeepSeekMoE [2], and find that Switch MoEs are weakest.

Thank you for pointing this out, we agree that other mixture of expert variants may perform better. However, we do not view this as solely a weakness of our work; indeed, using the simplest possible MoE variant in the paper allows us to focus on whether and how the technique works at all for language model feature learning, and thus provides a solid baseline that future work can improve upon. We have added a section to our conclusion discussing this future work, thank you again!

Additionally, this paper requires an auxiliary loss to balance MoE load. While this is popular in conventional MoE work [1, 2, 3], other works propose to do this [4]. In the context of SAEs, several improvements have involved dropping auxiliary loss terms (e.g. [4] was replaced by [5] by the Anthropic group, and [7] was replaced by [8] by the Google DeepMind group). So I am also concerned that auxiliary loss terms are worse than (e.g.) expert-choice variants of MoEs in training MoE SAEs.

This is another excellent point, thank you! We agree that an auxiliary loss term is less elegant, but for the same reasons as above we leave exploring additional mixture of expert architectures to future work. Also, as shown in the new Figure 7 in the appendix, Switch SAEs at least seem robust to the auxiliary weighting term $\alpha$, so this is not a sensitive parameter that needs to be carefully optimized.

Thanks for noting all my points as limitations in the paper, or doing the additional training runs / experiments to address some doubts. It is strange that attn and MLP sites have worse switch SAE performance, but they are less important than the residual stream site, so I'm confident this paper should be accepted.

Response to questions, part 2

The references seem biased towards recent years. Besides the Adam optimizer and GPU acceleration, no references exist from before 2016, despite sparse autoencoders, sparsity, topK, and MoEs being popular research topics much earlier than that. Are these not an important background to sparse SAEs?

Thank you for bringing this to our attention. We have added additional citations to our submission to address this. For sparse autoencoders, we added [7], [8], [9] to section 1 and section 2.3. For sparsity, we added [10], [11], [12], [13] to section 2.3. For TopK, we added [14] to section 1 and section 2.3. For MoE, we added [15], [16], [17], [18], [19], [20] to section 2.1.

Response to suggestions

Line 295: “For a wide range of sparsity (L0) values”. Specify the range.

Thank you for pointing this out! The range is 8 - 128; we have added this clarification to our submission.

Line 298: Clarify what is meant by a zero-ablation.

Thank you also for mentioning this! Zero-ablation corresponds to setting the residual stream to all 0s at that layer; we have now added this clarification to our submission. Note that this is a standard metric used for evaluating SAEs, see e.g. [2], [4], [21].

Line 317: “Switch SAEs can likely achieve greater acceleration on larger language models”. Expand on this.

We believe this is true for two main reasons. First, we see an increased speedup with larger SAEs in Figure 1. Second, mixture of expert models have been shown to work very well at large scales in language models, see e.g., Mixtral of Experts [22].

If our comments resolve your concerns, we would appreciate it if you would consider raising your score!

[1] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, Fedus et al.

[2] Towards monosemanticity: Decomposing language models with dictionary learning, Bricken et al.

[3] Update on how we train SAEs, Conerly et al.

[4] Improving dictionary learning with gated sparse autoencoders, Rajamanoharan et al.

[5] Scaling and evaluating sparse autoencoders, Gao et al.

[6] Automatically Interpreting Millions of Features in Large Language Models, Paulo et al.

[7] Sparse deep belief net model for visual area v2, Lee et al.

[8] Building high-level features using large scale unsupervised learning, Le et al.

[9] Zero-bias autoencoders and the benefits of co-adapting features, Konda et al.

[10] Sparse coding with an overcomplete basis set: A strategy employed by v1?, Olshausen et al.

[11] Sparse and redundant representations: from theory to applications in signal and image processing, Elad.

[12] Method of optimal directions for frame design, Engan et al.

[13] K-svd: An algorithm for designing overcomplete dictionaries for sparse representation, Aharon et al. [14] K-sparse autoencoders, Makhzani & Frey.

[15] Adaptive mixtures of local experts, Jacobs et al.

[16] Hierarchical mixtures of experts and the em algorithm, Jordan & Jacobs.

[17] Mixtures of gaussian processes, Tresp.

[18] A parallel mixture of svms for very large scale problems, Collobert et al.

[19] Infinite mixtures of gaussian process experts, Rasmussen & Ghahramani.

[20] Expert gate: Lifelong learning with a network of experts, Aljundi et al.

[21] Open Source Sparse Autoencoders for all Residual Stream Layers of GPT2-Small, Bloom.

[22] https://mistral.ai/news/mixtral-of-experts/.

Thank you for your thoughtful comments. We greatly appreciate you taking the time to review our work.

Response to weaknesses

The hyperparameter $\alpha$ is set to 3 (line 258) without further discussion. How was this value chosen? Is this a reasonable default in most settings or is tunning required based on the model and dataset?

Thank you for bringing this to our attention. We chose $\alpha = 3$ based on a preliminary hyperparameter sweep, where we found that results were not very sensitive to $\alpha$ and that $\alpha = 3$ was a reasonable default. We trained 16-expert and 32-expert Switch SAEs at a fixed sparsity (L0 = 64) and width (24576 features) and swept $\alpha$ across the values [0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100]. We have added a plot (Figure 7) corresponding to this hyperparameter sweep to Appendix A.2. We apologize for not including this in our initial submission.

The abstract claims a “substantial Pareto improvement” for a fixed training budget. This corresponds to the FLOP-matched experiments shown in the bottom left plot of Figure 3. It is not directly obvious that this constitutes a substantial improvement. The difference looks small enough to warrant repeated experiments to ensure its significance. A table of improvements with confidence intervals would make the difference much clearer.

The substantial Pareto improvement is much more obvious in the top left plot of Figure 3, showing reconstruction MSE. The substantial Pareto improvement for a fixed training budget is also apparent in our scaling laws (specifically, the left subplot of Figure 1). The Pareto improvement is less visually obvious in the bottom left plot of Figure 3 because the fraction of loss recovered (FLR) metric is saturated. We have inverted the y-axis (plotting 1 - FLR as opposed to FLR) in our updated submission to highlight the difference more clearly. In general, SAEs trained to convergence are very stable in terms of their reconstruction MSE and FLR.

Figure 4 could use an explanatory caption.

Thank you for bringing this to our attention. We have added an explanatory caption to our submission. The caption now reads “Switch SAE feature geometry experiments, measured via cosine similarity between SAE decoder vectors. We find that Switch SAEs with more experts exhibit more feature duplication, but that this effect diminishes for larger L0s. Furthermore, between-expert similarities show that experts specialize moderately; expert 0, for example, has low similarity with most other experts.”

Response to questions, part 1

Why are the proportion of activations routed and the routing probability used in the auxiliary loss? Intuitively, both convey the same information, with one being the distribution and the other samples from it.

Thank you for this excellent question. The purpose of the auxiliary loss is to encourage the model to send an equal number of activations to each expert in order to reduce overhead. Thus, we aim to drive the f-vector towards the uniform distribution. However, we cannot directly optimize this objective because the f-vector is not differentiable. So, we instead optimize the scaled dot-product between the f-vector and the P-vector, which is feasible because the P-vector is differentiable. Note that this auxiliary loss is not original to our work; it comes from [1].

A batch size of 8192 is used during training (line 269). Is a large batch size required for stable training of switch sparse autoencoders, and does this differ from the baseline methods?

Thank you for asking this question. A batch size of 8192 is typical for training sparse autoencoders. [2] similarly uses a batch size of 8192, [3] and [4] use a batch size of 4096 and [5] uses a much greater batch size of 131072.

Finally a broader question about the utility of scaling SAEs. Do the learned features need to be manually interpreted? If so, how is increasing the number of parameters a useful approach to interpretability if it also increases the required manual annotation?

Thank you for expressing this concern. In the right subplot of Fig. 1, we demonstrate that the gap between TopK and Switch SAE performance at a fixed width shrinks as we scale the number of parameters. For this reason, we expect that Switch SAEs will require a very small amount of additional features to achieve the same reconstruction loss on frontier language models.

In addition, there exists a separate line of work dedicated to efficient automatic interpretability. For example, recent work [6] open-sourced an automated pipeline for efficiently interpreting millions of features.

Thank you for your thoughtful comments. We greatly appreciate you taking the time to review our work.

Response to weaknesses

However, the architecture itself and the training technique are straightforward applications of existing work.

We agree that mixture-of-experts and sparse autoencoders are well established techniques in the literature. However, this work is the first to demonstrate that these ideas can be synergistically combined to enable more efficient large-scale dictionary learning on model activations. We believe our contributions are important in that they show not only scaling laws and performance tradeoffs, but also examine more subtle questions of feature duplication, geometry and interpretability.

The author conducts experiments comparing the Pareto frontiers of proposed Switch SAEs against a single SAE baseline, showing their better reconstruction performance at similar sparsity levels.

As a small note, in Fig. 3, we compare against three separate SAE baselines (ReLU SAE [1], Gated SAE [2] and TopK SAE [3]), not just one.

I think it is somewhat expected for the mixture model due to expert specialization and the additional router parameters.

Thank you for expressing this concern. The router parameters make up a very small proportion of the overall parameters. Across all our experiments, the router parameters make up between 0.002% and 0.3% of the total parameters. We apologize for not making this more clear and have added additional text to our submission to help clarify this. Thus, expert specialization is indeed the reason the Switch SAE performs better, but we do not think this was necessarily expected; it depends on how clusterable the model’s feature manifold is, and one of the main findings of this work (see e.g. Figure 5) is that empirically it is clusterable!

Moever, the training of this architecture is more complex and costly.

Thank you for bringing up this concern. In the left subplot of Fig. 1, we demonstrate that training a Switch SAE is approximately 10x less costly than training a TopK SAE with the same reconstruction error. The training procedure is also only slightly more complex, and thus we believe it will be well worth the associated efficiency gains in many cases. We also plan to release the non-anonymized link to our code once we are allowed to, which should make future re-implementations in modern SAE training libraries significantly easier.

The worse performance on the same width setting is indeed a drawback of this work, which might limit its applications.

We agree with you that this is a limitation of our current work. However, as shown in the right subplot of Fig. 1, the gap between TopK and Switch SAE performance at a fixed width shrinks as we scale the number of parameters. Thus, we expect that this drawback will be imperceptible at large scales and will not limit its applications. We apologize for not making this more clear and have added additional text to our submission to help clarify this.

In addition, we believe that future work might further decrease the gap in the same-width setting. In our conclusion, we cite “feature deduplication techniques, more sophisticated routing architectures, and multiple active experts” as possible directions.

Response to questions

Even though only one expert is active, I wonder what sparsity patterns look like for other experts and how much they overlap. Would it make sense to compare using global sparsity?

We apologize, but we are not sure what you mean here. We apply the TopK sparsity activation function only to the active expert. Do you mean examining the number of features that would activate if we did not apply the TopK function, and comparing the chosen and non-chosen experts? We are happy to run this experiment if you clarify!

If our comments resolve your concerns, we would appreciate it if you would consider raising your score!

[1] Towards Monosemanticity: Decomposing Language Models With Dictionary Learning, Bricken et al.

[2] Improving Dictionary Learning with Gated Sparse Autoencoders, Rajamanoharan et al.

[3] Scaling and evaluating sparse autoencoders, Gao et al.

Thank you for your thoughtful comments. We are glad that you enjoyed our paper!

Response to weaknesses

The switch SAEs require more features for the same reconstruction loss, which down the road: for automatic interpretability and SAE feature based interventions presents significant drawbacks. The cost of automatic annotation scales linearly with the number of features to annotate and in my opinion having a giant number of SAE features also makes them less human-compatible/interpretable.

Thank you for expressing this concern. In the right subplot of Fig. 1, we demonstrate that the gap between TopK and Switch SAE performance at a fixed width shrinks as we scale the number of parameters at a fixed sparsity. For this reason, we expect that Switch SAEs will require a very small amount of additional features to achieve the same reconstruction loss for the very wide SAEs trained on frontier language models.

In addition, there exists a separate line of work dedicated to efficient automatic interpretability. For example, recent work [1] open-sourced an automated pipeline for efficiently interpreting millions of features.

Response to questions

Do you think it is possible to set up switch SAE in a way that each SAE becomes a domain expert? E.g., there could be one SAE per language or something like that.

Thank you for asking this question. Typically, mixture-of-expert models are not cleanly interpretable (e.g., there isn’t a “chemistry” expert versus a “physics” expert). It may be possible to enforce that each SAE focuses on a different domain via a different router architecture, but in this work we focus on efficiency and feature quality. Other recent works, e.g. [2], [3], [4], also examine post-hoc clustering of SAE features; each cluster can then be seen as a domain expert.

[1] Automatically Interpreting Millions of Features in Large Language Models, Paulo et al.

[2] Not all language model features are linear, Engels et al.

[3] The Geometry of Concepts: Sparse Autoencoder Feature Structure, Li et al.

[4] HDBSCAN is Surprisingly Effective at Finding Interpretable Clusters of the SAE Decoder Matrix, Lim et al.