Efficient Dictionary Learning with Switch Sparse Autoencoders
Authors: Anish Mudide
Joshua Engels
Eric J Michaud
Max Tegmark
Christian Schroeder de Witt
Venue: ICLR 2025 Poster (2025)
Keywords: interpretability
mixture of experts
sparse autoencoders
Primary area: interpretability and explainable AI

Abstract: Sparse autoencoders (SAEs) are a recent technique for decomposing neural network activations into human-interpretable features. However, in order for SAEs to identify all features represented in frontier models, it will be necessary to scale them up to very high width, posing a computational challenge. In this work, we introduce Switch Sparse Autoencoders, a novel SAE architecture aimed at reducing the compute cost of training SAEs. Inspired by sparse mixture of experts models, Switch SAEs route activation vectors between smaller "expert" SAEs, enabling SAEs to efficiently scale to many more features. We present experiments comparing Switch SAEs with other SAE architectures, and find that Switch SAEs deliver a substantial Pareto improvement in the reconstruction vs. sparsity frontier for a given fixed training compute budget. We also study the geometry of features across experts, analyze features duplicated across experts, and verify that Switch SAE features are as interpretable as features found by other SAE architectures.

Published as a conference paper at ICLR 2025
EFFICIENT DICTIONARY LEARNING WITH SWITCH SPARSE AUTOENCODERS
Anish Mudide
Joshua Engels
Massachusetts Institute of Technology
Massachusetts Institute of Technology
Eric J. Michaud
Max Tegmark
Massachusetts Institute of Technology
Massachusetts Institute of Technology
Christian Schroeder de Witt University of Oxford
ABSTRACT
Sparse autoencoders (SAEs) are a recent technique for decomposing neural net- work activations into human-interpretable features. However, in order for SAEs to identify all features represented in frontier models, will be necessary to scale them up to very high width, posing computational challenge. In this work, we introduce Switch Sparse Autoencoders, novel SAE architecture aimed at re- ducing the compute cost of training SAEs. Inspired by sparse mixture of experts models, Switch SAFS route activation vectors between smaller "expert" SAFs, en- abling SAEs efficiently scale many more features. We present experiments comparing Switch SAES with other SAE architectures, and find that Switch SAES deliver substantial Pareto improvement in the reconstruction vs. sparsity frontier for a given fixed training compute budget. We also study the geometry of features across experts, analyze features duplicated across experts, and verify that Switch SAE features are as interpretable as features found by other SAE architectures.
INTRODUCTION
Recently, large language models have achieved impressive performance on many tasks (Brown al., 2020), but they remain largely inscrutable to humans. Mechanistic interpretability aims to open this metaphorical black box and rigorously explain model computations (Olah al., 2020). Specifi- cally, much work in mechanistic interpretability has focused on understanding features, the specific innereiiteerreeaable variables a model uses for computation (Olah et al., 2020: Park et al., 2023; Engels et al., 2024).
Early mechanistic attempts to understand features focused on neurons, but this work was stymied by the fact that neurons tend to be polysemantic: they are firequently activated by several completely dif- ferent types of inputs, making them hard to interpret (Olah et al., 2020). One theory for why neurons are polysemantic is superposition, the idea that language models represent many more concepts than they have available dimensions (Elhage et al., 2022). To minimize interference, the superposition hypothesis posits that features are encoded as almost orthogonal directions in the model's hidden state space.
Bricken et al. (2023) and Cunningham al. (2023) propose disentangle these superimposed model representations into monosemantic features by training unsupervised sparse autoencoders (Lee al., 2007: Le, 2013; Konda et al., 2014) on intermediate language model activations. The success of this technique has led to an explosion of recent work (Templeton et al., 2024; Gao et al., 2024) that has focused on scaling sparse autoencoders to frontier language models such as Claude 3 Sonnet (Anthropic, 2024a) and GPT-4 (Achiam et al., 2023). Despite scaling SAES to 34 million features, Templeton et al. (2024) estimate that there likely remains orders of magnitude more fea- tures. Furthermore, Gao et al. (2024) train SAES on series of language models and find that larger
*Correspondence amudide @ mit.cdu.

---

Published as a conference paper at ICLR 2025
4
(Tapk)
Switch SAE. experts
Reconstruction MSE
3
Switch SAC, experts
2
10a
10a
10
1013
1024
1015
10
1017
107
10a
FLOPs
Parameters
Figure 1: Scaling laws for Switch SAFs. We train dense Topk SAFs and Switch SAEs of varying size with fixed k 32. Left: Switch SAEs achieve better reconstruction than dense SAEs at fixed compute budget. Right: Switch SAEs require more features in total (and therefore more parameters) achieve the same reconstruction dense SAEs when trained convergence, although this gap narrows for larger SAES.
models require more features to achieve the same reconstruction error. As model sizes continue to grow, current training methodologies are quickly become computationally intractable.
Each training step of a sparse autoencoder generally consists six major computations: the en- coder forward pass, the encoder gradient, the decoder forward pass, the decoder gradient, the latent gradient and the pre-bias gradient. Gao et al. (2024) introduce kernels that leverage the sparsity of the TopK activation function dramatically optimize all computations except the encoder forward pass, which is not sparse. After implementing these optimizations, Gao et al. (2024) find that the training time is bottlenecked by the dense encoder forward pass and the memory is bottlenecked by storing the latent pre-activations.
In this work, we introduce the Switch Sparse Autoencoder, which our knowledge the first work to solve these dual memory and FLOP bottlenecks. The Switch SAE combines the Switch layer (Fedus et al., 2022) with the TopK SAE (Makhzani Frey. 2013; Gao et al., 2024). At a high level, the Switch SAE is composed of many small expert SAEs and a trainable routing network that determines which expert SAE will process given input. We demonstrate that the Switch SAE delivers improvement in training FLOPs vs. training loss over existing methods.
1.1 CONTRIBUTIONS
1. In Section 3, describe the Switch Sparse Autoencoder architecture and compare it to existing SAE architectures. We also describe our training methodology, which balances reconstruction and expert utilization.
2. In Section 4.1, describe scaling laws for reconstruction MSE with respect to FLOPs and parameters. We show that Switch SAEs achieve a lower MSE compared to TopK SAEs using the same amount of training compute.
3. In Section 4.2, benchmark Switch SAEs against ReLU, Gated and TopK SAEs the sparsity- reconstruction Pareto frontier.
In Section 4.3, we study feature duplication Switch SAEs and visualize the global structure of Switch SAE features using t-SNE.
5. In Section 4.4, demonstrate that Switch SAE features are interpretable as TopK SAE fea- tures.
2

