🤖 AI Summary
This work addresses the excessive parameter count and memory overhead in conventional sparse autoencoders caused by dense decoders, which hinder efficient mechanistic interpretability. The authors propose the Expander Sparse Autoencoder, which introduces a left-$d$-regular expander graph mask into both the decoder and its tied encoder to enforce a sparse connectivity structure. This design substantially reduces learnable parameters while preserving the scale of the sparse coding problem. Theoretical analysis establishes that under expansion and column-flatness conditions, the method guarantees identifiability of $k$-sparse codes. Experiments on the Qwen2.5-3B model demonstrate that with $d=7$, the decoder’s parameters are reduced by a factor of 293 while retaining 84% of the dense model’s cross-entropy loss recovery capability, achieving a significantly improved trade-off between memory efficiency and reconstruction fidelity.
📝 Abstract
Sparse autoencoders (SAEs) decompose internal activations of neural networks into sparse linear combinations of learned features by fitting an overcomplete dictionary $\mathbf{W}\in\mathbb{R}^{m\times n}$ with $m<n$, and inferring a sparse code $\mathbf{x}\in\mathbb{R}^n$ from $\mathbf{h}\approx\mathbf{W}\mathbf{x}$. This inference problem closely resembles the canonical setup of compressed sensing, but dense decoders requires $O(mn)$ learned values, which becomes costly at large feature counts. We introduce Expander SAEs: TopK SAEs whose decoder and tied encoder are supported on a left-$d$-regular expander mask with $d\ll m$, learning only $dn$ decoder values while keeping the sparse-coding problem $(m,n,k)$ fixed. The same structure reduces storage and turns the matching-pursuit correlation step $\mathbf{W}^\top \mathbf{r}$ in OMP into an $O(dn)$ gather-and-reduce operation. Our experiments show that across Pythia-70M/160M, Qwen2.5-3B, and Llama-3.2-1B residual-stream activations, varying $d$ traces a consistent storage--fidelity frontier, and that at the most compressed modern-LM setting, Qwen2.5-3B with $d=7$ uses $293\times$ fewer learned decoder values than the full dense decoder while retaining $84$% of dense CE-loss recovered. Control experiments show that the improved storage--fidelity tradeoff is driven by sparse, diverse decoder support structure rather than by fewer learned decoder values, and that when sparse and dense decoders are compared at matched parameter count, part of the remaining gap comes from encoder amortisation. On the theoretical side, we show that expansion and column flatness are sufficient for identifiability of noiseless $k$-sparse codes, and we derive complementary sufficient conditions under which OMP recovers the support exactly.