Efficient Dictionary Learning with Switch Sparse Autoencoders

## Abstract

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.

## 1 Introduction

Recently, large language models have achieved impressive performance on many tasks (Brown et al., ), but they remain largely inscrutable to humans. Mechanistic interpretability aims to open this metaphorical black box and rigorously explain model computations (Olah et al., ). Specifically, much work in mechanistic interpretability has focused on understanding features, the specific human-interpretable variables a model uses for computation (Olah et al., ; Park et al., ; Engels et al., ).

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 frequently activated by several completely different types of inputs, making them hard to interpret (Olah et al., ). 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., ). 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. () and Cunningham et al. () propose to disentangle these superimposed model representations into monosemantic features by training unsupervised sparse autoencoders (SAEs) on intermediate language model activations. The success of this technique has led to an explosion of recent work (Templeton et al., ; Gao et al., ) that has focused on scaling sparse autoencoders to frontier language models such as Claude 3 Sonnet (Anthropic, ) and GPT-4 (Achiam et al., ). Despite scaling SAEs to 34 million features, Templeton et al. () estimate that there likely remains orders of magnitude more features. Furthermore, Gao et al. () train SAEs on a series of language models and find that larger models require more features to achieve the same reconstruction error. As model sizes continue to grow, current training methodologies are set to quickly become computationally intractable.

Each training step of a sparse autoencoder generally consists of six major computations: the encoder forward pass, the encoder gradient, the decoder forward pass, the decoder gradient, the latent gradient and the pre-bias gradient. Gao et al. () introduce kernels that leverage the sparsity of the TopK activation function to dramatically optimize all computations except the encoder forward pass, which is not sparse. After implementing these optimizations, Gao et al. () 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 to our knowledge is the first work to solve these dual memory and FLOP bottlenecks. The Switch SAE combines the Switch layer (Fedus et al., ) with the TopK SAE (Gao et al., ). 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 a given input. We demonstrate that the Switch SAE delivers an improvement in training FLOPs vs. training loss over existing methods.

### 1.1 Contributions

- 1.
In Section 3, we 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, we 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, we benchmark Switch SAEs against ReLU, Gated and TopK SAEs on the sparsity-reconstruction Pareto frontier.
- 4.
In Section 4.3, we study feature duplication in Switch SAEs and visualize the global structure of Switch SAE features using t-SNE.
- 5.
In Section 4.4, we demonstrate that Switch SAE features are as interpretable as TopK SAE features.

Figure: Figure 1: Scaling laws for Switch SAEs. We train dense TopK SAEs and Switch SAEs of varying size with fixed $k=32$. Switch SAEs achieve better reconstruction than dense SAEs at a fixed compute budget, but require more features in total (and therefore more parameters) to achieve the same reconstruction as dense SAEs when trained for 100k steps (approximately to convergence).
Refer to caption: x1.png

## 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. Shazeer et al. () introduce the Sparsely-Gated Mixture-of-Experts (MoE) layer, the first general purpose conditional computation architecture. A Mixture-of-Experts layer consists of (1) a set of expert networks and (2) a routing network that determines which experts should be active on a given input.
Shazeer et al. () use MoE to scale LSTMs to 137 billion parameters, surpassing the performance of previous dense models on language modeling and machine translation benchmarks. Building on this work, Fedus et al. () 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. () 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., ) and Grok-1 (xAI, ). To the best of our knowledge, we are the first to apply conditional computation to training SAEs.

### 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., ) and TPUs (Jouppi et al., ) 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 et al., ) utilize sparsity to only evaluate a subset of the parameters for a given wide MLP layer. Other techniques such as Flash Attention (Dao et al., ) and Reformers (Kitaev et al., ) use hashing or structured matrices to reduce the time complexity of a transformer’s attention mechanism. See Appendix A in Dao et al. () 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 Autoencoder Optimizations

Since the initial works proposing SAEs to separate model representations (Cunningham et al., ; Bricken et al., ), there have been many proposed improvements to the SAE architecture. These have included alternative activation functions like ProLu (Taggart, ), TopK (Gao et al., ), and Batch-TopK (Bussmann et al., ), architectural changes to solve feature suppression caused by the L1 penalty (Rajamanoharan et al., ; ), and changes to the optimization objective itself (Braun et al., ). There are many metrics along which these papers evaluate their improvements, but the most common metric is a Pareto frontier of SAE latent sparsity (measured as average L0), mean squared error, and feature interpretability; thus, these are the metrics we focus on in this paper. We also compare primarily against TopK SAEs (Gao et al., ) as our baseline, as recent work (Anthropic, ) has shown that these achieve state of the art results on these metrics.