---

Published as a conference paper at ICLR 2025
2 RELATED WORK
2.1 MIXTURE OF EXPERT MODELS
In a standard deep learning model, every parameter is used for every input. An alternative approach is conditional computation, where only a subset of the parameters are active depending on the input, allowing models with more parameters without the commensurate increase in computational cost. Jacobs et al. (1991) propose to train multiple networks, where each network is dedicated to disjoint subset of all possible inputs, Since this seminal work on mixture-of-experts, significant follow-up work has been dedicated to exploring different architectures and configurations (Jordan & Jacobs, 1994: Tresp. 2000; Collobert et al., 2001; Rasmussen & Ghahramani, 2001; Aljundi et al., 2017) 7). Shazeer al. (2017) introduce the Sparsely- Gated Mixture-of Soxtuure-ff-xpertss (MoE) layer. the first general purpose conditional computation architecture component. A Mfxture-of- layer consists of (1) a set of expert networks and (2) routing network that determines which experts should be active on a given input. Shazeer et al. (2017) use MoE to scale LSTMs to 137 billion parameters, surpassing the performance previous dense models on language modeling and machine translation benchmarks. Building this work, Fedus al. (2022) introduce the Switch layer, a simplification to the MoE layer which routes to just a single expert and thereby decreases computational cost and increases training stability. Fedus et al. (2022) use Switch layers in place of MLP layers to scale transformers to over a trillion parameters. Recent state of the art language models have used MoE layers, including Mixtral 8x7B (Jiang et al., 2024) and Grok-1 (xAI, 2024). To the best of our knowledge, we are the first to apply conditional computation to training SARS.
2.2 DEEP LEARNING TRAINING OPTIMIZATIONS
Our method fits into the wider literature focused on accelerating deep learning training. One type of training speedup uses hardware accelerators like GPUs (Raina et al., 2009) and TPUs (Jouppi et al., 2017) to optimize highly parallelizable dense matrix operations. Algorithmic improvements, on the other hand, consist of architectural or implementation tricks to speed up forward and backwards passes on fixed hardware. Techniques such as MoE layers (discussed above) and Slide (Chen al., 2020) utilize sparsity to only evaluate subset of the parameters for a given wide MLP layer. Other techniques such as Flash Attention (Dao et al., 2022) and Reformers (Kitaev et al., 2020) use hashing or structured matrices to reduce the time complexity of transformer's s attention mechanism. See Appendix in Dao et al. (2022) for a comprehensive overview of the literature on algorithmic training optimizations. Our work differs from these papers in that we are concerned with not only whether the training optimization results in a speedup, but also whether SAE quality is preserved.
2.3 SPARSE AUTOENCODERS AND IMPROVEMENTS
Sparse dictionary learning is the task of decomposing input data into sparse linear combination of basic elements called atoms, which together form a dictionary (Olshausen & Field, 1997; Elad, 2010). There exist wide variety of techniques for performing dictionary learning, such as the method of optimal directions (Engan et al., 1999) and K-SVD (Aharon et al., 2006). However, such methods are not scalable large language models and may not be faithful the models internals. As such, recent work has focused on applying sparse autoencoders (Lee et al., 2007; Le, 2013; Konda et al., 2014), simple approximation of sparse dictionary learning, language models.
Since the initial works proposing SAES to separate model representations (Cunningham et al., 2023; Bricken al., 2023). there have been many proposed improvements to the SAE architecture. These have included alternative activation functions like ProLu (Taggart, 2024). TopK (Makhzani & Frey, 2013; Gao al., 2024), and Batch-TopK (Bussmann et al., 2024), architectural changes to solve feature suppression caused by the LI penalty (Rajamanoharan et al., 2024a;b), and changes to the optimization objective itself (Braun et al., 2024). There are many metrics along which these papers evaluate their improvements, but the most common metric is Pareto frontier of SAE latent sparsity (measured as average LO), mean squared error, and feature interpretability; thus, these are the metrics focus in this paper. We also compare primarily against TopK SAEs (Gao et al., 2024) as baseline, as recent work (Anthropic, 2024c) has shown that these achieve state of the art results on these metrics.
3

---

