💬 NLP

Geometry-Adaptive Explainer for Faithful Dictionary-Based Interpretability under Distribution Shift

This paper introduces the Geometry-Adaptive Explainer (GAE), a gradient-free method that realigns dictionary-based interpretability tools with out-of-distribution activation subspaces to significantly improve causal faithfulness under distribution shifts.

Original authors: Sungjun Lim, Heedong Kim, Andrew Lee, Kyungwoo Song

Published 2026-05-22

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sungjun Lim, Heedong Kim, Andrew Lee, Kyungwoo Song

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a highly skilled translator who has spent years learning a specific dialect of a language. They are excellent at translating stories written in that dialect. However, one day, they are asked to translate a story written in a slightly different dialect where the grammar rules and word order have subtly shifted.

Because the translator is so used to the old rules, they try to force the new story into the old structure. The result? The translation is confusing, inaccurate, and doesn't capture the true meaning of the new story.

This is exactly the problem the paper "Geometry-Adaptive Explainer for Faithful Dictionary-Based Interpretability under Distribution Shift" tackles, but instead of a human translator, it's about AI models and the tools we use to understand how they think.

Here is a breakdown of the paper's ideas using simple analogies:

1. The Problem: The "Rigid Map"

AI researchers use tools called Dictionary-Based Explainers (like Sparse Autoencoders) to understand what's happening inside a neural network. Think of these explainers as a dictionary or a map.

How they work: The map is drawn based on how the AI behaves when it sees "normal" data (like standard English sentences). The map tells us, "When the AI sees this pattern, it activates this specific feature."
The Issue: When the AI encounters Out-of-Distribution (OOD) data (like a new slang term, a different topic, or a weirdly phrased sentence), the internal "geometry" of the AI changes. It's as if the landscape itself has rotated.
The Result: The old map no longer fits the new landscape. The AI is using different "directions" to think, but the explainer is still trying to read the map from the old orientation. This creates a "Faithfulness Gap." The explanation is no longer faithful to what the AI is actually doing.

2. The Discovery: It's a Geometry Problem

The authors realized this isn't just a random error; it's a geometric misalignment.

The Analogy: Imagine the AI's internal thoughts are a cloud of points floating in 3D space. When the data changes, the whole cloud rotates.
The old explainer is a rigid frame that was built to fit the cloud in its original position.
When the cloud rotates, the frame no longer captures the points correctly. The paper proves that the more the cloud rotates (the "second-moment shift"), the bigger the gap between the frame and the cloud becomes, and the worse the explanation gets.

3. The Solution: GAE (The "Rotating Frame")

The authors propose a new method called GAE (Geometry-Adaptive Explainer). Instead of throwing away the old map and drawing a brand new one from scratch (which takes a lot of time and computing power), GAE does something clever:

Step 1: The Spin. It takes the old map and physically rotates it to match the new orientation of the AI's thoughts. It uses a mathematical trick (Orthogonal Procrustes) to find the perfect angle to align the old frame with the new cloud of data.
Step 2: The Fine-Tune. Once the frame is rotated, it makes tiny adjustments to the individual "rulers" inside the frame so they fit the specific points of the new data perfectly, without breaking the original structure of the map.

The Best Part: GAE does this without any training.

Traditional methods are like trying to learn a new language by reading a million books (taking hours or days of computer time).
GAE is like looking at a few sentences, realizing the grammar shifted, and instantly rotating your mental map to match. It takes seconds and requires no gradient updates (no heavy math training).

4. The Results: Fast and Accurate

The paper tested GAE on large language models (like GPT-2 and Pythia) with different types of "shifts" (new time periods, financial documents, and adversarial tricks).

Speed: While other methods took minutes or hours to adapt, GAE finished in under 3 seconds.
Accuracy: Even though it didn't "train" in the traditional sense, GAE produced explanations that were more faithful (more accurate to the AI's actual behavior) than methods that spent hours retraining the model.
The "Circuit" Test: In one experiment, they looked at how the AI decided to predict the word "American." The old map (Fixed) pointed the AI in the wrong direction. GAE rotated the map, and suddenly, the explanation correctly pointed toward the concept of nationality, even though the underlying "features" (the words the AI noticed) stayed exactly the same.

Summary

The paper argues that when AI models face new types of data, their internal "thinking space" rotates. Our old tools for understanding them get stuck in the old orientation.

GAE is a quick, math-based fix that simply rotates our understanding tools to match the new reality. It's a way to keep our explanations accurate and trustworthy without needing to spend days retraining the tools, making it a practical solution for keeping AI interpretable in a changing world.

Technical Summary: Geometry-Adaptive Explainer for Faithful Dictionary-Based Interpretability under Distribution Shift

Problem Statement

Mechanistic interpretability relies on dictionary-based explainers, such as Sparse Autoencoders (SAEs) and Transcoders, to decompose neural network hidden activations into sparse, interpretable feature directions. A critical requirement for these explanations is faithfulness: the explainer must accurately reflect the computations the model actually performs.

However, existing explainers are typically trained on in-distribution (ID) data. When a model encounters out-of-distribution (OOD) inputs, the geometric structure of its hidden activations changes. The paper identifies that distribution shift rotates the OOD-active subspace (the directions the model actively uses) away from the ID-trained explainer subspace. This misalignment, termed the faithfulness gap, causes ID-trained dictionaries to fail in reconstructing OOD activations faithfully. Prior work has addressed this through heuristic retraining or residual modules but lacks a principled diagnosis of the underlying geometric misalignment or a method to correct it without full retraining.

