Enhancing Neural Network Interpretability with Feature-Aligned Sparse Autoencoders

TL;DR

This paper introduces Mutual Feature Regularization (MFR), a technique that encourages parallel sparse autoencoders to learn similar, input-relevant features, improving interpretability and reconstruction performance on synthetic, EEG, and GPT-2 data.

Contribution

The paper proposes MFR, a novel regularization method that aligns features learned by multiple SAEs, enhancing interpretability and reconstruction accuracy across different data types.

Findings

MFR improves SAE reconstruction loss by up to 21.21% on GPT-2 Small.

MFR helps SAEs learn input features more effectively on synthetic data.

Enhanced interpretability of neural network features through aligned SAE training.

Abstract

Sparse Autoencoders (SAEs) have shown promise in improving the interpretability of neural network activations, but can learn features that are not features of the input, limiting their effectiveness. We propose \textsc{Mutual Feature Regularization} \textbf{(MFR)}, a regularization technique for improving feature learning by encouraging SAEs trained in parallel to learn similar features. We motivate \textsc{MFR} by showing that features learned by multiple SAEs are more likely to correlate with features of the input. By training on synthetic data with known features of the input, we show that \textsc{MFR} can help SAEs learn those features, as we can directly compare the features learned by the SAE with the input features for the synthetic data. We then scale \textsc{MFR} to SAEs that are trained to denoise electroencephalography (EEG) data and SAEs that are trained to reconstruct GPT-2 Small activations. We show that \textsc{MFR} can improve the reconstruction loss of SAEs by up to 21.21\% on GPT-2 Small, and 6.67\% on EEG data. Our results suggest that the similarity between features learned by different SAEs can be leveraged to improve SAE training, thereby enhancing performance and the usefulness of SAEs for model interpretability.