Published as a conference paper at ICLR 2025
3
THE SWITCH SPARSE AUTOENCODER
3.1 BASELINE SPARSE AUTORNCODER
Sparse autoencoders are trained to reconstruct language model activations Rf by decomposing them into sparse linear combinations of M > d unit-length feature vectors f, f2.... f....., fas E R
SAE architectures consist of an encoder W eme MRx, decoder W dec E Rdx bias terms (e.g., byre E I) and an activation function. The TopK SAE (Gao et al., 2024) outputs reconstruction x of the input activation x, given by
2 = TopK (enneWXexxcc bpre))
(1)
& Wee bpre
(2)
The latent vector I represents how strongly each feature is activated. Since is sparse, the decoder forward pass can be optimized by suitable kernel. The loss is the reconstruction error
C= ||x x1 -
We additionally benchmark against the ReLU SAE (Anthropic, 2024b) and the Gated SAE (Raja- manoharan et al., 2024a). The ReLU SAE uses the ReLU activation function and applies an LI penalty to the feature activations to encourage sparsity. The Gated SAE avoids activation shrink- age (Wright & Sharkey, 2024) by separately determining which features should be active and how strongly activated they should be.
3.2 SWITCH SPARSE AUTOENCODER ARCHITECTURE
The Switch SAE consists of N expert SAES {E,), as well as routing network that determines which expert SAE should be used for a given input (Fig. 2). Each expert SAE E, resembles TopK SAE with no bias term:
(x Topk (Wencx)
(3)
The router, defined by trainable parameters W router E R xd and bionte E RRd computes probabil- ity distribution pxx) E N over the N experts given by p(x) = softmax Wrowter( byouter)) We route the imput x to the most probable expert 1" (x) arg max (x)... The output is given by,
= Pixx
bpre) bpre-
(4)
The Switch SAE thus avoids the dense W matrix multiplication by leveraging conditional com- putation. When comparing Switch SAES to existing SAE architectures, we evaluate two cases: (1) FLOP-matchedd Switch SAEs, which perform the same number of FLOPs per activation but increase Switch SAEs, which reduce the
the number of features by factor of N, and (2) FLOPs per activation by factor of N while keeping the number of features constant.
3.3 SWITCH SPARSE AUTOENCODER TRAINING
We train the Switch SAE end-to-end. When computing X, we weight (xx(x - bpre) by Pa (x) to make the router differentiable. We adopt many of the training strategies described Templeton et al. (2024) and Gao et al. (2024) (see Appendix A. for details).
The TopK SAE enforces sparsity via its activation function and thus directly optimizes the recon- struction MSE. Following Fedus et al. (2022), we train our Switch SAEs using a weighted com- bination the reconstruction MSE and an auxiliary loss which encourages the router to send an equal number of activations to each expert reduce overhead. For batch B with T activations, the auxiliary loss is given by Laxx N ΣN" Si P.. Si represents the proportion of activations that are routed to expert i, and P. represents the proportion of router probability that is assigned to expert i. Formally, f. & ExeB 1(i (x) i) and P. + ExEB Pi (x). The auxiliary loss is minimized when the batch of activations is routed uniformly across the N experts. The recon- struction loss is defined be Liesun + ExEB ||x xII3. Note that Lrewm x d. Let a represent tunable load balancing hyperparameter that scales the auxilliary loss. The total loss Lioal is given Laux. We optimize Cional using Adam (Kingma, 2014). We find that results are not overly sensitive to the value of a, but that 3 is reasonable default based on hyperparameter sweep (see Appendix A.2 for details).
by Lioall =
Lreson +
Our code can be found httpt:ttp////....... //g1 thu.. www.h.c.oo/meee. comd/ de/ sw/ taitcch_see
4

---

Published as a conference paper at ICLR 2025
Reconstructed Activation Vector
Pi- Ei (x) bges) + Dpr
Decoder
Decoder 2
Spurse Feature Activations
Encoder
Encoder
Expert
Expert
Expert
Ed (x)
Wide TopK (*)
(x) arg (x
p(x) softmax W rooser (X
brauter) ))
Router
X
Activation Vector
Figure 2: Switch Sparse Autoencoder Architecture. The router computes a probability distribution the expert SAEs and routes the input activation vector the expert with the highest probability. The figure depicts the architecture for d = 5, N = 3, M = 12.
4 RESULTS
We train sparse autoencoders on the residual stream activations of GPT-2 small, which have a di- mension of 768 (Radford et al., 2019). Early layers of language models handle detokenization and feature engineering, while later layers refine representations for next-token prediction (Lad al., 2024). In this work. we follow Gao et al. (2024) and train SAEs on activations from layer 8, which we expect hold rich feature representations not yet specialized for next-token prediction. Using text data from Open OeenWeeText (Gokaslan & Cohen, 2019), we train for 100K steps using a batch size of 8192, for total of approximately 820 million tokens.
4.1 SCALING LAWS FOR RECONSTRUCTION ERROR
We first study scaling laws for Switch SAEs, comparing them to dense TopK SABs at fixed sparsity k 32. In Fig. show scaling curves of reconstruction MSE error with respect both FLOPs and number of parameters for Switch SAEs with 8 and 64 experts. We find that Switch SAES have favorable scaling with respect FLOPs compared to dense TopK SAEs. In fact, Switch SAEs using ~ OOM less compute can often achieve the same reconstruction MSE as TopK SAEs. However, Switch SAEs perform worse fixed width relative to dense TopK SAES. Increasing the number of experts improves compute efficiency but reduces parameter efficiency. We hypothesize that this trade-off a result of features needing to be duplicated across multiple Switch SAE experts, which discuss in more detail later. Nevertheless, in the right subplot of Fig. 1, we show that the gap between TopK and Switch SAE performance at a fixed width shrinks as we scale the number of parameters. Thus, for large-scale experiments, this gap may be imperceptible; since this is the regime which efficient training is most useful, we believe that this is not a significant weakness of Switch SAES.
5

---

Published as a conference paper at ICLR 2025
FLOP-Matched
Width-Matched
TopK SAE
TopK SAE
ReLU SAE
ReLU SAE
Reconstruction MSE
Gated SAE
Gated SAE
Switch
Switch SAE. experts
Switch SAE. 32 experts
Switch SAE, experts
Switch SAE, 64 experts
experts
10a
10a
1 FLR (Fraction of Loss Recovered)
10a
102
10a
102
TopK
SAF
10%
ReLU SAE
10%
ReLLI SAE
Gated SAE
Gated SAE
experts
experts
Switch SAE experts
Switch SAE 32 experts
SAE experts
Switch SAL 128 experts
1%
1%
10a
10a
10a
10a
Sparsity (LO)
Sparsity (LO)
Figure 3: Pareto frontier of sparsity versus (top) reconstruction mean squared error and (bottom) FLR [fraction loss recovered]. FLOP-matched Switch SAEs Pareto-dominate TopK SAEs using the same amount of compute (left). Width-matched Switch SAES perform slightly worse than TopK SABs but Pareto-dominate ReLU SAEs while performing fewer FLOPs (right).

---

