Research map / Learning theory and algorithms / Optimal transport

Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression

NeurIPS 2026 · Learning theory and algorithms › Optimal transport · approach: Neural optimal transport · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Frames SAE feature matching across layers and circuit compression as cross-manifold semantic distance, representing each feature as an activation-weighted distribution over hidden states compared with Wasserstein distance in a shared reference space.

Problem

Matching semantically similar SAE features across layers and compressing large feature circuits into interpretable supernodes are hard to scale; existing decoder-vector matching fails across dissimilar layer distributions and manual supernode grouping does not scale.

Approach

Represent each SAE feature as an activation-weighted empirical distribution over the hidden states that activate it, project these into a shared reference space (target layer hidden states for matching; concatenated all-layer hidden states for circuits), and use optimal transport (Wasserstein distance) as the semantic distance for nearest-neighbor matching and clustering.</multi_core_idea>

Key results

• On GPT2 layer 5 to layer 6 matching, the method reaches LLM Eval 2.53 and Acc 68.2%, versus 2.46 and 66.2% for FeatFlow.
• On Gemma-2-2B layer 12 to layer 13 matching, Acc is 58.6% versus 20.4% for FeatFlow.
• Theory provides activation-rescaling invariance, Wasserstein stability under measure perturbation, and margin-based recovery guarantees.

Surprising findings

• Naive single-centroid projection onto the same reference space is competitive with the full Wasserstein method on several matching setups.

Limitations

• Uses plain L2 ground cost, assuming flat Euclidean geometry.
• Relies on LLM judges and small evaluation sets (100 pairs per setup, 200 circuits).

Experimental setup

Datasets: The Pile

Benchmarks: feature matching (GPT2, Gemma-2-2B) circuit compression digit addition subtle-feature matching

Metrics: LLM Eval Accuracy CE loss Variance Explained human blind-ranking preference

Baselines: Match Match ML FeatFlow FeatFlow ML Attrib LLM Sel Naive

Builds on: SAE Match Feature Flow Attribution Patching Sinkhorn algorithm

Closest papers by meaning

• Scaling and evaluating sparse autoencoders (ICLR 2025 · similarity 0.91)
• SAEBench: A Comprehensive Benchmark for Sparse Autoencoders in Language Model Interpretability (ICML 2025 · similarity 0.91)
• Compute Optimal Inference and Provable Amortisation Gap in Sparse Autoencoders (ICML 2025 · similarity 0.90)
• Sparse Autoencoders, Again? (ICML 2025 · similarity 0.90)
• Improving Sparse Decomposition of Language Model Activations with Gated Sparse Autoencoders (NeurIPS 2024 · similarity 0.90)
• Universal Sparse Autoencoders: Interpretable Cross-Model Concept Alignment (ICML 2025 · similarity 0.90)