## 3 The Switch Sparse Autoencoder

### 3.1 Baseline Sparse Autoencoder

Sparse autoencoders are trained to reconstruct language model activations $\mathbf{x}\in\mathbb{R}^{d}$ by decomposing them into sparse linear combinations of $M\gg d$ unit-length feature vectors $\mathbf{f}_{1},\mathbf{f}_{2},\ldots,\mathbf{f}_{M}\in\mathbb{R}^{d}$. SAE architectures consist of an encoder $\mathbf{W}_{\text{enc}}\in\mathbb{R}^{M\times d}$, a decoder $\mathbf{W}_{\text{dec}}\in\mathbb{R}^{d\times M}$, bias terms (e.g., $\mathbf{b}_{\text{pre}}\in\mathbb{R}^{d}$) and an activation function. The TopK SAE (Gao et al., ) outputs a reconstruction $\hat{\mathbf{x}}$ of the input activation $\mathbf{x}$, given by

$$ $\displaystyle\mathbf{z}$ $\displaystyle=$ $\displaystyle\text{TopK}(\mathbf{W}_{\text{enc}}(\mathbf{x}-\mathbf{b}_{\text{ pre}}))$ (1) $\displaystyle\hat{\mathbf{x}}$ $\displaystyle=$ $\displaystyle\mathbf{W}_{\text{dec}}\mathbf{z}+\mathbf{b}_{\text{pre}}$ (2) $$

The latent vector $\mathbf{z}\in\mathbb{R}^{M}$ represents how strongly each feature is activated. Since $\mathbf{z}$ is sparse, the decoder forward pass can be optimized by a suitable kernel. The loss is the reconstruction error $\mathcal{L}=\|\mathbf{x}-\hat{\mathbf{x}}\|_{2}^{2}$.

We additionally benchmark against the ReLU SAE (Anthropic, ) and the Gated SAE (Rajamanoharan et al., ). The ReLU SAE uses the ReLU activation function and applies an L1 penalty to the feature activations to encourage sparsity. The Gated SAE avoids activation shrinkage (Wright & Sharkey, ) 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_{i}\}_{i=1}^{N}$ as well as a routing network that determines which expert SAE should be used for a given input (Fig. 2). Each expert SAE $E_{i}$ resembles a TopK SAE with no bias term:

$$ $E_{i}(\mathbf{x})=\mathbf{W}_{i}^{\text{dec}}\text{TopK}(\mathbf{W}_{i}^{\text {enc}}\mathbf{x})$ (3) $$

The router, defined by trainable parameters $\mathbf{W}_{\text{router}}\in\mathbb{R}^{N\times d}$ and $\mathbf{b}_{\text{router}}\in\mathbb{R}^{d}$, computes a probability distribution $\mathbf{p}(\mathbf{x})\in\mathbb{R}^{N}$ over the $N$ experts given by $\mathbf{p}(\mathbf{x})=\text{softmax}(\mathbf{W}_{\text{router}}(\mathbf{x}-
\mathbf{b}_{\text{router}}))$. We route the input $\mathbf{x}$ to the most probable expert $i^{*}(\mathbf{x})=\underset{i}{\arg\max}\ p_{i}(\mathbf{x})$. The output $\hat{\mathbf{x}}$ is given by,

$$ $\hat{\mathbf{x}}=p_{i^{*}(\mathbf{x})}E_{i^{*}(\mathbf{x})}(\mathbf{x}-\mathbf {b}_{\text{pre}})+\mathbf{b}_{\text{pre}}.$ (4) $$

The Switch SAE thus avoids the dense $\mathbf{W}_{\text{enc}}$ matrix multiplication by leveraging conditional computation. When comparing Switch SAEs to existing SAE architectures, we evaluate two cases: (1) FLOP-matched Switch SAEs, which perform the same number of FLOPs per activation but increase the number of features by a factor of $N$, and (2) width-matched Switch SAEs, which reduce the FLOPs per activation by a factor of $N$ while keeping the number of features constant.

Figure: Figure 2: Switch Sparse Autoencoder Architecture. The router computes a probability distribution over the expert SAEs and routes the input activation vector to the expert with the highest probability. The figure depicts the architecture for $d=5$, $N=3$, $M=12$.
Refer to caption: x2.png

### 3.3 Switch Sparse Autoencoder Training

