PLOT: Progressive Localization via Optimal Transport in Neural Causal Abstraction

Abstract: Causal abstraction offers a principled framework for mechanistic interpretability, aligning a high-level causal model with the low-level computation realized by a neural network through counterfactual intervention analysis. Existing methods such as distributed alignment search (DAS) learn expressive subspace interventions, but the relevant neural site is unknown a priori, so finding a handle requires a computationally burdensome search over candidate sites. We introduce PLOT (Progressive Localization via Optimal Transport), a transport-based framework that localizes causal variables from the output effect geometry of abstract and neural interventions. PLOT fits an optimal transport coupling between abstract variables and candidate neural sites, yielding a global soft correspondence that can be calibrated into intervention handles. In simple settings, a single coupling over individual neurons suffices. In larger models, PLOT is applied progressively, moving from coarse sites such as tokens, timesteps, or layers to finer supports such as coordinate groups or PCA spans, and optionally guiding DAS based on the localized signal. Across experiments of increasing complexity, transport-only PLOT handles are exceedingly fast and competitive on accuracy, while PLOT-guided DAS reaches DAS-level accuracy at a fraction of full DAS runtime, providing an efficient localization engine for causal abstraction research at scale.

• The paper presents PLOT, a framework that leverages optimal transport to efficiently localize neural representations of abstract causal variables.
• It employs a hierarchical localization strategy, refining intervention handles from coarse layer-level sites to fine-grained subspaces.
• Empirical results show that PLOT matches Distributed Alignment Search accuracy while reducing runtime by up to 30×, boosting interpretability and scalability.

Progressive Localization via Optimal Transport for Neural Causal Abstraction

Introduction and Context

Mechanistic interpretability in neural networks seeks to establish connections between high-level, interpretable latent variables and their concrete representations within neural architectures. Within the causal abstraction framework, the underlying question is whether interventions on abstract variables—defined in an idealized causal model—can be matched by interventions at specific internal sites within a neural network, as evaluated via counterfactual analysis. Current distributed alignment methods, notably Distributed Alignment Search (DAS), parameterize and learn such internal interventions but are hampered by computational challenges stemming from exhaustive site and subspace searches.

"PLOT: Progressive Localization via Optimal Transport in Neural Causal Abstraction" (2605.06979) introduces Progressive Localization via Optimal Transport (PLOT), a framework that leverages optimal transport (OT) to efficiently localize the neural realization of abstract variables. PLOT matches transformations (effects) induced by interventions in the abstract space with those in the network, using an OT-based soft correspondence. This process can be applied hierarchically—first to coarse site families (e.g., tokens, layers), then to more fine-grained encodings (e.g., principal components)—resulting either in direct, executable intervention handles or in the focused restriction of search spaces for methods like DAS.

Methodological Framework

Causal Abstraction & Counterfactual Correspondence

The central criterion for causal abstraction is whether neural interventions, derived from a hypothesized correspondence, reliably reproduce the effects of abstract counterfactual interventions across a calibration bank of input pairs. The authors implement this using interchange interventions, constructing "effect signatures" that quantify the changes in network output following localized interventions.

Optimal Transport-Based Localization

Given a set of candidate neural sites and abstract variables, PLOT computes the effect signatures under both abstract and neural interventions, projects them into a common feature space, and fits a discrete optimal transport coupling (entropic regularization is used for tractability). The resulting coupling matrix denotes a soft alignment: for each abstract variable, the entries rank the candidate neural sites by their ability to reproduce the causal effect.

PLOT supports both classic entropic OT and unbalanced OT (UOT), the latter allowing for the exclusion of irrelevant neural sites from the correspondence—a critical property when the underlying neural representation is highly distributed or sparse.

Figure 1: PLOT as a progressive localization engine. PLOT moves from coarse localization across sites (e.g., tokens, layers) to finer resolutions (e.g., PCA spans), calibrating effect-driven handles at each stage.

Progressive and Hierarchical Localization

In settings involving large or layered networks, PLOT is applied hierarchically:

1. Coarse Localization: Identify broad regions (tokens, layers, time steps) with strongest correspondences.
2. Refinement: Within the high-mass regions, further dissect into coordinate groups, subspaces, or PCA factors.
3. Calibration: Convert the support into actionable neural intervention handles, optionally passing the reduced locality as the initialization/focus for DAS or similar subspace learners.

   Figure 2: HEQ site-level intervention handles learned by OT (before and after calibration) and DAS (in the canonical and rotated basis). OT mass can be spread across layers pre-calibration but is sparsified post-calibration.

