Group Equivariance Meets Mechanistic Interpretability: Equivariant Sparse Autoencoders

TL;DR

This paper introduces adaptive equivariant sparse autoencoders that incorporate group symmetries to improve interpretability and downstream task performance in scientific data analysis.

Contribution

It develops adaptive equivariant SAEs that learn symmetry-aware features, enhancing interpretability and performance over traditional SAEs.

Findings

Features transform predictably under rotations.

Adaptive SAEs outperform regular SAEs in probing tasks.

Incorporating symmetries improves interpretability.

Abstract

Sparse autoencoders (SAEs) have proven useful in disentangling the opaque activations of neural networks, primarily large language models, into sets of interpretable features. However, adapting them to domains beyond language, such as scientific data with group symmetries, introduces challenges that can hinder their effectiveness. We show that incorporating such group symmetries into the SAEs yields features more useful in downstream tasks. More specifically, we train autoencoders on synthetic images and find that a single matrix can explain how their activations transform as the images are rotated. Building on this, we develop adaptively equivariant SAEs that can adapt to the base model's level of equivariance. These adaptive SAEs discover features that lead to superior probing performance compared to regular SAEs, demonstrating the value of incorporating symmetries in mechanistic interpretability tools.