We simultaneously train the router and all the expert sparse autoencoders. When computing $\hat{\mathbf{x}}$, we weight $E_{i^{*}(\mathbf{x})}(\mathbf{x}-\mathbf{b}_{\text{pre}})$ by $p_{i^{*}(\mathbf{x})}$ to make the router differentiable. We adopt many of the training strategies described in Templeton et al. () and Gao et al. () (see Appendix A for details).(^1^11Our code can be found at [https://github.com/amudide/switch_sae](https://github.com/amudide/switch_sae))

The TopK SAE enforces sparsity via its activation function and thus directly optimizes the reconstruction MSE. Following Fedus et al. (), we train our Switch SAEs using a weighted combination of the reconstruction MSE and an auxiliary loss which encourages the router to send an equal number of activations to each expert to reduce overhead. For a batch $\mathcal{B}$ with $T$ activations, the auxiliary loss is given by
$\mathcal{L}_{\text{aux}}=N\cdot\sum_{i=1}^{N}f_{i}\cdot P_{i}$. $f_{i}$ represents the proportion of activations that are routed to expert $i$, and $P_{i}$ represents the proportion of router probability that is assigned to expert $i$. Formally,
$f_{i}=\frac{1}{T}\sum_{\mathbf{x}\in\mathcal{B}}\mathbbm{1}{\{i^{*}(\mathbf{x}
)=i\}}$ and $P_{i}=\frac{1}{T}\sum_{\mathbf{x}\in\mathcal{B}}p_{i}(\mathbf{x})$. The auxiliary loss is minimized when the batch of activations is routed uniformly across the $N$ experts. The reconstruction loss is defined to be $\mathcal{L}_{\text{recon}}=\frac{1}{T}\sum_{\mathbf{x}\in\mathcal{B}}\|\mathbf
{x}-\hat{\mathbf{x}}\|_{2}^{2}$. Note that $\mathcal{L}_{\text{recon}}\propto d$. Let $\alpha$ represent a tunable load balancing hyperparameter that scales the auxilliary loss. The total loss $\mathcal{L}_{\text{total}}$ is given by
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{recon}}+\alpha\cdot d\cdot
\mathcal{L}_{\text{aux}}$. We optimize $\mathcal{L}_{\text{total}}$ using Adam (Kingma, ). In our experiments, we set $\alpha=3$.

## 4 Results

Figure: Figure 3: Pareto frontier of sparsity versus (top) reconstruction mean squared error and (bottom) 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 SAEs but Pareto-dominate ReLU SAEs while performing fewer FLOPs (right).
Refer to caption: x3.png

We train sparse autoencoders on the residual stream activations of GPT-2 small (Radford et al., ), which have a dimension of $768$. Early layers of language models handle detokenization and feature engineering, while later layers refine representations for next-token prediction (Lad et al., ). In this work, we follow Gao et al. () and train SAEs on activations from layer 8, which we expect to hold rich feature representations not yet specialized for next-token prediction. Using text data from OpenWebText (Gokaslan & Cohen, ), we train for 100K steps using a batch size of 8192, for a 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 SAEs at fixed sparsity $k=32$. In Fig. 1, we show scaling curves of reconstruction MSE error with respect to both FLOPs and number of parameters for Switch SAEs with 8 and 64 experts. We find that Switch SAEs have favorable scaling with respect to FLOPs compared to dense TopK SAEs. In fact, Switch SAEs using $\sim$1 OOM less compute can often achieve the same reconstruction MSE as TopK SAEs. However, Switch SAEs perform worse at 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 is a result of features needing to be duplicated across multiple Switch SAE experts, which we discuss in more detail later.

### 4.2 Sparsity vs. reconstruction error

We now study Switch SAE performance in the reconstruction error vs. sparsity frontier. We benchmark the Switch SAE against the ReLU SAE (Anthropic, ), the Gated SAE (Rajamanoharan et al., ) and the TopK SAE (Gao et al., ). We present results for two settings:

- 1.
FLOP-matched: The ReLU, Gated and TopK SAEs are trained with $32\cdot 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\cdot N$ features. The Switch SAE performs roughly the same number of FLOPs per activation compared to the TopK SAE.
- 2.
Width-matched: Each SAE is trained with $32\cdot 768=24576$ features. We train Switch SAEs with 16, 32, 64 and 128 experts. Each expert of the Switch SAE with $N$ experts has $\frac{24576}{N}$ features. The Switch SAE performs roughly $N$ times fewer FLOPs per activation compared to the TopK SAE.

For a wide range of sparsity (L0) values, we report the reconstruction MSE and the proportion of cross-entropy loss recovered when the sparse autoencoder output is patched into the language model. A loss recovered value of 1 corresponds to a perfect reconstruction, while a loss recovered value of 0 corresponds to a zero-ablation.

#### 4.2.1 FLOP-Matched Results

For the FLOP-matched experiments, we train Switch SAEs with 2, 4 and 8 experts, with the results 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 loss recovered. 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.

#### 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 loss recovered, while for the most part outperforming Gated and ReLU SAEs. When L0 is low, Switch SAEs perform particularly well. This suggests that the high frequency features that improve reconstruction fidelity the most for a given activation might lie within the same cluster.

These results demonstrate that Switch SAEs can reduce the number of 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.

Figure: (a) Fraction of SAE decoder vectors with nearest neighbor cosine sim $>0.9$, width-matched SAEs.
Refer to caption: x4.png

### 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 on 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 co-occur). We test this hypothesis using a common (see e.g. Braun et al. ()) SAE evaluation metric: nearest neighbor cosine similarity.

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 a nearest neighbor above a given cosine similarity, since highly similar decoder vectors are likely duplicates. In Fig. 4(a), we compare SAEs trained with different numbers of experts and different sparsities on this measure with a threshold of $0.9$, and find that as soon as we train with experts this measure jumpy sharply from a baseline of 0 0 in TopK SAEs to $5$ to $10$ percent. In other words, strict duplicates likely reduce Switch SAE capacity by up to $10\%$. This effect is less prevalent at larger L0s. However, we are not just worried about exact duplicates, but also a general shift towards redundancy: in Fig. 4(b), we find that a more global metric, cosine similarity of the nearest neighbor averaged across all features, shows a similar pattern across sparsity and number of experts.

Another measure of Switch SAE quality is expert specialization, that is, how similar are 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 one 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 a TopK SAE of the same size. On the other end, no specialization should result in identical experts. In Fig. 4(c), we plot this metric for all pairs of experts in a $16$ expert Switch SAE with $\text{L0}=64$. Some experts, e.g. expert 0 0, seem unique, while other pairs of experts, e.g. $10$ and $7$, are highly similar. All pairs are more similar than random blocks of the same size in a 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, we can also directly observe feature duplication – around the periphery of the plot we find a variety of isolated points which upon closer inspection are actually tight groupings of multiple features from different experts.

Figure: Figure 5: t-SNE projections of encoder and decoder features for a Switch SAE with 65k total features and 8 experts. Encoder features appear to cluster together by expert. Away from these clusters, we see a variety of isolated points which are in fact tight groups of features which have been duplicated across multiple experts. We do not observe as clear clusters in the decoder features.
Refer to caption: extracted/5917263/fig/tsne-threepanel-zoom.png

### 4.4 Automated Interpretability

Bricken et al. () evaluate how interpretable sparse autoencoder features using automated interpretability: they generate explanations for each feature by giving top activating examples for that feature to a language model, and then ask a language model (in a new context) to predict, using only the description of the feature, on which tokens the described feature fires. More recently, Juang et al. () introduce a more compute efficient method that measures detection, whether the explanation allows a language model to predict whether a feature fires at all on a given context. Juang et al. () additionally measure detection at each decile of feature activation, as well as on contexts where the feature does not fire at all (where to be correct the model should report that the feature does not fire). Using the Juang et al. () implementation, we find that FLOP-matched SAE features have similar interpretability as TopK SAEs, but width-matched SAEs have a higher true positive rate on contexts where the features fire, while at the same time having a 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.

Figure: 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.
Refer to caption: x7.png

## 5 Conclusion

Switch SAEs are a promising approach towards scaling sparse autoencoders, as they allow an improvement in the FLOPs vs. MSE scaling law without a significant reduction in feature interpretability. We believe that Switch SAEs may find their best application for huge training runs on large clusters of GPUs, since in this setting each expert can be placed on a separate GPU, leading to significant wall clock training speed ups. The most significant limitation of the Switch SAE is the reduction in performance of the SAE at a fixed number of parameters; future work to bridge this gap might investigate feature deduplication techniques, more sophisticated routing architectures, and multiple active experts. Overall, we believe that this work provides a case study for applying existing training optimization techniques to the setting of sparse autoencoders for feature extraction from language models; we hope that the investigations in this paper serve as a guide for adapting more such techniques.

#### Acknowledgments

We used infrastructure from the [https://github.com/saprmarks/dictionary_learning](https://github.com/saprmarks/dictionary_learning) repository to train our SAEs. We would like to thank Samuel Marks and Can Rager for advice on how to use the repository. We would also like to thank Achyuta Rajaram, Jacob Goldman-Wetzler, Michael Pearce, Gitanjali Rao, Satvik Golechha, Kola Ayonrinde, Rupali Bhati, Louis Jaburi, Vedang Lad, Adam Karvonen, Shiva Mudide, Sandy Tanwisuth, JP Rivera, Constantin Venhoff 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 Fellowship EP/W002981/1, and Armasuisse Science+Technology. AM was
greatly helped by the MATS program, funded by AI Safety Support.

## Appendix A Switch Sparse Autoencoder Training Details

We initialize the rows of $\mathbf{W}_{\text{enc}}^{i}$ to be parallel to the columns of $\mathbf{W}_{\text{dec}}^{i}$ for all $i$. We initialize both $\mathbf{b}_{\text{pre}}$ and $\mathbf{b}_{\text{router}}$ to the geometric median of a batch of samples, but we do not tie $\mathbf{b}_{\text{pre}}$ and $\mathbf{b}_{\text{router}}$. We additionally normalize the decoder column vectors to unit-norm at initialization and after each gradient step. We remove gradient information parallel to the decoder feature directions to minimize interference with the Adam optimizer. We set the learning rate based on the $\frac{1}{\sqrt{M}}$ scaling law from Gao et al. () and linearly decay the learning rate over the last 20% of training. We do not include neuron resampling (Bricken et al., ) or the AuxK loss (Gao et al., ), but future work could explore employing these strategies to prevent dead latents in larger models. We optimize $\mathcal{L}_{\text{total}}$ with Adam using the default parameters $\beta_{1}=0.9$, $\beta_{2}=0.999$.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
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. 

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. 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. 

The worse performance on the same width setting is indeed a drawback of this work, which might limit its applications.
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.
weaknesses: See above.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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.
soundness: 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.
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.
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.

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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
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.
- 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?
- 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?
- 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?

**Minor suggestions that do not individually affect the score**
- Line 295: “For a wide range of sparsity (L0) values”. Specify the range.
- Line 298: Clarify what is meant by a zero-ablation.
- Line 317: “Switch SAEs can likely achieve greater acceleration on larger language models”. Expand on this.
soundness: 2
strengths: The strengths of the paper are its simple approach and side studies.
- The approach enables a trade-off between model parameters and computational requirements which should be useful to practitioners.
- 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.
- The limitations (parameter efficiency and feature duplication) are discussed in an informative way.
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.
weaknesses: The main weakness of the paper is its quantitative evaluation.
- 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?
- 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.
- Figure 4 could use an explanatory caption.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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
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.

(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)

2. 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.

3. Do you have any experiments on different models or different sites?
soundness: 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.

2. **The evaluation is clear and relatively comprehensive**. Addressing FLOP-matched results and feature interpretability compared to baselines is very clear.

3. **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.
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.

This 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.

Here are my references for all sections in this review:

[1] Switch MoE: https://arxiv.org/abs/2101.03961

[2] DeepSeek MoE: https://arxiv.org/abs/2401.06066

[3] GShard: https://arxiv.org/abs/2006.16668

[4] Expert-choice; https://arxiv.org/abs/2202.09368

[5] Ghost Grads: https://transformer-circuits.pub/2024/jan-update#dict-learning-resampling

[6] Updated Anthropic training method: https://transformer-circuits.pub/2024/april-update/index.html#training-saes

[7] Gated SAEs: https://arxiv.org/abs/2404.16014

[8] JumpReLU SAEs: https://arxiv.org/abs/2407.14435

[9] Anthropic Set Of Metrics Where JumpReLU does not reproduce: https://transformer-circuits.pub/2024/august-
update/index.html#interp-evals

[10] Scaling Monosemanticity: https://transformer-circuits.pub/2024/scaling-monosemanticity/

[11] SFC: https://arxiv.org/abs/2403.19647
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.

2. **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.

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.

I 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.

3. **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).

4. **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.

Note 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.

Meta-review: The paper introduces an improvement to sparse autoencoders (SAEs) which are a (currently) popular set of tools used in LLM interpretability. SAEs suffer from the fact that the encoder is dense, and therefore compute-heavy; the contribution of this paper is to sparsify the encoder using a "switch" / routing layer that selects one of several smaller encoders/"experts". This enables improvements; for iso-FLOPs, the proposed switch SAEs are easier to train and obtain lower reconstruction cost. The features obtained by the resulting SAEs achieve comparable results on interpretability tasks as standard SAEs. 

The paper is very well-written. Although the idea is a straightforward composition of existing tools (SAEs + switch layers), it is very well executed and may spur further work.

The drawback, as pointed out by a couple of reviewers, is that the number of features blows up, which may affect human-interpretability. To this point, the authors make the reasonable case that reconstruction error affects downstream tasks (including interpretability!), so on net this tradeoff might be worth it. I'd encourage the authors to probe this point a bit more in detail while preparing the final version.

Decision: Accept (Poster)