Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression
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
- Sparse Autoencoders Find Highly Interpretable Features in Language Models
- Scaling and evaluating sparse autoencoders
- Mechanistic Permutability: Match Features Across Layers
- Analyze Feature Flow to Enhance Interpretation and Steering in Language Models
- Sparse Feature Circuits: Discovering and Editing Interpretable Causal Graphs in Language Models
- Gemma Scope: Open Sparse Autoencoders Everywhere All At Once on Gemma 2
- Attribution Patching Outperforms Automated Circuit Discovery
- gpt-oss-120b & gpt-oss-20b Model Card
- Jumping Ahead: Improving Reconstruction Fidelity with JumpReLU Sparse Autoencoders
- BatchTopK Sparse Autoencoders
- Overcoming Sparsity Artifacts in Crosscoders to Interpret Chat-Tuning
- Do I Know This Entity? Knowledge Awareness and Hallucinations in Language Models
- I Have Covered All the Bases Here: Interpreting Reasoning Features in Large Language Models via Sparse Autoencoders
- The Pile: An 800GB Dataset of Diverse Text for Language Modeling
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks