Group Equivariance Meets Mechanistic Interpretability: Equivariant Sparse Autoencoders

TL;DR

This work addresses mechanistic interpretability for scientific data by introducing adaptive, group-equivariant sparse autoencoders (SAEs). The authors show that a single linear map $M$ can explain activation-space transformations under a rotation group, enabling an invariant encoder and a learned equivariance map to adapt to the base model's symmetry. Their adaptive equivariant SAEs discover more informative features and improve downstream probing tasks compared to regular SAEs, albeit with a trade-off in reconstruction sparsity. This approach offers a pathway to leverage domain symmetries to enhance interpretability tools in scientific ML, with reproducible methods and scalable concepts for future work on larger models and more complex groups.

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.

Figures (7)

• Figure 1: Left: We train an invariant SAE that maps activations of the transformed inputs to the same latents, and optimize the matrix $\mathbf{M}$ to estimate how the activations transform, achieving equivariance. Right: Transforming the decoder dictionary $\mathbf{D} \mapsto \mathbf{M}\mathbf{D}$ allows us to observe which features discovered by the SAE are invariant or equivariant with respect to input transformations. We provide a more detailed figure of the mathematical operations involved in Figure \ref{['fig:method_detailed']} in the Appendix.
• Figure 2: Sample images from our dataset. Each image contains one of eight possible shapes in each of its four quadrants.
• Figure 3: Latent and reconstruction probing performance, with truncation length 32 and the CNN autoencoder, comparing a regular SAE (one-layer encoder), an SAE with $\vert G \vert$ times the latents (wide), and an SAE with a two-layer encoder with our equivariant SAE. Performance of probes over the base model activations are duplicated for easier comparisons and reflect the upper bound.
• Figure 4: Sparsity/reconstruction performance of SAEs for varying TopK values. The x axis corresponds to the base autoencoders' reconstruction performance when their intermediate activations are passed through the SAEs, and the y axis denotes the L1 norm of the SAE latent activations.
• Figure 5: Left: A neural network$\psi$transforms an orbit $G\mathbf{x}$ into more complex, non-linear structures. The invariant SAE maps all activations in a transformed orbit to the same latent point, and reconstructs a canonical activation. The matrix$\mathbf{M}$ then transforms the reconstructions back to their original forms, achieving equivariance by approximating the transformed orbit in activation space. Right: Transforming the learned dictionary $\mathbf{D} \mapsto \mathbf{M}^p \mathbf{D}$ allows us to observe which features discovered by the SAE are invariant or equivariant with respect to input transformations.

Theorems & Definitions (4)

• Definition 1: Group
• Definition 2: Group action
• Definition 3: Orbit
• Definition 4: Invariant and equivariant functions