Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression

arXiv:2605.28567 · cs.LG, cs.AI · Submitted 2026-05-27 · Read on arXiv

cs.LG, cs.AI

Submitted: 2026-05-27

Updated: 2026-10-21

Comments: Neurips 2026

License: http://creativecommons.org/licenses/by/4.0/

The gist: Sparse autoencoders (SAEs) have become a central tool for interpreting language models.

Terminology

Abstract

Sparse autoencoders (SAEs) have become a central tool for interpreting language models. However, two key SAE analyses that remain difficult to scale are (1) matching semantically similar features across multi-layers and (2) compressing large feature circuits into interpretable supernodes. Although these have been treated as separate problems, we show that both are instances of a more fundamental challenge, which we frame as the estimation of semantic distances between SAE features that lie on different activation manifolds. We introduce a distributional framework for this problem, in which each feature is represented not by a single decoder vector like in the literature, but by an activation-weighted distribution over the hidden states that express it. By projecting these distributions into a shared reference space and comparing them with Wasserstein distance, our method provides a unified semantic metric for cross-layer feature comparison. We prove that our representation is invariant to activation rescaling, stable under perturbations, and recovers true matches under finite-sample margin conditions. Empirically, our method outperforms decoder-vector and LLM-based baselines and captures subtle functional distinctions between related features. Notably, our method compresses large feature circuits into interpretable supernodes automatically.

Sources

Related papers