Combining Causal Models for More Accurate Abstractions of Neural Networks

TL;DR

This work addresses the challenge of partial faithfulness in mechanistic interpretability by proposing a framework that combines multiple simple high-level causal models into a single input-dependent abstraction. Using causal models, exact and approximate transformations, interventional tests, and distributed alignment search, the authors show how to construct a faithfulness-strength trade-off and to greedily partition inputs among models to maximize explanatory power. They apply the framework to GPT-2-small on arithmetic and boolean logic toy tasks, demonstrating that combined models yield stronger, more faithful interpretability hypotheses than any single model, especially at high faithfulness levels. The approach offers a principled path to richer, more robust mechanistic explanations and has potential to scale to larger models and more complex tasks, aiding debugging and trust in neural computations.

Abstract

Mechanistic interpretability aims to reverse engineer neural networks by uncovering which high-level algorithms they implement. Causal abstraction provides a precise notion of when a network implements an algorithm, i.e., a causal model of the network contains low-level features that realize the high-level variables in a causal model of the algorithm. A typical problem in practical settings is that the algorithm is not an entirely faithful abstraction of the network, meaning it only partially captures the true reasoning process of a model. We propose a solution where we combine different simple high-level models to produce a more faithful representation of the network. Through learning this combination, we can model neural networks as being in different computational states depending on the input provided, which we show is more accurate to GPT 2-small fine-tuned on two toy tasks. We observe a trade-off between the strength of an interpretability hypothesis, which we define in terms of the number of inputs explained by the high-level models, and its faithfulness, which we define as the interchange intervention accuracy. Our method allows us to modulate between the two, providing the most accurate combination of models that describe the behavior of a neural network given a faithfulness level.

Figures (8)

• Figure 1: Interchange Intervention Accuracy (IIA) for the intermediate variables $P$ of different high-level causal models across GPT layers when using a 256-dimensional alignment subspace. Early layers (1-4) show high accuracy for models representing separate variables ($\mathcal{M}^X$, $\mathcal{M}^Y$, $\mathcal{M}^Z$), indicating no summation has occurred. The complete summation model ($\mathcal{M}^{XYZ}$) maintains perfect accuracy across all layers. From layers 5-11, there is a gradual transition where partial summation models ($\mathcal{M}^{XY}$, $\mathcal{M}^{YZ}$, $\mathcal{M}^{XZ}$) show intermediate accuracy, suggesting the network does not transition between computational states in discrete steps.
• Figure 2: Analysis of layer 7 causal models in fine-tuned GPT model on arithmetic tasks. The x-axis quantifies faithfulness with the interchange intervention accuracy achieved for the hypothesis. The y-axis quantifies strength as the proportion of inputs assigned to a non-trivial high-level causal model. Results compare the performance of individual causal models versus their combinations across different faithfulness thresholds. Combined models demonstrate stronger hypotheses at high faithfulness levels (0.9-1.0), while all models converge in performance at lower faithfulness thresholds (0.8).
• Figure 3: Two hypotheses for solving the boolean task.
• Figure 4: Interchange Intervention Accuracy (IIA) for the intermediate variables $P$ of different high-level causal models of the boolean logic task across the layers of GPT 2-small when searching within 256-dimensional subspaces of the neural representations.
• Figure 5: Analysis of layer 10 causal models in fine-tuned GPT2-small model on boolean logic tasks. Results compare the performance of individual causal models versus their combinations across different faithfulness thresholds. Combined models demonstrate stronger hypotheses at high faithfulness levels (0.9-1.0).

Theorems & Definitions (2)

• Definition 1: Combined Causal Models
• Definition 2: Model Strength