Published as a conference paper at ICLR 2025
4.2 SPARSITY VS. RECONSTRUCTION ERROR
We now study Switch SAE performance in the reconstruction error sparsity frontier. We bench- mark the Switch SAE against the ReLU SAE (Bricken et al., 2023), the Gated SAE (Rajamanoharan et al., 2024a) and the TopK SAE (Gao et al., 2024). We present results for two settings:
The ReLU, Gated and TopK SAES are trained with 32 768 24576 features. We train Switch SAEs with 2, 4 and 8 experts. Each expert of the Switch SAE with N experts has 24576 features, for a total of 24576 N features. The Switch SAE performs roughly the same number of FLOPs per activation compared to the TopK SAE.
FLOP-matched:
2. Width-matched: Each SAE is trained with 32 768 24576 features. We train Switch SAEs with 16, 32, 64 and 128 experts. Each expert of the Switch SAE with N experts has 24576 features. The Switch SAE performs roughly times fewer FLOPs per activation N
compared to the TopK SAE.
Note that the router parameters make up small, insignificant proportion of the total parameters. Across all our experiments, the router parameters make up between 0.002% and 0.3% of the total parameters. For a wide range of sparsity values (LO between and 128), we report the reconstruction MSE and the fraction of cross-entropy loss recovered (FLR) when the sparse autoencoder output is patched into the language model. A FLR value of corresponds to a perfect reconstruction, while a FLR value of 0 corresponds zero-ablation (setting the residual stream the zero vector).
4.2.1 FLOP-MATCHED RESULTS
For the shown in the left two plots of Fig. 3. The Switch SAEs are a Pareto improvement over the TopK, Gated and ReLU SAEs in terms of both MSE and FLR. 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. For a detailed discussion of FLOP calculations, see Appendix A.3.
experiments, we train Switch SAEs with 2, 4 and 8 experts, with the results
4.2.2 WIDTH-MATCHED RESULTS
For the width-matched experiments, we train Switch SAEs with 16, 32, 64 and 128 experts, with the results shown in the right two plots of Fig. 3. The Switch SAEs consistently underperform compared to the TopK SAE in terms of MSE and FLR, while for the most part outperforming Gated and ReLU SAEs. When LO is low, Switch SAEs perform particularly well. This suggests that the high frequency features that improve reconstruction fidelity the most for given activation might lie within the same cluster.
These results demonstrate that Switch SAEs can reduce the number FLOPs per activation by up to 128x while still retaining the performance of a ReLU SAE. We further believe that Switch SAES can likely achieve greater acceleration on larger language models.
To show the generality of Switch SAE training, in Appendix B we train on four layers of GPT-2 at residual, attention, and MLP outputs. We also train on a single layer of Gemma 2 2B (Team al., 2024). We find that Switch SAEs perform well at residual layers all positions and models tested, but worse on MLP and attention outputs.
4.3 FEATURE GEOMETRY
4.3.1 FEATURE SIMILARITY
We are interested in why Switch SAES achieve a worse reconstruction MSE than TopK SAEs of the same size. One hypothesis is that because any given forward pass only a single expert is activated, some experts will have to learn duplicate features, reducing SAE capacity (this is necessary since there is no perfect split of features into clusters such that features in different clusters never occur). We test this hypothesis using common (see e.g. Braun et al. (2024)) SAE evaluation metric: nearest neighbor cosine similarity.
7

---

Published as a conference paper at ICLR 2025
0.45
16
Experts
0.00
Expert
Expert
0.40
0.06
32
0.4
0.35
0.04
0.30
0.3
0.25
0.00
0.20
Sparsity (LO)
Sparsity (LO)
Expert
(a) Fraction of SAE decoder tors with nearest neighbor cosine sim 0.9, width-matched SAES.
(b) Average cosine sim with near- neighbor for decoder vectors,
(c) Average cosine sim between near- neighbors experts (16 perts, 64).
width-matched SAEs.
Figure 4: 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 LOs. Furthermore, between-expee seteenn-epeert similarities show that experts specialize moderately; expert O, for example, has low similarity with most other experts.
One way to use this metric to measure the number of duplicate features in an SAE is to measure the proportion of vectors in the decoder that have nearest neighbor above given cosine similarity, since highly similar decoder vectors are likely duplicates. In Fig. 4a, we compare SAEs trained with different numbers of experts and different sparsities on this measure with a threshold of 0.9, and find that soon as we train with experts this measure jumpy sharply from baseline of 0 in TopK SAEs to to percent. In other words, strict duplicates likely reduce Switch SAE capacity by up to 10%. This effect is less prevalent larger LOs. However, duplicates, but also general shift towards redundancy: in Fig. 4b, we find that a more global metric, cosine similarity of the nearest neighbor averaged across all features, shows similar pattern across sparsity and number of experts.
are not just worried about exact
Another measure of Switch SAE quality is expert specialization, that is, how similar the sets of features each expert learns. We quantify by averaging the cosine similarity of each feature in expert A with its nearest neighbor in expert B. On end, if the features are perfectly clustered, then each pair of experts should have the same cosine similarity as random blocks of the same size in TopK SAE of the same size. On the other end, no specialization should result in identical experts. In Fig. 4c, we plot this metric for all pairs of experts in a 16 expert Switch SAE with LO = 64. Some experts, e.g. expert O, seem unique, while other pairs of experts, e.g. 10 and are highly similar. All pairs are more similar than random blocks of the same size in TopK SAE, which has a value of about 0.2 under this same metric.
4.3.2 T-SNE VISUALIZATION
To visualize the global structure of SAE features, we show t-SNE projections of the encoder and decoder feature vectors in Fig. 5. We find that encoder features from the same expert cluster together, while decoder features tend to be more diffuse. In the encoder feature t-SNE projection, directly observe feature duplication around the periphery of the plot we find a variety of isolated points which upon closer inspection actually tight groupings of multiple features from different experts.
can also
4.4 AUTOMATED INTERPRETARILII
Bricken et al. (2023) evaluate how interpretable sparse autoencoder features using automated inter- pretability: they generate explanations for each feature by giving top activating examples for that feature language model, and then ask a language model (in new context) predict, using only the description the feature, on which tokens the described feature fires. More recently, Juang et al. (2024) introduce more compute efficient method that measures detection, whether the explanation allows language model to predict whether a feature fires at all on a given context. Juang et al. (2024) additionally measure detection each decile of feature activation, well contexts
8

