Features that Make a Difference: Leveraging Gradients for Improved Dictionary Learning
deep learning;sparse autoencoders;dictionary learning;mechanistic interpretability;large language models;activation function;feature extraction
I have read and agree with the venue's withdrawal policy on behalf of myself and my co-authors.
We find that incorporating gradients into the Sparse Autoencoder architecture improves SAE performance.
@misc{
olmo2024features,
title={Features that Make a Difference: Leveraging Gradients for Improved Dictionary Learning},
author={Jeffrey Olmo and Jared Wilson and Max Forsey and Bryce Hepner and Thomas Vincent Howe and David Wingate},
year={2024},
url={https://openreview.net/forum?id=dpBIUI7D9c}
}
Sparse Autoencoders (SAEs) are a promising approach for extracting neural network representations by learning a sparse and overcomplete decomposition of the network's internal activations. However, SAEs are traditionally trained considering only activation values and not the effect those activations have on downstream computations. This limits the information available to learn features, and biases the autoencoder towards neglecting small but functionally significant aspects of activations. To address this, we introduce Gradient SAEs (g-SAEs), which modify the $k$-sparse autoencoder architecture by augmenting the topK activation function to rely on the gradients of the input activation when selecting the $k$ elements. For a given sparsity level, g-SAEs produce reconstructions that are more faithful to original network performance when propagated through the network, while facilitating more complete use of the SAE's capacity. On average, g-SAE features are more functionally influential than features recovered with traditional SAEs. By considering the downstream effects of activations, our approach leverages the dual nature of neural network features as both $\textit{representations}$, retrospectively, and $\textit{actions}$, prospectively. While previous methods have approached the problem of feature discovery primarily focused on the former aspect, we believe that g-SAEs represent a step towards accounting for the latter as well.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
~Jeffrey_Olmo1 ~Jared_Wilson1 ~Max_Forsey1 ~Bryce_Hepner1 ~Thomas_Vincent_Howe1 ~David_Wingate1
Jeffrey Olmo Jared Wilson Max Forsey Bryce Hepner Thomas Vincent Howe David Wingate
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
deep learning sparse autoencoders dictionary learning mechanistic interpretability large language models activation function feature extraction
I certify that there is no acknowledgement section in this submission for double blind review.
olmo|features_that_make_a_difference_leveraging_gradients_for_improved_dictionary_learning
interpretability and explainable AI
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Features that Make a Difference: Leveraging Gradients for Improved Dictionary Learning
ICLR 2025 Conference Withdrawn Submission
ICLR.cc/2025/Conference/Withdrawn_Submission