🤖 machine learning

Equivariant Sparse Autoencoders: Mechanistic Interpretability of Neural Networks on Symmetric Data

This paper introduces Equivariant Sparse Autoencoders, which incorporate data symmetries to resolve the unidentifiability issues of standard SAEs on symmetric datasets, thereby discovering more useful and interpretable features even when reconstruction quality is lower.

Original authors: Ege Erdogan, Ana Lucic

Published 2026-08-10

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ege Erdogan, Ana Lucic

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand how a giant, super-smart robot thinks. You can't ask it questions, so instead, you have to peek inside its brain while it works. This field is called mechanistic interpretability, and it's like being a detective trying to figure out the secret code of a black box. The robot's brain is made of layers of tiny switches (neurons) that light up in complex patterns. The problem is that these patterns are messy. Often, the robot mixes many different ideas into a single switch, a phenomenon researchers call "superposition." It's like trying to read a book where every sentence is a jumbled mix of three different stories; you know something is happening, but you can't tell what.

To fix this mess, scientists use a tool called a Sparse Autoencoder (SAE). Think of an SAE as a super-organized librarian. Its job is to take that messy jumble of light-up switches and sort them out into a neat list of single, clear concepts. If the robot is looking at a picture of a cat, the librarian tries to find the specific "cat" switch and turn it on, while turning off everything else. For a long time, this worked well for things like text, where the rules are mostly the same no matter how you read a word. But what happens when the data is symmetric? In the real world, many things look the same even if you spin them around or flip them. A cat is still a cat whether it's facing left or right. If your librarian doesn't know about this rule, they might get confused, thinking a "left-facing cat" and a "right-facing cat" are two completely different, unrelated things. This paper asks: What happens when we teach our librarian to respect these symmetries?

The researchers, Ege Erdogan and Ana Lucic from the University of Amsterdam, decided to upgrade the standard librarian tool to handle these spinning and flipping rules. They call their new tool an Equivariant Sparse Autoencoder. Instead of just trying to sort the messy lights into a list, they built a system that understands: "Hey, if you rotate the picture, the 'cat' concept doesn't disappear; it just moves to a different spot in the list." They tested this idea on a toy model, a dataset of geometric shapes, and real-world images of galaxies and blood cells.

Here is the surprising twist they found: The new, symmetry-aware librarians were actually worse at perfectly reconstructing the original images than the old, messy ones. If you measured them by how well they could rebuild the picture from their notes, the old tools won. However, when the researchers asked a different question—"Which notes are actually useful for understanding what the robot is seeing?"—the new tools were the clear winners. The old librarians were creating notes that looked perfect for rebuilding the image but were actually less useful when it came to identifying the actual concepts. The new librarians, even though their notes were a bit messier to look at, gave a much truer picture of what the robot was thinking.

The paper suggests that for data with symmetries (like rotating objects), the standard way of judging these tools—checking how well they can rebuild the input—is misleading. In fact, the better a tool is at rebuilding the image, the less useful its features might be for understanding the underlying concepts. By building the rule of symmetry directly into the tool, the researchers found features that were much better at helping computers identify shapes, galaxy types, and cell types, even if the final image reconstruction wasn't perfect. They didn't prove this is the ultimate solution for every AI, but their simulations and experiments on real data strongly suggest that ignoring symmetry makes our tools for understanding AI less reliable, and that we should stop relying solely on "reconstruction quality" as our main scorecard.

Technical Summary: Equivariant Sparse Autoencoders

Problem Statement

Machine learning models, particularly in scientific domains, often operate on data with inherent symmetries (e.g., rotations in protein folding, galaxy images, or cell microscopy). While Sparse Autoencoders (SAEs) have emerged as a primary tool for mechanistic interpretability by decomposing dense neural activations into sparse, interpretable features, they face a critical limitation when applied to symmetric data: unidentifiability.

Standard SAEs assume a Linear Representation Hypothesis (LRH) where activations are sparse linear combinations of features. However, when input data possesses symmetries (group actions $G$ ), a vanilla SAE must learn distinct feature directions for the same semantic concept under different transformations (e.g., a "cat" feature for every possible rotation). This leads to:

Feature Entanglement: The model may learn non-orthogonal, uninterpretable features to reconstruct disparate inputs.
Unidentifiability: Different explanations can fit the data equally well, but only one may align with the true underlying structure.
Metric Misalignment: Reconstruction quality (MSE) often fails to correlate with feature usefulness, as SAEs may learn reconstruction-specific features that do not correspond to ground-truth concepts.

The paper posits that current SAEs fail to account for these symmetries, leading to suboptimal interpretability in scientific domains where such symmetries are prevalent.

Methodology: Equivariant SAEs

The authors propose Equivariant Sparse Autoencoders, which extend the Linear Representation Hypothesis to explicitly account for group symmetries.

1. Extended Linear Representation Hypothesis

The authors reformulate the LRH to state that for a group $G$ acting on inputs $X$ , the neural network activations $\psi(x)$ transform predictably. Specifically, if $x$ is transformed by $g \in G$ , the activations transform via a linear map $\rho(g)$ :
$\psi(gx) = \rho(g) \sum_{i=1}^C z(x)_i D_i = \sum_{i=1}^C z(x)_i (\rho(g) D_i)$
Here, the feature magnitudes $z(x)$ remain invariant (sparse coefficients), while the feature vectors $D_i$ transform according to the group action.

2. Architecture Design

The Equivariant SAE decouples the learning of invariant features from the learning of transformation rules:

Invariant SAE: An encoder-decoder pair trained to map all activations within an orbit (e.g., all rotations of an image) to a single canonical latent representation. This learns the invariant features.
Linear Transformation Module: A separate set of learnable matrices ( $M, R$ ) that model how the canonical reconstruction transforms back to the original input space. For a group generated by $g, h$ , the model learns matrices such that:
$M_p R_q D(E(\psi(g^p h^q x))) = \psi(g^p h^q x)$
This allows the SAE to capture the symmetry structure without forcing the dictionary itself to be equivariant, which would require an impractical number of parameters.

3. Training Strategy

Training occurs in two independent steps (which can be parallelized):

Invariant SAE Training: The encoder and decoder are optimized to minimize reconstruction loss on canonical inputs while enforcing invariance. The authors experiment with three loss objectives:
- $L_{Canonical}$ : Reconstruct a canonical sample from any orbit member.
- $L_{Latent}$ : Enforce invariance directly on the latent codes ( $E(\psi(x)) \approx E(\psi(gx))$ ).
- $L_{Output}$ : Enforce invariance on the SAE outputs.
Transformation Learning: The matrices $M$ and $R$ are optimized to minimize the error between the predicted transformed activations and the actual transformed activations of the base model.

Key Contributions

Theoretical Extension: The paper extends the Linear Representation Hypothesis to symmetric data, formally defining how feature coefficients and vectors behave under group actions.
Algorithmic Innovation: The design of Equivariant SAEs that separate invariant feature learning from transformation modeling, avoiding the parameter explosion of group-equivariant dictionaries.
Empirical Validation: Demonstration across synthetic and real-world datasets (shapes, galaxy images, blood cells) that Equivariant SAEs outperform vanilla, Archetypal, and Group Cross-coder baselines in feature recovery and downstream probing tasks.
Metric Insight: The discovery that reconstruction quality is inversely correlated with feature usefulness in symmetric settings. Equivariant SAEs often achieve lower reconstruction accuracy (higher MSE) than vanilla SAEs but yield significantly more interpretable and useful features.

Experimental Results

The authors evaluated their method on:

Toy Models: Synthetic data with known ground-truth equivariant features.
Shapes Dataset: Synthetic images with $C_4$ (rotations) and $D_4$ (rotations + flips) symmetries, using MLP, CNN, and ViT base models.
GalaxyMNIST: Real-world galaxy images with $C_4/D_4$ symmetries.
MLL23: Peripheral blood cell images using a DinoBloom foundation model.

Key Findings:

Feature Recovery: On toy models, Equivariant SAEs achieved superior feature recovery (Mean Cosine Similarity) and detection (F1 Score) compared to baselines, especially in overcomplete regimes where interference is high.
Downstream Probing: Across all datasets, probes trained on Equivariant SAE latents achieved higher accuracy in classifying input properties (e.g., shape presence, orientation, cell type) than probes trained on vanilla SAE latents. In some cases (e.g., GalaxyMNIST), they matched the performance of probes trained directly on the base model activations.
Reconstruction vs. Utility: A consistent trend emerged where Equivariant SAEs had higher reconstruction error (worse MSE) than vanilla SAEs. However, this "worse" reconstruction corresponded to features that were more aligned with the true semantic structure and more useful for downstream tasks.
Transformation Modeling: The learned linear matrices ( $M, R$ ) successfully explained over 97% of the variance in activation transformations for most models, validating the hypothesis that input symmetries induce linear transformations in activation space.

Significance and Claims

The paper claims that aligning the priors of interpretability methods with the known structure of the data (specifically symmetries) is crucial for obtaining faithful explanations. The authors highlight two major implications:

Reconstruction is not Interpretability: Relying solely on reconstruction metrics (like MSE) to evaluate SAEs is dangerous, particularly for symmetric data. A model that reconstructs inputs poorly may still discover the most useful and interpretable features.
Domain-Specific Interpretability: For scientific applications (e.g., biology, astronomy) where data symmetries are fundamental, standard SAEs are insufficient. Equivariant SAEs provide a pathway to discover features that are robust to these symmetries, potentially leading to novel scientific insights and more reliable AI systems.

The work suggests that future interpretability research must move beyond generic sparse coding and incorporate domain-specific structural priors to avoid the pitfalls of unidentifiability.