🤖 AI Summary
This work addresses the limitations of sparse autoencoders (SAEs) in interpretability due to “dead features” and training instability. The authors propose an alignment-based training method that imposes geometric constraints on the encoder and decoder, enforcing unit inner products between corresponding features. This approach reveals, for the first time, a bimodal distribution in SAE feature alignment scores. To mitigate training collapse without introducing additional parameters or data, they introduce a reparameterization scheme. Combined with Top-K and BatchTop-K architectures and p-annealing, the method achieves consistent Pareto improvements across diverse models, dictionary sizes, and sparsity levels, significantly reducing dead features while enhancing reconstruction quality and cross-seed stability.
📝 Abstract
Sparse autoencoders (SAEs) are one of the main methods to interpret the inner workings of deep neural networks (DNNs), decomposing activations into higher-dimensional features. However, they exhibit critical shortcomings where a large fraction of features are never activated and are unstable. Despite variants of SAEs that attempt to mitigate these issues, they require additional data, resampling, or training. We propose the \textbf{aligned training}, a parameter-free reparameterization of SAEs that simultaneously improves reconstruction quality, eliminates dead features, and significantly enhances stability across training seeds. Our approach is motivated by an overlooked observation that SAE feature quality, measured by the inner product between encoder and decoder directions (which we call the \textbf{alignment score}), follows a bimodal distribution across all modern architectures. The proposed aligned training enforces a geometric constraint between the encoder and decoder such that their inner product equals one for every feature, which removes a source of degeneracy in the SAE training without adding any hyperparameters. Across multiple models, dictionary sizes, and sparsity levels, the aligned training shows Pareto improvements on the SAEBench benchmarks. Beyond improving dead features, stability and reconstruction, our method readily integrates with techniques in mechanical interpretability such as Top/BatchTop-K architectures and p-Annealing. Overall, the aligned training substantially improves feature quality and stability of SAE without computational complexity or cost.