Universal Sparse Autoencoders: Interpretable Cross-Model Concept Alignment

1. Problem Statement

Current interpretability methods for Deep Neural Networks (DNNs), particularly Sparse Autoencoders (SAEs), typically operate in a single-model, post-hoc manner. They learn a dictionary of concepts for one specific model and task. This approach has several limitations:

• Lack of Unification: It is difficult to compare internal representations across different architectures (e.g., ViT vs. CNN), training paradigms (e.g., supervised vs. self-supervised), or datasets.
• Inefficiency: Existing methods for finding "universal" features often involve computationally expensive post-hoc matching or filtering of concepts discovered independently in each model.
• Scalability: There is no unified framework to simultaneously reconstruct and interpret the internal activations of multiple models, hindering the analysis of shared "foundation" concepts.

The authors propose a framework to discover a universal concept space that aligns the internal representations of multiple diverse models, enabling simultaneous interpretation and cross-model analysis.

2. Methodology: Universal Sparse Autoencoders (USAEs)

The core innovation is the Universal Sparse Autoencoder (USAE), a framework that jointly learns a shared, overcomplete dictionary of concepts from multiple models.

Architecture and Training

Unlike standard SAEs that map activations of a single model to a latent space and back, USAEs are designed to handle $M$ different models, each with potentially different feature dimensions ($d_i$).

• Shared Latent Space: The model learns a single sparse code space $Z \in \mathbb{R}^{n \times m}$ (where $m$ is the dictionary size).
• Model-Specific Encoders: Each model $i$ has its own encoder $\Psi^{(i)}_\theta$ that maps its specific activations $A^{(i)}$ into the shared space $Z$.
• Model-Specific Decoders: Each model $i$ has its own decoder (dictionary) $D^{(i)}$ that reconstructs its original activations from the shared code $Z$.
• Training Objective: The training minimizes a universal loss that aggregates reconstruction errors across all models.
  • In each iteration, a single model $i$ is selected to encode a batch of activations into $Z$.
  • This shared $Z$ is then decoded by all $M$ decoders to reconstruct the activations of every model $j$.
  • The loss is the sum of Frobenius norms of the reconstruction errors:
    $$L_{Universal} = \sum_{j=1}^{M} \| A^{(j)} - \hat{A}^{(j)} \|_F = \sum_{j=1}^{M} \| A^{(j)} - \Psi^{(i)}_\theta(A^{(i)}) D^{(j)} \|_F$$
  • Only the encoder $\Psi^{(i)}_\theta$ and decoder $D^{(i)}$ of the selected model are updated in that step, but the shared $Z$ forces alignment.

Key Application: Coordinated Activation Maximization (CAM)

The authors introduce a novel visualization technique enabled by the shared space. Instead of maximizing a concept in one model, they optimize inputs for all models simultaneously to activate the same concept dimension $k$. This reveals how different architectures encode the same semantic concept (e.g., "dog face") and highlights architectural biases.

3. Experimental Setup

• Models: Three diverse vision models were used:
  1. DinoV2: Self-supervised, image-level and patch-level discriminative learning.
  2. SigLIP: Vision-language contrastive learning.
  3. ViT (ImageNet): Supervised classification.
• Data: Trained on ImageNet training set; validated on ImageNet, DTD (textures), and CelebA (faces).
• Configuration: TopK SAE with a dictionary size of 6144 (8x expansion of the largest feature dimension, 768).

4. Key Results

A. Discovery of Universal Concepts

The USAE successfully discovered a broad spectrum of interpretable concepts shared across all three models:

• Low-level: Colors (yellow, blue), textures, and geometric shapes.
• Mid-level: Structural relationships (foreground/background, thin objects).
• High-level: Object parts (dog faces, animal jaws) and groupings (crowds).
• Significance: These concepts persist across different training objectives and architectures, suggesting a "Platonic" representation of visual concepts.

B. Quantitative Validation

• Cross-Model Reconstruction: The USAE achieved positive $R^2$ scores when reconstructing activations of one model using the encoded features of another, proving the existence of a shared feature space.
• Universality vs. Importance: The authors introduced Firing Entropy and Co-Fire Proportion metrics. They found a strong positive correlation ($r=0.89$ for high-frequency concepts) between a concept's universality (firing across models) and its importance (energy contribution to reconstruction). The most important concepts are the most universal.
• Consistency with Independent SAEs: Universal training recovers 23–38% of concepts found in independently trained SAEs (depending on the model). Crucially, the most universal concepts show higher overlap with independent SAEs, suggesting universal training naturally selects for fundamental visual features.

C. Model-Specific Nuances

• DinoV2: Exhibited unique, low-entropy concepts related to 3D geometry, perspective, and depth cues (e.g., vanishing points, converging lines). This aligns with its self-distillation training regime which emphasizes geometric consistency.
• SigLIP: Showed unique concepts firing for both visual and textual elements (e.g., the word "star" and the shape of a star), reflecting its vision-language training.

D. Generalization

The USAE trained on ImageNet demonstrated strong zero-shot generalization to out-of-distribution datasets (DTD and CelebA), maintaining reconstruction accuracy and concept consistency.

5. Key Contributions

1. Framework: Introduction of USAEs, the first method to jointly learn a shared, interpretable concept space across multiple distinct DNNs via end-to-end training.
2. Methodological Shift: Moving from post-hoc concept matching to simultaneous alignment, enabling efficient cross-model analysis.
3. Novel Application: Coordinated Activation Maximization, allowing simultaneous visualization of how different models encode the same concept.
4. Insights: Empirical evidence that the most important features in vision models are also the most universal, and that specific training regimes (like self-distillation in DinoV2) induce unique geometric representations.

6. Significance and Impact

• Risk Mitigation: Provides a unified lens for analyzing diverse models, crucial for managing risks in large-scale AI deployment where multiple models interact.
• Regulatory Compliance: Offers a structured way to audit model internals against regulatory frameworks requiring transparency.
• Scientific Insight: Supports the "Platonic Representation Hypothesis," suggesting that different models converge on similar high-level abstractions despite different training paths.
• Future Directions: Opens avenues for cross-model transfer learning, model distillation, and more robust interpretability tools that are not tied to a single architecture.

In summary, this paper establishes a foundational method for cross-model interpretability, demonstrating that a single, shared dictionary can effectively decode the internal logic of diverse deep neural networks, revealing both their commonalities and their unique architectural biases.