---

Published as a conference paper at ICLR 2025
Encoder features
Decoder features
Expert
Expert
Expert
Expert Expert
Expert Expert
Expert
Figure 5: t-SNE projections of encoder and decoder features for Switch SAE with 65k total features and 8 experts. Encoder features appear to cluster together by expert. Away from these clusters, we variety of isolated points which in fact tight groups of features which have been duplicated across multiple experts. We do not observe as clear clusters in the decoder features.
where the feature does not fire (where to be correct model should report that the feature does not fire). Using the Juang et al. (2024) implementation, we find that FL.OP-mathhed SAE features have similar interpretability as TopK SAEs, but width-matched SAEs have a higher true positive rate contexts where the features fire. while at the same time having lower true negative rate (Fig. 6). We interpret this result as Switch SAEs being biased towards having duplicate frequent features, which may both be easier to detect and more prone to false positives.
FLOP-Matched
Width-Matched
1.0
Topk
Topk
0.9
Switch: 2e
Switch: 16e
Switch: 4e
Switch: 32e
0.8
Switch: Be
Switch: 64e
Switch: 128e
Accuracy
0.7
0.6
0.5
0.4
0.3
0.2
Not
88
Q10
Not
00
Quantiles
Quantiles
Figure 6: Automated interpretability detection results across SAE feature activation quantiles for 1000 random features, 95% confidence intervals shown. "Not" means that the context comes from a random set of text where the feature was not activating.
CONCLUSION
Switch SAEs promising approach towards scaling sparse autoencoders, they allow im- provement in the FLOPs vs. MSE scaling law without significant reduction in feature interpretabil- ity. We believe that Switch SAEs may find their best application for huge training large

---

Published as a conference paper at ICLR 2025
clusters of GPUs, since in this setting each expert be split to separate GPU, leading to significant wall clock training speed ups. Overall, we believe that this work provides a case study for applying existing machine learning training optimization techniques to the setting of sparse au- toencoders for feature extraction from language models; we hope that the investigations in this paper serve guide for adapting
such techniques.
Limitations: The most significant limitation of the Switch SAE is the reduction in performance of the SAE at fixed number of parameters (especially for attention and MLP layers, see Ap- pendix B); future work to bridge this gap might investigate feature deduplication techniques, more sophisticated routing architectures, and multiple active experts. Additionally, we investigate only the simplest possible mixture of experts architecture; while this allows us to focus on the question of whether and how mixture of experts training works all for sparse autoencoders, leaves open the question of the maximum performance gain mixture of experts style training might allow. Thus, future work could examine more sophisticated mixture of experts architectures like GShard (Lep- ikhin et al., 2020), DeepSeekMoE (Dai et al., 2024), and expert choice routing (Zhou et al., 2022) to further optimize performance. Finally, we do not study downstream evaluations of Switch SAEs in this paper; for instance, it is possible that the increased duplication of features between experts complicates feature steering or circuit discovery.
ACKNOWLEDGMENTS
We used the dictionary learning repository (Marks et al., 2024) to train our SAEs. We would like to thank Samuel Marks and Can Rager for advice on how use the repository. We would also like to thank Jacob Goldman--etzzede Achyuta Rajaram, Michael Pearce, Gitanjali Rao, Satvik Golechha, Kola Ayonrinde, Rupali Bhati, Louis Jaburi, Vedang Lad, Adam Karvonen, Shiva Mudide, Sandy Tanwisuth, JP Rivera and Juan Gil for helpful discussions. This work is supported by Erik Otto, Jaan Tallinn, the Rothberg Family Fund for Cognitive Science, the NSF Graduate Research Fellowship (Grant No. 2141064), IAIFI through NSF grant PHY-2019786. UKRI grant: Turing AI Fellow- ship EP/W002981 1/1, and Armasuisse Science+ Technology. AM was greatly helped by the MATS program, funded by AT Safety Support.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023
Michal Aharon, Michael Elad, and Alfred Bruckstein. K-svd: An algorithm for designing over- complete dictionaries for sparse representation. IEEE Transactions signal processing, 14(1)): 4311-4322, 2006.
Rahaf Aljundi, Punarjay Chakravarty, and Tinne Tuytelaars. Expert gate: Lifelong learning with a network of experts. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3366-3375, 2017.
Anthropic. The claude 3 model family: Opus, sonnet, haiku, 2024a. URL httpthttp:///wttt /Rthttps::/wwwcdd anthropic com @aeebbabbllee Card_Claude_3 .dff.
Transformer Circuits Team Anthropic. Transformer circuits april 2024 update: Update on
how we train saes, 2024b. URL https://transforrt april- ppril-update/se www.uupddeeeieeeeeee hwtp//nddxxecm///m html terlltttrat #training saes. Accessed: 2024-09-16.
Transformer Circuits Team Anthropic. Transformer circuits august 2024 update: Evals case study,
pwbppppb/2//0/ /2044/ httpa//0///guuuuuuatttttttt /ugpssttuuuate// 2024c. URL /ttt:s://taanssffrer--iiic wtt.tt...ubbbbb/ index. htmllmllllmmmm.mm -evals-case-stteeeeeeee wwweaaasuuuee Accessed: 2024-09-16.
Dan Braun, Jordan Taylor, Nicholas Goldowsky-Doe and Lee Sharkey. Identifying functionally important features with end-to-end sparse dictionary learning. arXiv preprint arXiv:2405. 12241, 2024.
10

---

