Technical Summary: Semantic Optimal Transport for Sparse Autoencoder Feature Matching and Circuit Compression

1. Problem Definition

Sparse Autoencoders (SAEs) have emerged as a primary tool for interpreting Large Language Models (LLMs) by decomposing latent spaces into interpretable features. However, two critical challenges in scaling SAE analysis remain:

1. Cross-Layer Feature Matching: Identifying semantically similar features across different layers of an LLM. Existing methods typically rely on comparing decoder vectors. This approach fails when matching features across non-consecutive layers because decoder vectors are trained to reconstruct distinct activation distributions, making direct vector comparison inaccurate.
2. Feature Circuit Compression: Feature circuits (causal graphs connecting features across layers) often contain hundreds of nodes, making manual inspection impractical. Current compression methods rely on manually grouping semantically similar nodes into "supernodes," a process that does not scale.

The core problem identified by the authors is the estimation of semantic distance between features residing on different activation manifolds. Existing approaches treat features as single vectors (e.g., decoder weights), which discards the rich context of where and how strongly a feature activates, leading to information loss and an inability to handle multi-manifold comparisons.

2. Methodology

The authors propose a distributional framework that reframes feature comparison as an Optimal Transport (OT) problem.

2.1 Feature Representation as Distributions

Instead of representing a feature by a single decoder vector, the method represents each feature $i$ at layer $\ell$ as an activation-weighted empirical distribution over the hidden states where it activates.

• Given a corpus of token positions $T$, the method selects the top-$K$ token positions where feature $i$ activates most strongly.
• Each selected token $t$ is assigned a weight $w_{i,t}$ proportional to its activation magnitude.
• The feature is represented as a distribution $\hat{\mu}_{i,T} = \sum w_{i,t} \delta_{x_t}$, where $x_t$ is the hidden state at token $t$.

2.2 Shared Reference Space

To compare features across different layers (which reside on different manifolds), the method projects these layer-specific distributions into a shared reference space $\mathcal{R}$.

• For cross-layer matching, the reference space is typically the hidden state space of the target layer. Activations from the source layer are projected to the corresponding token positions in the target layer.
• For circuit compression, the reference space is the concatenation of hidden states across all layers.

2.3 Semantic Distance via Wasserstein Distance

Once projected into the shared space, the semantic distance between two features is computed using the Wasserstein distance (Earth Mover's Distance).

• The distance $W_c(\hat{\mu}, \hat{\nu})$ measures the minimum cost to transport the probability mass of one feature's activation distribution to the other.
• The ground cost $c(z, z')$ measures the expense of matching two context representations in the shared space.
• This approach captures the geometric structure of activation patterns rather than just point-wise vector similarity.

2.4 Downstream Tasks

• Feature Matching: For a target feature, the method finds the source feature with the minimum Wasserstein distance (nearest neighbor).
• Circuit Compression: The method treats circuit nodes as points in the distributional space and applies clustering (e.g., Agglomerative clustering) to group nodes with small pairwise Wasserstein distances into supernodes, minimizing within-cluster semantic distance.

3. Key Contributions and Theoretical Guarantees

The paper provides three main theoretical guarantees for the proposed framework:

1. Activation Rescaling Invariance: The projected feature distribution is invariant to global scaling of activation values. Since the representation relies on the relative ranking and proportional distribution of activations (Top-K selection and normalized weights), arbitrary scalar shifts in activation magnitudes do not alter the semantic distance.
2. Stability under Perturbations: The Wasserstein distance between feature distributions is stable under smooth perturbations of the distributions. The change in distance is bounded by the Lipschitz constants of the ground cost and the Wasserstein distance between the perturbed and original distributions. This ensures robustness against noise in token selection or SAE retraining.
3. Margin-Based Recovery: The authors establish conditions under which the empirical nearest-neighbor matching and semantic grouping perfectly recover the true (oracle) matches. If the "margin" (the gap between the best match and the second-best match) exceeds twice the estimation error, the empirical decision is guaranteed to be correct.

4. Experimental Results

The method was evaluated on SAEs trained on GPT-2 and Gemma-2-2B, comparing against baselines including SAE Match (decoder vector L2), Feature Flow (decoder vector cosine), Attribution Patching, and LLM-based selection.

• Feature Matching: The proposed method outperformed all baselines in LLM Evaluation (semantic similarity ratings) and Matching Accuracy across consecutive, middle-to-last, and first-to-last layer setups. Crucially, it maintained competitive performance on downstream metrics (Cross-Entropy loss and Variance Explained), demonstrating that it identifies features with similar functional roles, not just similar activation tokens. It notably succeeded in matching subtle "digit addition" features where other methods failed.
• Circuit Compression: The method achieved higher accuracy in compressing feature circuits into intuitive supernodes compared to manual grouping and other baselines. It successfully identified modular feature groups (e.g., "induction heads" and "token detection") that consistently perform the same functional role across different circuits.
• Ablation Study: Experiments confirmed that representing features as distributions preserves more semantic information than collapsing them into a single centroid vector, even when projected onto the same manifold.

5. Significance and Claims

The paper claims to provide a unified semantic metric that solves both cross-layer feature matching and circuit compression as instances of the same fundamental problem: estimating semantic distance across manifolds.

• Scalability: By automating the grouping of features into supernodes, the method removes the reliance on human labeling, making the analysis of large-scale feature circuits feasible.
• Robustness: The distributional approach captures the "context" of feature activation, allowing for accurate matching even when features evolve significantly across layers (dissimilar distributions), a scenario where vector-based methods fail.
• Interpretability: The ability to automatically compress complex circuits into interpretable supernodes enhances the understanding of model mechanisms, as demonstrated in case studies on "Dallas" (geography), "Rhyming," and "Opposite" circuits.

The authors position this work as a foundational step toward scalable, automated mechanistic interpretability, moving beyond the limitations of vector-centric feature comparison.