Universal Sparse Autoencoders: Interpretable Cross-Model Concept Alignment

TL;DR

Universal Sparse Autoencoders (USAEs) train a single overcomplete TopK sparse autoencoder whose shared latent code can be encoded from — and decoded back to — the activations of multiple pretrained vision models at once (DinoV2, SigLIP, supervised ViT). The training trick is round-robin: at each step pick one model, encode its features into the shared code, then decode through every model’s dictionary and sum the reconstruction losses. The result is a single concept dictionary in which the same dimension co-fires across architectures for low- to high-level visual concepts (colors, textures, “dog face”, “mouth-open animal jaws”), with concept universality positively correlated to importance. The framework also enables coordinated activation maximization — visualizing how one shared concept index materializes differently in each model.

Key claims

A single TopK SAE with model-specific encoder/decoder pairs and a shared $K$ -sparse code can reconstruct three heterogeneous vision backbones (DinoV2, SigLIP, ViT-supervised) with positive off-diagonal $R^2$ , i.e. cross-model reconstruction is non-trivial [§4.2, Fig. 5].
The training procedure picks one model uniformly per step, encodes its batch into the shared code, then computes reconstruction loss against all decoders before backprop — chosen as a memory/throughput compromise vs. full-multi-model loss or random pairs [§3.2, Fig. 3, Eq. 5–6].
The learned dictionary is bimodal in firing entropy: >1000 dictionary atoms fire with maximum cross-model entropy (fully universal), a smaller group favors two of three models, and very few are model-specific [§4.3, Fig. 6(a)].
Co-Fire Proportion (fraction of a concept’s fires that coincide across all three models on the same input) is moderately correlated with energy-based concept importance overall ( $R \approx 0.5$ ), and strongly correlated once restricted to concepts with >250 co-fires — “the most important concepts are also the most universal” [§4.3, Fig. 6(c)].
DinoV2 carries more unique features than SigLIP or ViT: lowest CFP (0.266 vs 0.344 / 0.326), highest self-reconstruction $R^2$ , and a separate low-entropy cluster of geometric concepts; SigLIP’s idiosyncratic cluster is textual features [§4.2, §4.3, Appendix A.2.1].
Hungarian-aligned cosine similarity between USAE decoder atoms and independently trained per-model SAE atoms is far above a same-statistics Gaussian baseline (AUC 0.30–0.41 vs 0.13), but the fraction of high-similarity matches is modest (23–38% with cos > 0.5) — universal training discovers concepts that single-model SAEs miss [§4.4, Fig. 7, Table 1].
Coordinated activation maximization — optimizing one input per model to maximize the same USAE dimension — produces visually coherent yet model-specific renderings of shared concepts (e.g. “curves”, “round objects”, “objects from above”), with DinoV2’s optimizers consistently appearing at larger spatial scale [§4.5, Fig. 8].
Dictionary size $D = 32{,}768$ ; trained on ImageNet final-layer features of three timm backbones; total training cost is ~3 days on a single RTX 6000 [§4 Implementation Details].

Method

A standard TopK SAE on one model takes activations $A \in \mathbb{R}^{N \times d}$ , computes a sparse code $Z = \mathrm{TopK}(A W_E)$ , and reconstructs $\hat{A} = Z W_D$ under Frobenius loss. USAE generalizes this to $M$ models by giving each model its own encoder $W_E^{(m)} : \mathbb{R}^{d_m} \to \mathbb{R}^{D}$ and decoder $W_D^{(m)} : \mathbb{R}^{D} \to \mathbb{R}^{d_m}$ , but the TopK-sparse latent $Z \in \mathbb{R}^{N \times D}$ is shared across all models [§3.1, Eq. 3–4]. At each training step a single source model $m^\star$ is sampled, a batch of its activations is encoded to $Z = \mathrm{TopK}(A^{(m^\star)} W_E^{(m^\star)})$ , and $Z$ is decoded through every model’s $W_D^{(m)}$ . The total loss sums Frobenius reconstruction error across all $M$ targets; gradients update only the chosen $W_E^{(m^\star)}$ plus all decoders [§3.2, Eq. 5–6]. This balances per-step memory against the slow convergence of randomly paired single-encoder/single-decoder updates.

Concept importance is measured by reconstruction energy — squared norm of the average activation of an atom times its decoder vector [§4.1, Eq. 9]. Universality is measured by (a) firing entropy across the $M$ per-model firing rates of an atom, normalized to $[0,1]$ [§4.3, Eq. 12–13], and (b) co-fire proportion: the fraction of an atom’s fires that coincide across all $M$ models on the same token [§4.3, Eq. 11, Eq. 14]. Coordinated activation maximization fixes a single dictionary index $z$ and optimizes one input image per model to maximize $\mathrm{TopK}(\phi^{(m)}(x^{(m)}) W_E^{(m)})_z$ , using Maco-style phase-only optimization with magnitude prior [§3.3, Eq. 7–8].

Results

Cross-model $R^2$ : positive off-diagonal entries in the 3×3 source→target reconstruction matrix; diagonals highest, with DinoV2 self-reconstruction the strongest [§4.2, Fig. 5].
Single-model SAE baselines ( $R^2$ , paper’s own runs): DinoV2 ≈ 0.8, SigLIP ≈ 0.7, ViT ≈ 0.7 [§4.2].
USAE-vs-independent-SAE concept alignment (cos > 0.5 after Hungarian matching): ViT 0.38, DinoV2 0.26, SigLIP 0.23; Gaussian baseline 0.00 [§4.4, Table 1, Fig. 7].
Concept-universality / importance correlation: $R \approx 0.5$ over all 32,768 dimensions, $R \gg 0.5$ once restricted to the high-co-fire tail [§4.3, Fig. 6(c)].
Per-model co-fire proportion (top concepts): SigLIP 0.344, ViT 0.326, DinoV2 0.266 [§4.3, Fig. 6(b)].

Why it’s interesting

Slots cleanly into Mechanistic Interpretability as the cross-model SAE primitive — the open question flagged on that concept page (“when is parameter-space decomposition the right primitive vs activation-space SAEs?”) gets a new orthogonal axis: when should the SAE be per-model vs shared across models? The bimodal firing-entropy result and the universality⇄importance correlation are concrete empirical support for the Platonic-representation hypothesis (Huh et al. 2024) that the page’s TL;DR alludes to but does not cite. Complements Natural Language Autoencoders Produce Unsupervised Explanations of LLM Activations — both move beyond per-model atom dictionaries (NLAs to free-text readouts within one model; USAEs to a single dictionary across many models) — and contrasts with Training large language models on narrow tasks can lead to broad misalignment‘s “shared persona direction” story by showing that vision models exhibit a directly analogous shared-concept geometry on a much larger atom budget.