Nelson Elhage, Tristan Hume, Catherine Olsson, Nicholas Schiefer, Tom Henighan, Shauma Kravec, Zac Hadfield-Dodds, Robert Lasenby, Dawn Drain, Carol Chen, Roger Grosse, Sam McCandlish, Jared Kaplan, Dario Amodei, Martin Wattenberg, and Christopher Olah. Toy models of super- pub/ Published as a conference paper at ICLR 2025
Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jemmyn, Tom Con- erly, Nick Tumne, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac HatfieldDoddss, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah E Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https//transformer- circuits. circuits.pub/2022. http://00/mmooommmmatefeerreeeee
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz. Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Rad- ford, Ilya Sutskever, and Dario Amodei. Language Models are Few-Shot Learners. In Ad- vances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Asso- ciates, Inc., 2020. URL httpsss/pprooeedigggg neuripp cc wwc//appeeee httpppppppp/p/aaaaee http///2220/h/////// html.
Bart Bussmann, Patrick Leask, and Neel Nanda. BatchTTpp:: A simple improvement for TopK- a fomm SAEs. Al Aligmment Forum, 2024. URL immpplee-imrrooeeeeett for- worr-topk- saes.
Wkx6yW2bbAs htt.c//////accttttp....o.
Beidi Chen, Tharun Medini, James Farwell, Charlie Tai, Anshumali Shrivastava, et al. Slide: In defense of smart algorithms over hardware acceleration for large-scale deep learning systems. Proceedings of Machine Learning and Systems, 2:291-306, 2020.
Ronan Collobert, Samy Bengio, and Yoshua Bengio. A parallel mixture of svms for very large scale problems. Advances in Neural Information Processing Systems, 14, 2001.
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen- coders find highly interpretable features in language models. arXiv preprint arXiv:2309.0866 086000 2023.
Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. Deepseekmoe: Towards ultimate expert specialization in mixture- of-experts language models. arXiv preprint MxXiv:2401.00660 2024.
Tri Dao. Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
Michael Elad. Sparse and redundant representations: from theory to applications in signal and image processing. Springer Science & Business Media, 2010.
position. Transformer Circuits Thread, 2022. ht 2002/toy_model/ince h/iddx html.
Kjersti Engan, Sven Ole Aase, and J Hakon Husoy. Method of optimal directions for frame design. In 1999 IEEE International Conference on Acoustics, Speech, and Signal Processing. Proceed- ings. ICASSP99 (Cat. No. 99CH36258), volume 5, pp. 2443-2446. IEEE, 1999.
Joshua Engels, Isaac Liao, Eric J Michaud, Wes Gurnee, and Max Tegmark. Not all language model features are linear. arXiv preprint arXiv:2405. 14860, 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
Leo Gao, Tom Dupré la Tour, Henk Tillman, Gabriel Goh, Rajan Troll, Alec Radford, Ilya Sutskever, Jan Leike, and Jeffrey Wu. Scaling and evaluating sparse autoencoders, 2024. URL https: / //axxiv.org/aaaaa http://///aa........./ 04093.
11

---

Published as a conference paper at ICLR 2025
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. httppp/sttt//////ooo00.....htu 1o/ OpenWebTextCozpus, 2019.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-88) 87, 1991.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bam- ford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral experts. arXiv preprint arXiv:2401.04088 2024.
Michael Jordan and Robert A Jacobs. Hierarchical mixtures of experts and the em algorithm. Neural computation, 6(2):181-214, 1994.
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of a tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pp. 1-12, 2017.
Caden Juang, Gonçalo Paulo, Jacob Drori, and Nora Belrose. Open source automated interpretability for sparse autoencoder features. https: /blog. e eleuther a http//uuonneerrree 2024.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint #rXiv:2001.04451, 2020.
Kishore Konda, Roland Memisevic, and David Krueger. Zero-bias autoencoders and the benefits of co-adapting features. arXiv preprint arXiv:1402.333: 2014.
Vedang Lad, Wes Gumnee, and Max Tegmark. The remarkable robustness of Ilms: Stages of infer- ence? arXiv preprint arXiv:2406. 19384, 2024.
Quoe Le. Building high-level features using large scale unsupervised learning. In 2013 IEEE international conference on acoustics, speech and signal processing, pp. 8595-8598. IEEE, 2013.
Honglak Lee, Chaitanya Ekanadham, and Andrew Ng. Sparse deep belief net model for visual area v2. Advances in neural information processing systems, 20, 2007.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 2020.
Alireza Makhzani and Brendan Frey. K-sparse autoencoders. arXiv preprint arXiv: 2013.
Samuel Marks, Adam Karvonen, and Aaron Mueller. dictionary learning. https: / github. 2024.
com/ http///pnccckk/dcciiiorar.eeaaaa
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to cireuits. Distill, 2020. doi: 10.23915/disilll.0004..00 hutps:/distill. hwtpa:iii....aboo00oooouuu//ooo.
Bruno A Olshausen and David I Field. Sparse coding with an overcomplete basis set: A strategy employed by v1? Vision research, 37(23):3311 -325, 1997.
Kiho Park, Yo Joong Choe, and Victor Veitch. The linear representation hypothesis and the geometry of large language models. arXiv preprint arXiv:2311.03653. 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Rajat Raina, Anand Madhavan, and Andrew Y Ng. Large-scale deep unsupervised learning using graphics processors. In Proceedings of the 26th annual international conference on machine learning, pp. 873-880, 2009.
12

---

