Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders
Senthooran Rajamanoharan, Tom Lieberum, Nicolas Sonnerat, Arthur Conmy, Vikrant Varma, János Kramár, Neel Nanda
2024-07-22
Summary
This paper introduces JumpReLU Sparse Autoencoders (SAEs), a new type of model that improves how well we can reconstruct data while keeping the model simple and efficient. The authors focus on making these models better at understanding and processing information from language models.
What's the problem?
Sparse autoencoders are used to analyze data from language models, but they face a challenge: they need to be both accurate (faithful) in reconstructing the original data and simple (sparse) by using fewer features. These two goals often conflict, making it hard to achieve high performance without sacrificing one for the other.
What's the solution?
The authors propose JumpReLU SAEs, which use a modified activation function called JumpReLU instead of the traditional ReLU. This change allows the model to ignore certain inputs more effectively, improving reconstruction quality without losing interpretability. They also implement a training method that directly encourages sparsity without relying on simpler but less effective techniques. Their experiments show that JumpReLU SAEs outperform other models like Gated and TopK SAEs in terms of reconstruction quality while remaining easy to train.
Why it matters?
This research is significant because it enhances the capabilities of sparse autoencoders, making them more effective for analyzing complex data from language models. By improving how these models work, we can better understand and interpret the information processed by AI systems, which is crucial for developing advanced applications in natural language processing and beyond.
Abstract
Sparse autoencoders (SAEs) are a promising unsupervised approach for identifying causally relevant and interpretable linear features in a language model's (LM) activations. To be useful for downstream tasks, SAEs need to decompose LM activations faithfully; yet to be interpretable the decomposition must be sparse -- two objectives that are in tension. In this paper, we introduce JumpReLU SAEs, which achieve state-of-the-art reconstruction fidelity at a given sparsity level on Gemma 2 9B activations, compared to other recent advances such as Gated and TopK SAEs. We also show that this improvement does not come at the cost of interpretability through manual and automated interpretability studies. JumpReLU SAEs are a simple modification of vanilla (ReLU) SAEs -- where we replace the ReLU with a discontinuous JumpReLU activation function -- and are similarly efficient to train and run. By utilising straight-through-estimators (STEs) in a principled manner, we show how it is possible to train JumpReLU SAEs effectively despite the discontinuous JumpReLU function introduced in the SAE's forward pass. Similarly, we use STEs to directly train L0 to be sparse, instead of training on proxies such as L1, avoiding problems like shrinkage.