Efficient Dictionary Learning with Switch Sparse Autoencoders

TL;DR

This work tackles the computational challenges of scaling sparse autoencoders to extract interpretable features from large language-model activations. It introduces the Switch Sparse Autoencoder, a mixture-of-experts architecture that routes activations to multiple small expert SAEs via a routing network, reducing FLOPs while preserving interpretability. The study develops end-to-end training with a reconstruction objective augmented by a load-balancing term, and analyzes scaling laws, sparsity-reconstruction tradeoffs, feature geometry, and automated interpretability, using GPT-2 residual activations. Results show Switch SAEs achieve Pareto improvements in reconstruction vs. compute over dense TopK SAEs, with some feature duplication across experts that diminishes with scale, suggesting practical applicability for large-scale training with potential wall-clock speedups on multi-GPU clusters.

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.

Figures (7)

• Figure 1: Scaling laws for Switch SAEs. We train dense TopK SAEs and Switch SAEs of varying size with fixed $k=32$. Left: Switch SAEs achieve better reconstruction than dense SAEs at a fixed compute budget. Right: Switch SAEs require more features in total (and therefore more parameters) to achieve the same reconstruction as dense SAEs when trained to convergence, although this gap narrows for larger SAEs.
• 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$.
• Figure 3: Pareto frontier of sparsity versus (top) reconstruction mean squared error and (bottom) 1 - FLR [fraction of 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).
• 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 L0s. Furthermore, between-expert similarities show that experts specialize moderately; expert 0, for example, has low similarity with most other experts.
• 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.