Published as a conference paper at ICLR 2025
Senthooran Rajamanoharan, Arthur Conmy, Lewis Smith, Tom Lieberum, Vikrant Varma, János Kramár, Rohin Shah, and Neel Nanda. Improving dictionary learning with gated sparse autoen- coders. arXiv preprint arXiv:2404.1611, 2024a.
Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, and Neel Nanda. Jumping ahead: Improving reconstruction fidelity with jumprelu sparse autoencoders. arXiv preprint arXiv:2407. 14435, 2024b.
Carl Rasmussen and Zoubin Ghahramani. Infinite mixtures of gaussian process experts. Advances in neural information processing systems, 14, 2001.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 010117001...8,,, 2017.
Glen Taggart ProLU: A nonlinearity for sparse autoencoders. Al Alignment Forum, 2024. URL https: /wwp://wwwwwt/.o// aliggment prolu-a-pareto meexxeto-improveeeeet improvement -tw.reeeeetttooreeeeeeeeee -oree-oorrreseeeeeeee autarscceat
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju, Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models a practical size. arXiv preprint 2xXiv:2408.001 2024.
Adly Templeton. Tom Conerly, Jonathan Marcus, Jack Lindsey. Trenton Bricken, Brian Chen, Adam Pearce, Craig Citro, Emmanuel Ameisen, Andy Jones, Hoagy Cunningham, Nicholas I. Turner, Callum McDougall, Monte MacDiarmid, C. Daniel Freeman, Theodore R. Sumers, Edward Rees, Joshua Batson, Adam Jermyn, Shan Carter, Chris Olah, and Tom Henighan. Scaling monosemanticity: Extracting interpretable features from claude sonnet. Trans- former Circuits Thread, 2024. URL https:/ /tte:://rrre scallingmoooosmmattiiiiii index.ndxx..iiii html.
pub/2024/
Volker Tresp. Mixtures of gaussian processes. Advances in neural information processing systems, 13, 2000.
Benjamin Wright and Lee Sharkey. Addressing feature suppression in SAEs. Al Aligument Forum, 2024. URL https: wwwp://wwwww/..o al http://wwwl.gmmntffrrr... addressing- feature suppression in in-siiiooniinissss------ saes.
xAI. Grok-1 model card, 2024. URL https::/tt..... /tlpa////////gog/////eeeeaaaa
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
A IMPLEMENTATION DETAILS
A.1 SWITCH SPARSE AUTOENCODER TRAINING DETAILS
We initialize the rows of W ene to be parallel the columns of W dec for all i. We initialize both Dyre and byaatr to the geometric median of batch of samples, but we do not tie bye and broute-- We additionally normalize the decoder column vectors unit-norm initialization and after each gradient step. We remove gradient information parallel to the decoder feature directions to mini- mize interference with the Adam optimizer. We the learning rate based on the 1 VM scaling law from Gao et al. (2024) and linearly decay the learning rate over the last 20% of training. We do not include neuron resampling (Bricken et al., 2023) or the AuxK loss (Gao et al., 2024), but future work could explore employing these strategies to prevent dead latents in larger models. We optimize Lusa with Adam using the default parameters 0.9. B2 0.999.
13

---

