Identifying Functionally Important Features with End-to-End Sparse Dictionary Learning

TL;DR

This paper introduces end-to-end sparse dictionary learning, a method that trains sparse autoencoders to identify functionally important features in neural networks by directly aligning their output distributions with the original model.

Contribution

The paper proposes a novel end-to-end training approach for sparse autoencoders that ensures learned features are functionally relevant to the network's behavior.

Findings

E2E SAEs explain more network performance with fewer features.

E2E SAEs require fewer active features per data point.

E2E SAEs improve interpretability without sacrificing accuracy.

Abstract

Identifying the features learned by neural networks is a core challenge in mechanistic interpretability. Sparse autoencoders (SAEs), which learn a sparse, overcomplete dictionary that reconstructs a network's internal activations, have been used to identify these features. However, SAEs may learn more about the structure of the datatset than the computational structure of the network. There is therefore only indirect reason to believe that the directions found in these dictionaries are functionally important to the network. We propose end-to-end (e2e) sparse dictionary learning, a method for training SAEs that ensures the features learned are functionally important by minimizing the KL divergence between the output distributions of the original model and the model with SAE activations inserted. Compared to standard SAEs, e2e SAEs offer a Pareto improvement: They explain more network performance, require fewer total features, and require fewer simultaneously active features per datapoint, all with no cost to interpretability. We explore geometric and qualitative differences between e2e SAE features and standard SAE features. E2e dictionary learning brings us closer to methods that can explain network behavior concisely and accurately. We release our library for training e2e SAEs and reproducing our analysis at https://github.com/ApolloResearch/e2e_sae