用语义最优传输实现 SAE 特征匹配与电路压缩

View PDF HTML (experimental)

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.

Submission history

From: Tue Minh Cao [view email]
[v1] Wed, 27 May 2026 14:54:58 UTC (5,371 KB)
[v2] Wed, 21 Oct 2026 19:48:02 UTC (5,381 KB)
[v3] Fri, 30 Oct 2026 02:01:16 UTC (5,381 KB)