Methodology: Geometry-Adaptive Explainer (GAE)

The authors propose the Geometry-Adaptive Explainer (GAE), a closed-form, post-hoc method that realigns an ID-trained explainer with the OOD-active subspace using only unlabeled OOD activations. GAE operates without gradient updates or iterative training.

Theoretical Foundation

The paper formalizes the problem using the second-moment matrix of hidden activations, $M_e = \mathbb{E}[h(X)h(X)^\top]$ .

Faithfulness Gap ( $\Delta$ ): Defined as the Frobenius distance between the explainer's subspace ( $\Pi_{dec}$ ) and the OOD-active subspace ( $\Pi_{OOD}$ ).
Geometric Mechanism: The authors prove (Proposition 1) that the faithfulness gap of an ID-trained explainer is upper-bounded by the magnitude of the second-moment shift ( $\|M_{OOD} - M_{ID}\|_F$ ).
Control of Degradation: They further prove (Proposition 2) that the explainer-dependent component of OOD faithfulness loss is tightly controlled by the square of the faithfulness gap. Thus, reducing the gap is necessary and sufficient for improving OOD faithfulness.

Algorithm

GAE adapts the dictionary in two steps:

Step 1: Subspace Rotation (Orthogonal Procrustes):
The ID-trained dictionary $W_{dec}^{ID}$ is rotated to align its column space with the estimated OOD-active subspace $\hat{\Pi}_{OOD}$ (derived from the top- $r$ eigenvectors of the OOD second-moment matrix). The rotation matrix $T^\star$ is chosen to minimize the distance to the original dictionary, preserving the original feature structure as much as possible. This step eliminates the faithfulness gap ( $\Delta = 0$ ) by construction.
Step 2: Constrained Decoder Refit:
While Step 1 aligns the subspace, it does not optimize sample-level reconstruction. Step 2 refits the decoder weights via a closed-form constrained ridge regression. This step minimizes reconstruction error on OOD activations while penalizing deviations from the Step 1 rotated dictionary and ensuring the decoder remains within the OOD-active subspace.

The entire pipeline is analytical, requiring no backpropagation.

Key Contributions

Geometric Diagnosis: The paper identifies the faithfulness gap as the geometric mechanism driving OOD faithfulness degradation in dictionary-based explainers. It provides theoretical bounds showing that this gap grows with second-moment shift and directly controls the reducible portion of OOD faithfulness loss.
GAE Algorithm: A training-free, closed-form method that realigns the explainer's dictionary with the OOD-active subspace. It offers a theoretical guarantee (Theorem 1) that the adapted explainer strictly improves upon the unadapted ID explainer, with improvement growing quadratically with the initial misalignment.
Empirical Superiority: GAE matches or surpasses all training-based baselines (including full retraining from scratch and fine-tuning) in causal faithfulness across multiple models and OOD settings, despite requiring zero gradient computations and only $\sim2,000$ unlabeled OOD samples.

Experimental Results

The authors evaluated GAE on GPT-2 Small and Pythia-1.4B using both Transcoders and Top-K SAEs. The evaluation covered three types of distribution shifts:

Temporal: Web text collected after the model's pretraining cutoff (FineWeb).
Domain: Financial filings with specialized vocabulary (Edgar).
Adversarial: Hallucination-inducing prompts (HaluEval).

Performance Metrics:

Causal Faithfulness: Measured via Normalized AOPC (nAOPC) and Normalized Comprehensiveness (nComp).
Reconstruction Quality: Measured via Delta Cross-Entropy ( $|\Delta CE|$ ).

Key Findings:

GAE vs. Training-Based Baselines: GAE outperformed methods like Retrain (training from scratch on OOD data) and Finetune (warm-started fine-tuning) in the majority of settings. For instance, on GPT-2 Transcoders under adversarial shift, GAE improved nComp by 4.6% over the strongest baseline (SAEBoost) and reduced reconstruction error ( $|\Delta CE|$ ) by 93%.
Efficiency: While training-based baselines required processing millions of tokens and hours of GPU time, GAE completed adaptation in under 3 seconds (0.5s for GPT-2, 2.9s for Pythia-1.4B) using only a single forward pass over 2,048 samples.
Mechanism Validation: Ablation studies confirmed that Step 1 (subspace rotation) successfully closed the faithfulness gap, while Step 2 was necessary to restore causal coherence (feature attribution) without sacrificing the geometric alignment.
Case Studies: In circuit attribution tasks (e.g., predicting "American" or "Henry"), GAE corrected the decoder's directionality, ensuring that identified features pushed the model toward the correct semantic class, whereas the fixed ID explainer often pointed away from the target class.

Significance and Claims

The paper claims that GAE provides a principled, efficient, and theoretically grounded solution to the problem of OOD faithfulness in mechanistic interpretability. By framing the issue as a geometric misalignment rather than a data coverage problem, the authors demonstrate that faithful explanations can be restored without the computational cost of retraining.

The authors note that GAE is particularly significant for on-the-fly adaptation scenarios where deployment distributions change dynamically, and retraining is infeasible. They modestly acknowledge limitations, specifically that the method relies on a top- $r$ SVD truncation (discarding residual feature information) and has not yet been evaluated on larger-scale models. However, the results suggest that geometric realignment is a powerful and underutilized lever for maintaining interpretability robustness.