Research map / Architectures and efficiency / RNNs and normalization
Efficient Dictionary Learning with Switch Sparse Autoencoders
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces Switch Sparse Autoencoders, which route activations to one of several small TopK expert SAEs (Switch layer style) to cut training FLOPs and scale feature count for interpretability.
Problem
Training sparse autoencoders at very large width is computationally expensive because the dense encoder forward pass and latent pre-activation memory dominate cost, limiting how many features can be learned in frontier models.
Approach
Apply Switch-layer mixture-of-experts conditional computation to SAEs: a trainable router sends each activation vector to a single expert TopK SAE (no bias), with the output weighted by the router probability and trained end-to-end with a load-balancing auxiliary loss.
Key results
- FLOP-matched Switch SAEs Pareto-dominate TopK, Gated and ReLU SAEs in reconstruction MSE and FLR at the same compute budget.
- Width-matched Switch SAEs with 16-128 experts perform slightly worse than TopK SAEs but outperform ReLU and Gated SAEs while using up to 128x fewer FLOPs per activation.
- Feature duplication: Switch SAEs show 5-10% of decoder vectors with nearest-neighbor cosine similarity above 0.9, versus 0 for TopK SAEs; Switch SAE features are roughly as interpretable as TopK features in FLOP-matched settings.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Switch SAE instead of TopK SAE at fixed width | GPT-2 residual-stream SAEs · GPT-2 residual stream, layers 2, 4, 8, and 10 | FVE | Mixed |
| Use Switch SAE instead of TopK SAE at fixed width | GPT-2 attention-output SAEs · GPT-2 attention outputs, layers 2, 4, 8, and 10 | FVE | Hurt |
| Use Switch SAE instead of TopK SAE at fixed width | GPT-2 MLP-output SAEs · GPT-2 MLP outputs, layers 2, 4, 8, and 10 | FLR | Hurt |
| Use Switch SAE rather than TopK SAE while matching FLOPs | SAEs trained on GPT-2 small layer 8 residual-stream activations · Sparsity-reconstruction frontier | MSE | Helped |
| Use Switch SAE rather than TopK SAE while matching width | SAEs trained on GPT-2 small layer 8 residual-stream activations · Width-matched sparsity-reconstruction frontier | MSE | Hurt |
| Increase the number of Switch SAE experts | Switch SAE · Reconstruction at fixed compute | MSE | Helped |
| Increase the number of Switch SAE experts | Switch SAE · Feature duplication | nearest-neighbor cosine similarity | Hurt |
| Use Switch SAE rather than TopK SAE at fixed width | Gemma 2 2B residual-stream SAE · Gemma 2 2B layer 12 resid | FVE: 0.961 -> 0.955 | Hurt |
| Use Switch SAE rather than TopK SAE at fixed width | GPT-2 layer 2 MLP SAE · GPT-2 layer 2 mlp | FLR: 0.888 -> 0.701 | Hurt |
| Use Switch SAE rather than TopK SAE at fixed width | GPT-2 layer 4 residual-stream SAE · GPT-2 layer 4 resid | FVE: 0.997 -> 0.997 | No clear effect |
| Set router load-balancing weight α to 3 | 16- and 32-expert Switch SAEs · GPT-2 small, L0=64, fixed width 24576 | reconstruction MSE | Helped |
| Vary router load-balancing weight α | Switch SAE · GPT-2 small, L0=64, fixed width 24576 | reconstruction MSE | No clear effect |
Surprising findings
- Switch SAEs perform worse than dense TopK SAEs at fixed width (parameter efficiency), with the gap shrinking as scale increases.
- Switch SAEs trained on MLP and attention outputs suffer significant performance drops, while residual stream training is minimally affected.
Limitations
- Fixed-width reconstruction performance is worse than dense TopK SAEs, and feature duplication across experts reduces capacity.
- Only the simplest MoE architecture (top-1 routing) is studied; no multiple active experts or advanced routing.
- No downstream evaluations (e.g., steering or circuit discovery) are performed.
Experimental setup
Datasets: OpenWebText
Benchmarks: OpenWebText reconstruction (GPT-2 layer 8 residual stream)
Metrics: reconstruction MSE fraction of loss recovered L0 sparsity nearest-neighbor cosine similarity automated interpretability detection
Baselines: TopK SAE ReLU SAE Gated SAE
Builds on: Switch Transformer TopK SAE Sparsely-Gated Mixture-of-Experts Gated SAE ReLU SAE
Closest papers by meaning
- Identifying Functionally Important Features with End-to-End Sparse Dictionary Learning (NeurIPS 2024 · similarity 0.91)
- Towards Principled Evaluations of Sparse Autoencoders for Interpretability and Control (ICLR 2025 · similarity 0.90)
- Sanity Checks for Sparse Autoencoders: Do SAEs Beat Random Baselines? (NeurIPS 2026 · similarity 0.90)
- Learning Multi-Level Features with Matryoshka Sparse Autoencoders (ICML 2025 · similarity 0.90)
- Are Sparse Autoencoders Useful? A Case Study in Sparse Probing (ICML 2025 · similarity 0.88)
- Compute Optimal Inference and Provable Amortisation Gap in Sparse Autoencoders (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.