Two SAE latents, one feature? A single energy-weighted transport plan now prices whole-dictionary disagreement and locates it latent by latent
Synopsis
Comparing two sparse autoencoders (SAEs, dictionaries that decompose model activations into sparse learned components) requires answering three things at once: how much two whole dictionaries differ, which latents account for that difference, and how much is sampling noise from finite documents. One energy-weighted optimal-transport plan connects them: its cost is the dictionary-level discrepancy, the cost decomposes exactly into latent contributions, and transported fractions describe how each latent's activation energy is allocated across counterparts, with joint asymptotic confidence intervals. Across four language models it locates differences between SAEs trained on identical data with different seeds and separates topic-dependent changes into a weight term and a cost-matrix term.
Interpretation
The difference between two SAE dictionaries is now a single number that can be split by latent, showing which latents contribute to the disagreement. An aggregate score alone could not locate disagreement, and separate pairwise similarities did not account for how all latents share the available counterparts. On TinyStories-1M, five SAEs trained on identical data with different seeds give a mean transport cost of 1.02784 on animals and 1.05424 on family; the prespecified difference is 0.02639, or 2.57% of the animals mean, and the interval accounts for document sampling and dependence between SAE pairs evaluated on the same documents.
The method treats each latent's normalized activation energy on evaluation documents as mass and transports one dictionary's mass onto the other, with a per-pair price set by the chosen similarity geometry. The plan's row and column sums are fixed to the two dictionaries' energy weights, so how each latent's mass is shared across several counterparts comes from the plan rather than a nearest-neighbour pick. Across the five TinyStories-1M seeds, the weight term is estimated at 0.01680 and the cost-matrix term at 0.00956; they sum to 0.02639 and account for 64% and 36% of the estimated difference, with both intervals excluding zero.
The same documents determine both the energy weights and the costs, so errors in costs and transported fractions are handled jointly rather than estimated separately. Earlier comparisons did not fold the dependence between both dictionaries' marginals and the costs, all estimated from the same documents, into the intervals. In a synthetic experiment with independent samples from a known population, 95% simultaneous document-bootstrap intervals for three fractions and one cost cover 94.5%–95.5%; resampling the fixed TinyStories document pools gives 91.5%–97.0%.
Perspective
The demonstrated setting is fixed trained SAEs with positive retained latent energies, fixed dictionary sizes and regularization, and statistical guarantees for document sampling; experiments cover TinyStories-1M, TinyStories-33M, GPT-2 small and Pythia-160M at widths 64, 3072 and 6144.
The authors note that asymptotic validity does not guarantee coverage when rare documents have a large effect on estimates, and in the rare-document experiment full-plan coverage falls well below 95% at smaller sample sizes; transported fractions depend on the specified cost, input distribution and regularization parameter and do not by themselves establish semantic identity or equal effects on model outputs; sensitivity to regularization and self-comparison costs also limits treating reported values as independent of those choices.