The framework is agnostic to neural architecture and incorporates feature-extractive mappings for generality.

Empirical Evaluation

Hierarchical Equality (HEQ)

In the HEQ scenario, abstract variables correspond to explicit logical relationships (equality of input pairs). The neural model is a small MLP; PLOT is used to directly localize the relevant neurons for each variable.

• Accuracy: Single-stage PLOT achieves average exact counterfactual accuracy of 0.991, closely matching DAS (0.995).
• Efficiency: PLOT requires on average 4.4 seconds per run, which is approximately 30× faster than DAS (131 seconds).
• Interpretability: OT-based handles are highly localized; calibration sparsifies the correspondence to a small set of neurons.

  Figure 3: End-to-end accuracy and runtime summary. PLOT matches DAS-level accuracy while being substantially faster.

Binary Addition

For 4-bit binary addition with a GRUCell, abstract variables (internal carry bits) are represented non-trivially in distributed GRU states.

• Two-Stage Localization: (1) OT localizes each carry to a recurrent timestep; (2) Further refinement within the state via native coordinates, PCA projections, or localized DAS.
• Findings: PCA-based OT handles yield high accuracy (0.941), and PLOT-guided DAS matches full-DAS accuracy (0.974–0.984), with substantial reductions in runtime.
• Scale: Runtime improvements are amplified as dimensionality increases.

  Figure 4: Internal-carry exact accuracy and runtime summary for binary addition. PLOT-guided methods achieve high accuracy at a fraction of full DAS runtime.

  

  Figure 5: Binary-addition handles for $d=16$: PLOT (native, PCA) and DAS handles for carries $C_1$, $C_2$, $C_3$. OT mass is compact and well-localized post-calibration.

Multiple-Choice Question Answering (MCQA) with Gemma-2-2B

This large-scale benchmark demands identification of high-level variables (answer pointer, answer token) inside a deep transformer. PLOT is applied across three stages: layers (UOT), subspaces (native, PCA), and DAS-restricted refinement.

• Layer Localization: UOT successfully avoids assigning mass to irrelevant layers, consistently identifying task-relevant regions.
• Direct Handles: OT-based interventions are within 2% of full DAS on accuracy, but with 65× lower runtime.
• DAS Guidance: Restricting DAS to PLOT-localized layers/dimensions yields matching or better accuracy at over an order of magnitude faster runtime.

Theoretical Insights and Practical Implications

PLOT establishes OT as an efficient, architecture-agnostic meta-alignment mechanism for neural causal abstraction. By moving from an exhaustive search toward guided, effect-driven localization, it decouples global search from local expressiveness. Empirically, it shows that globally soft but progressively sharpened correspondences suffice to match the performance of rotation-parametric methods like DAS, but with drastically reduced computational load.

Notably, the use of UOT to ignore irrelevant support is essential in high-dimensional settings, mitigating the rigid, all-mass assignment inherent to standard OT and aligning with the theoretical perspective that neural real-world systems likely distribute abstract features non-uniformly.

Limitations and Future Directions

While PLOT is shown to be efficient and robust, there are clear axes for extension:

• Finer Hierarchies: Incorporation of deeper multi-scale site families, e.g., neurons $\to$ attention heads $\to$ blocks.
• Batching/Scalability: Adapting to very large models (e.g., >10B parameters) with distributed computation and approximate effect signatures [wang2023neuralegw, tsur2025nemot].
• Integration with Model Editing: PLOT-guided interventions could be used in online model repair/editing scenarios, enhancing the reliability of fact localization/edit [meng2022rome].
• Improved Calibration: Online or adaptive calibration schemes could further automate the extraction of most effective intervention supports.

Figure 6: Stability of PLOT in the entropic regularization parameter $\varepsilon$. Performance is robust except at very small values, which over-sharpen the coupling.

Conclusion

PLOT substantiates progressive, transport-based soft correspondence as a scalable alignment engine for neural causal abstraction. In all evaluated settings—from low-dimensional logic to deep transformers—PLOT rapidly and accurately identifies causal variables, efficiently produces high-quality intervention handles, and robustly guides local subspace learners. Its principled hierarchical and unbalanced localization is theoretically sound and pragmatically effective, making it a compelling foundation for future mechanistic interpretability and model editing pipelines in large-scale neural architectures.