Published as a conference paper at ICLR 2025
Switch SAE, 16 experts
1.25 10a
Switch SAE, 32 experts
Reconstruction MSE
1.2 103
1.15 103
1.1 10a
1.05 103
10
10
10
10
10a
102
Figure 7: Hyperparameter sweep for a. We train fixed-width Switch SAEs with 16 and experts on activations from GPT-2 small. The sparsity level is set to L0=64.
A.2 HYPERPARAMETER SEARCH
In our objective function, the balance between reconstruction error and the auxiliary loss (which encourages the router to send an equal number of activations to each expert) is controlled by the hyperparameter a. We train 16-expert and 32-expert Switch SAEs at fixed sparsity (LO 64) and fixed width (24576 features), but sweep across the values [0.001, 0.003, 0.01, 0.03, 0.1, 0.3, 1, 3, 10, 30, 100]. We find that the reconstruction MSE is not 3 performs the best in both the 16-expert and
overly sensitive the value of and that the 32-expert settings (Fig. 7). We default to a = for the rest of our experiments.
A.3 FLOP CALCULATIONS
single training step of deep neural network requires roughly 6 P FLOPs, where P the number of active parameters and D is the batch size. Let d be the dimension of the language model activations, M be the total number of features and N be the number of experts. For TopK SAE, P = d. For a Switch SAE, P = dN + 2Md = 2Md since the router parameters are negligible. Thus, we have
N
FLOPS Topk (d, 6 (2Md) D,
FLOPS (d, M, 6 (aN 2Md
(2Md)
D 22
D.
N
B TRAINING ON ADDITIONAL SITES AND MODELS
In Table I, we show the results of training on GPT-2 layers 2, 4, 8, and 12 on MLP outputs, attention outputs, and the residual stream. We also train single SAE Gemma 2B (Team al., 2024). All SAES are trained with & = 64 with fixed width; we train the Gemma 2 2B SAES with width 65536 and the GPT-2 SAEs with width 24576. We use experts for the Switch SAES. As the main paper, we find that the decrease in SAE performance for fixed width minimally effects residual stream training; however, Switch SAEs trained on MLP and attention outputs do suffer significant performance reduction.
14

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of Topk and Switch SAEs across different models, layers and component
types. FVE = Fraction of Variance Explained, FLR Fraction of Loss Recovered.
Model
Layer
Type
TopK FVE
Switch FVE
TopK FLR
Switch FLR
GPT-2
2
resid
0.999
0.998
0.997
0.996
GPT-2
2
attn
0.919
0.884
0.997
0.968
GPT-2
2
mip
0.999
0.998
0.888
0.701
GPT-2
4
resid
0.997
0,997
0,996
0.995
GPT-2
4
attn
0,849
0,805
0.925
0.883
GPT-2
4
mlp
0.877
0,819
0.853
0.777
GPT-2
8
resid
0,989
0.987
0.994
0.993
GPT-2
8
attn
0,846
0,803
0,938
0.848
GPT-2
8
mip
0.782
0.733
0.868
0.801
GPT-2
10
resid
0.973
0.970
0.982
0.977
GPT-2
10
attn
0.892
0,863
0.942
0.883
GPT-2
10
mlp
0.855
0,828
0.859
0,812
Gemma 2 2B
12
resid
0.961
0.955
0.991
0.989
15


Reviews:
[{"review_id": "P1KDvwY0Zj", "reviewer": "Reviewer_GyB2", "summary": "This paper introduces Switch sparse autoencoders for learning sparse representations of neural network activations. Switch SAEs include a mixture of TopK SAEs and a router to ensure that only one expert processes each input.", "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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "This work presents an idea to apply a mixture-of-experts layer to sparse autoencoder training and propose Switch sparse autoencoders. However, the architecture itself and the training technique are straightforward applications of existing work. \n\nThe 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. I think it is somewhat expected for the mixture model due to expert specialization and the additional router parameters. Moever, the training of this architecture is more complex and costly. \n\nThe worse performance on the same width setting is indeed a drawback of this work, which might limit its applications.", "weaknesses": "See above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mSTbwGQpp6", "reviewer": "Reviewer_uAzQ", "summary": "The paper introduces switch SAE's which is a combination of switch layers with top-k SAEs. Instead of decomposing a latent directly it is forwarded to one of multiple SAE encoder decoder pairs. The authors show that given a FLOP budget the switch SAE reaches a lower reconstruction loss faster, while for a fixed number of parameters the vanilla top-k SAEs achieve lower reconstruction loss. That is for a low reconstruction loss with switch SAEs more features in total are required (but they are cheaper to train). Switch SAE features are shown to also be interpretable.", "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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors explore a combination of methods that makes a lot of sense (switch layers and top-k SAEs) and analyze the resulting method in a rigorous and convincing way. They make the interesting findings that switch SAE's are in terms of FLOPS actually cheaper to train (for a given target reconstruction loss) and at the same time require more features in total (which will be an additional challenge for the usual subsequent automatic interpretation of the learnt features). Also they show that the switch SAE features are interpretable.", "weaknesses": "The idea maybe can be considered as a little bit incremental, but given the very good execution of the paper I don't think that's a big problem.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7QyVGC4Hlx", "reviewer": "Reviewer_g9rJ", "summary": "This paper proposes an extension to Sparse Autoencoders (SAEs) in the form of switch SAEs. Switch SAEs reduce the memory consumption of storing the sparse activations by splitting the representation space across multiple SAE experts, of which only one is used at a time. Experiments show that switch SAEs achieve better reconstructions with the same number of FLOPs as current methods at the cost parameter efficency. The paper also explores the features learned by the experts and observe that some features cluster within experts and others are duplicated across them.", "questions": "- 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.\n- 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?\n- 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?\n- 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?\n\n**Minor suggestions that do not individually affect the score**\n- Line 295: “For a wide range of sparsity (L0) values”. Specify the range.\n- Line 298: Clarify what is meant by a zero-ablation.\n- Line 317: “Switch SAEs can likely achieve greater acceleration on larger language models”. Expand on this.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The strengths of the paper are its simple approach and side studies.\n- The approach enables a trade-off between model parameters and computational requirements which should be useful to practitioners.\n- The auxiliary experiments are well-motivated and support the author’s hypotheses. Specifically, the discussions of redundant representations accompanied by t-SNE visualizations are illuminating.\n- The limitations (parameter efficiency and feature duplication) are discussed in an informative way.", "weaknesses": "The main weakness of the paper is its quantitative evaluation.\n- 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?\n- 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.\n- Figure 4 could use an explanatory caption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "J6xM2pBHg6", "reviewer": "Reviewer_x3AE", "summary": "This paper introduces **Switch Sparse Autoencoders** (Switch SAEs). Switch SAEs are a new SAE variant, which is a popular area of research in interpretability. Switch SAEs are more efficiently to train and run at inference time for similar reasons to why Mixture-of-Expert (MoE) LLMs are more efficient than dense LLMs.\n\nThis paper claims that, fixing the number of FLOPs, Switch SAEs pareto dominate several leading SAE variants. Fixing the total number of features SAEs have however, Switch SAEs seem less performant. The paper then shows Switch SAEs have as interpretable metrics as dense SAEs as measured by an automated metric. There are some surface level studies into feature duplication, suggesting that this sometimes happens. All experiments are on GPT-2 Small's residual stream at layer 8.\n\nHere are my references for all sections in this review:\n\n[1] Switch MoE: https://arxiv.org/abs/2101.03961\n\n[2] DeepSeek MoE: https://arxiv.org/abs/2401.06066\n\n[3] GShard: https://arxiv.org/abs/2006.16668\n\n[4] Expert-choice; https://arxiv.org/abs/2202.09368\n\n[5] Ghost Grads: https://transformer-circuits.pub/2024/jan-update#dict-learning-resampling\n\n[6] Updated Anthropic training method: https://transformer-circuits.pub/2024/april-update/index.html#training-saes\n\n[7] Gated SAEs: https://arxiv.org/abs/2404.16014\n\n[8] JumpReLU SAEs: https://arxiv.org/abs/2407.14435\n\n[9] Anthropic Set Of Metrics Where JumpReLU does not reproduce: https://transformer-circuits.pub/2024/august-\nupdate/index.html#interp-evals\n\n[10] Scaling Monosemanticity: https://transformer-circuits.pub/2024/scaling-monosemanticity/\n\n[11] SFC: https://arxiv.org/abs/2403.19647", "questions": "1. 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\nexperts and represent more features, performance continues to increase while keeping computational\ncosts and memory costs (from storing the pre-activations) roughly constant\" are very confusing.\n\n(My GUESS at what that sentence is gesturing at is that fixing Switch SAE width, increasing number of experts decreases the encoder forward pass FLOP but not the backward pass FLOP. And after thinking about it for a minute, I think I see how increasing number of experts doesn't increase FLOP so long as you only route acts to 1 expert, always. But this was a non-trivial inference and I think it should be in the paper)\n\n2. 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.\n\n3. Do you have any experiments on different models or different sites?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **This paper addresses an important and novel direction that may be impactful for future work.** There are already several very large training runs that have trained large SAEs, and a key bottleneck is the autoencoder training (getting and potentially storing billions of activations from frontier models is also a bottleneck, however). The paper provides promising evidence for the MoE SAE approach.\n\n2. **The evaluation is clear and relatively comprehensive**. Addressing FLOP-matched results and feature interpretability compared to baselines is very clear.\n\n3. **The authors do sensible sanity checks of the approach.** Duplicated features could be a serious issue, and the analysis of how frequent they were showed they weren't an enormous proportion of features, but still present. I like the transparency of this section.", "weaknesses": "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.\n\n2. **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.\n\nAdditionally, 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.\n\nI would like to hear why these design choices were made, and think the paper should be edited if accepted to the conference to discuss these variants for future work, and to hedge the conclusions more.\n\n3. **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).\n\n4. **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.\n\nNote that, despite how I found more strengths than weaknesses, these weaknesses are more technical and subtle. I think the paper overall is strong, and it required harder thinking to realise some issues with it.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)