Causal Abstraction: A Theoretical Foundation for Mechanistic Interpretability

Abstract: Causal abstraction provides a theoretical foundation for mechanistic interpretability, the field concerned with providing intelligible algorithms that are faithful simplifications of the known, but opaque low-level details of black box AI models. Our contributions are (1) generalizing the theory of causal abstraction from mechanism replacement (i.e., hard and soft interventions) to arbitrary mechanism transformation (i.e., functionals from old mechanisms to new mechanisms), (2) providing a flexible, yet precise formalization for the core concepts of polysemantic neurons, the linear representation hypothesis, modular features, and graded faithfulness, and (3) unifying a variety of mechanistic interpretability methods in the common language of causal abstraction, namely, activation and path patching, causal mediation analysis, causal scrubbing, causal tracing, circuit analysis, concept erasure, sparse autoencoders, differential binary masking, distributed alignment search, and steering.

• The paper introduces a novel framework that uses interventions to map low-level neural network behaviors to high-level causal models.
• It generalizes causal abstraction to cyclic structures and typed variables by decomposing models through marginalization, variable merge, and value merge.
• It unifies various XAI methods by providing a quantifiable metric, interchange intervention accuracy, to assess approximate causal abstractions.

This paper, "Causal Abstraction for Faithful Model Interpretation" (2301.04709), argues that the theory of causal abstraction offers a mathematical foundation for creating faithful and interpretable explanations of AI model behavior and internal structure. The core idea is that a high-level, human-intelligible causal model can be considered a faithful description of a complex, low-level AI model (like a neural network) if the high-level variables can be aligned with sets of low-level variables that play the same causal role. The authors use interventions on model-internal states to rigorously assess this alignment.

The main contributions are:

1. Generalization of Causal Abstraction: The paper extends existing notions of causal abstraction to accommodate cyclic causal structures (where causal influences can form loops) and typed high-level variables (where variables can be categorized, e.g., as Booleans or integers).
2. Interchange Interventions for Analysis: It details how multi-source interchange interventions can be used to conduct causal abstraction analyses. An interchange intervention involves running a model with a "base" input but setting some internal model states (e.g., activations in a neural network layer) to the values they would have taken if a different "source" input had been provided. This allows testing the causal role of specific internal representations.
3. Approximate Causal Abstraction: A notion of approximate causal abstraction is defined, allowing for a graded assessment of how well a high-level model describes a low-level one. This is crucial because perfect abstractions are rare in practice. This leads to a quantifiable metric called "interchange intervention accuracy."
4. Decomposition of Constructive Abstraction: The paper proves that constructive causal abstraction can be broken down into three fundamental operations:
   • Marginalization: Removing variables from a model.
   • Variable-merge: Grouping multiple low-level variables into a single high-level variable.
   • Value-merge: Grouping multiple values of a variable into a single abstract value.
5. Unifying XAI Methods: It formalizes several existing XAI methods—LIME, causal effect estimation, causal mediation analysis, iterated nullspace projection, and circuit-based explanations—as special cases of causal abstraction analysis. It also shows how integrated gradients can compute interchange interventions.

Key Concepts and Implementation

Causal Models

The paper defines a deterministic causal model $\mathcal{M}$ as a pair $(\mathbf{V}, \mathcal{F})$, where $\mathbf{V}$ is a set of variables and $\mathcal{F}$ is a set of structural functions $\{f_V\}_{V \in \mathbf{V}}$. Each $$f_V: \mathsf{Val}(\mathbf{V}) \rightarrow \mathsf{Val}(V)$$ assigns a value to variable $V$ based on the values of all other variables (though typically only depending on a subset, its "parents").

• Intervention: An intervention $do(\mathbf{I}=\mathbf{i})$ sets a subset of variables $\mathbf{I} \subseteq \mathbf{V}$ to specific values $\mathbf{i}$, replacing their original structural functions $f_I$ with constant functions.
• Solution Set: $Solve(\mathcal{M}_\mathbf{i})$ is the set of total settings $\mathbf{v}$ (assignments of values to all variables in $\mathbf{V}$) that satisfy all structural equations after an intervention $\mathbf{i}$. For acyclic models, this solution is unique.

Causal Abstraction

Given a low-level model $\mathcal{L}$ (e.g., a neural network) and a high-level model $\mathcal{H}$ (e.g., a symbolic algorithm), an alignment $\langle \Pi, \tau \rangle$ is defined.

• $$\Pi = \{\Pi_{X_{\mathcal{H}}}\}_{X_{\mathcal{H}} \in \mathbf{V}_{\mathcal{H}} \cup \{\bot\}}$$ is a partition of the low-level variables $\mathbf{V}_{\mathcal{L}}$. Each cell $\Pi_{X_{\mathcal{H}}}$ maps to a high-level variable $X_{\mathcal{H}}$, and $\Pi_{\bot}$ contains low-level variables not mapped to any high-level variable (marginalized out).
• $$\tau = \{\tau_{X_{\mathcal{H}}}\}_{X_{\mathcal{H}} \in \mathbf{V}_{\mathcal{H}}}$$ is a family of partial surjective maps. Each $$\tau_{X_{\mathcal{H}}}: \mathsf{Val}(\Pi_{X_\mathcal{H}}) \rightarrow \mathsf{Val}(X_{\mathcal{H}})$$ maps a setting of a low-level variable cluster to a value of the corresponding high-level variable.

This alignment induces a translation function $$\tau: \mathsf{Val}(\mathbf{V}_{\mathcal{L}}) \rightarrow \mathsf{Val}(\mathbf{V}_{\mathcal{H}})$$.

Causal Consistency: An alignment is causally consistent if for any valid low-level intervention $\mathbf{i}$ (one that has a corresponding high-level intervention $\tau(\mathbf{i})$):

$$\tau(Solve(\mathcal{L}_{\mathbf{i}})) = Solve(\mathcal{H}_{\tau(\mathbf{i})})$$

This means that intervening on the low-level model and then abstracting the result yields the same high-level state as abstracting the intervention and then applying it to the high-level model.

Low-Level: i_L ---(Solve L_i)--> v_L | | tau tau | | High-Level: i_H ---(Solve H_i)--> v_H

The diagram must commute.

Constructive Abstraction: $\mathcal{H}$ is a constructive abstraction of $\mathcal{L}$ under $\langle \Pi, \tau \rangle$ if causal consistency holds.

Interchange Intervention Analysis

This is a practical method to test for causal abstraction.

• Alignment Construction (Definition \ref{def:interchangealign}): For analyzing a neural network $\mathcal{N}$ with a high-level algorithm $\mathcal{A}$:
  1. Partition the low-level intermediate variables of $\mathcal{N}$ into cells $\Pi_{X_\mathcal{H}}$, each corresponding to an intermediate variable $X_\mathcal{H}$ in $\mathcal{A}$. Input and output variables are typically aligned by design.
  2. Define $\tau$ for inputs and outputs based on the task.
  3. For an intermediate high-level variable $X_\mathcal{H}$ and a low-level state $$\mathbf{z}_\mathcal{L} \in \mathsf{Val}(\Pi_{X_\mathcal{H}})$$, if $\mathbf{z}_\mathcal{L}$ is realized by $\mathcal{N}$ for some input $\mathbf{x}_\mathcal{L}$ (i.e., $$\mathbf{z}_\mathcal{L} = \mathsf{Proj}(Solve(\mathcal{N}_{\mathbf{x}_\mathcal{L}}), \Pi_{X_\mathcal{H}})$$), then define $$\tau_{X_\mathcal{H}}(\mathbf{z}_\mathcal{L}) = \mathsf{Proj}(Solve(\mathcal{A}_{\tau(\mathbf{x}_\mathcal{L})}), X_\mathcal{H})$$. Otherwise, $\tau_{X_\mathcal{H}}(\mathbf{z}_\mathcal{L})$ is undefined.

• Interchange Intervention (Definition \ref{def:interchange}):

  Given a neural network $\mathcal{N}$, a base input $\mathbf{b}$, source inputs $\mathbf{s}_1, \dots, \mathbf{s}_k$, and disjoint sets of intermediate low-level variables $$\mathbf{X}_{\mathcal{L}}^1, \dots, \mathbf{X}_{\mathcal{L}}^k$$: The intervention sets the input variables to $\mathbf{b}$, and each set $\mathbf{X}_{\mathcal{L}}^j$ to the value it would take if source input $\mathbf{s}_j$ were provided to $\mathcal{N}$.

  $$\mathsf{IntInv}(\mathcal{N}, \mathbf{b}, \langle \mathbf{s}_1, \dots ,\mathbf{s}_k \rangle, \langle \mathbf{X}_{\mathcal{L}}^1, \dots, \mathbf{X}_{\mathcal{L}}^k\rangle) = \mathbf{b} \cup \bigcup_{j=1}^k \mathsf{Proj}(Solve(\mathcal{N}_{\mathbf{s}_j}), \mathbf{X}_{\mathcal{L}}^j)$$

To perform the analysis, one applies such an intervention to $\mathcal{N}$ and the corresponding abstracted intervention to $\mathcal{A}$, then checks if their outputs (or subsequent states) match according to $\tau$.

Example: Hierarchical Equality Task

The paper uses a running example of a "hierarchical equality" task: input is two pairs of objects; output is True if both pairs have the same equality status (both equal or both unequal), False otherwise.

• Low-Level Model ($\mathcal{N}$): A fully-connected feed-forward neural network trained on this task. Its variables are neuron activations, and structural functions are neural computations (e.g., matrix multiply + ReLU). Inputs are vector embeddings of shapes.
• High-Level Model ($\mathcal{A}$): A symbolic, tree-structured algorithm:
  1. $V_1 = \text{equals}(X_1, X_2)$
  2. $V_2 = \text{equals}(X_3, X_4)$
  3. $O = \text{equals}(V_1, V_2)$ Variables $X_i$ are input shapes, $V_i$ are intermediate Boolean equality results, $O$ is the final Boolean output.

An alignment $\langle \Pi, \tau \rangle$ is proposed:

• Input neurons in $\mathcal{N}$ encoding the $k$-th shape are mapped to $X_k$ in $\mathcal{A}$.

• Specific hidden layer neurons in $\mathcal{N}$ (e.g., $H_{(2,2)}, H_{(2,3)}$) are mapped to $V_1$ in $\mathcal{A}$.
• Output logits in $\mathcal{N}$ are mapped to $O$ in $\mathcal{A}$.
• $\tau$ maps neuron activation patterns to symbolic values (e.g., specific activation vectors in $H_{(2,2)}, H_{(2,3)}$ map to True/False for $V_1$).

The paper demonstrates that for this specific network (trained using interchange intervention training), $\mathcal{A}$ is a constructive abstraction of $\mathcal{N}$. This means performing an interchange intervention on $\mathcal{N}$ (e.g., running with input $(\pentagon, \pentagon, \bigtriangleup, \square)$ but forcing the internal state corresponding to $V_1$ to be what it would be for input $(\square, \pentagon, \bigtriangleup, \bigtriangleup)$) and then abstracting the output gives the same result as abstracting the intervention and applying it to $\mathcal{A}$.

Decomposition of Constructive Abstraction (Theorem \ref{theorem:constructive})

A high-level model $\mathcal{H}$ is a constructive abstraction of a low-level model $\mathcal{L}$ if and only if $\mathcal{H}$ can be derived from $\mathcal{L}$ by:

1. Marginalization (Definition \ref{def:marginalize}): Removing a subset of variables $\mathbf{X}$ from $\mathcal{L}$. The structural functions of remaining variables are adjusted to reflect this removal, essentially integrating out the effect of $\mathbf{X}$.
2. Variable Merge (Definition \ref{def:variable-merge}): Partitioning $\mathbf{V}_\mathcal{L}$ into cells $\{\Pi_X\}_{X \in \mathbf{W}}$. Each cell $\Pi_X$ becomes a new variable $X$ in $\mathbf{W}$, whose value space is the Cartesian product of the value spaces of variables in $\Pi_X$. Structural functions are composed accordingly.
3. Value Merge (Definition \ref{def:value-merge}): For each variable $X$, a function $\delta_X: \mathsf{Val}(X) \rightarrow B_X$ maps original values to new, potentially coarser-grained values. This is valid only if collapsed values play the same causal role (i.e., if $\delta_X(x) = \delta_X(x')$, then substituting $x$ for $x'$ in any intervention context leads to outputs that are also equivalent under $\delta$).

This decomposition provides a constructive way to understand how a simpler model can emerge from a more complex one. For the hierarchical equality example:

1. Marginalize: Neurons in $\mathcal{N}$ not part of the aligned clusters (i.e., in $\Pi_\bot$) are removed.
2. Variable Merge: Clusters of neurons in $\mathcal{N}$ (e.g., $\{R_1, R_2\}$ for input $X_1$, or $\{H_{(2,2)}, H_{(2,3)}\}$ for intermediate $V_1$) are merged into single variables corresponding to $\mathcal{A}$'s variables.
3. Value Merge: Continuous activation values of merged neuron clusters in $\mathcal{N}$ are mapped to discrete symbolic values (e.g., $\{\pentagon, \bigtriangleup, \square\}$ for inputs, $\{True, False\}$ for intermediates/output) of $\mathcal{A}$.

Approximate Abstraction (Section \ref{section:approximate})

Since perfect abstraction is rare, the paper defines $\alpha$-on-average constructive abstraction.

• Requires a distance metric $\mathsf{Distance}_{\mathcal{H}}$ between high-level total settings.
• An alignment is $\alpha$-on-average causally consistent if the expected distance (over a uniform distribution of interventions $\mathbf{i} \in \mathsf{Domain}(\tau)$) between $\tau(Solve(\mathcal{L}_{\mathbf{i}}))$ and $Solve(\mathcal{H}_{\tau(\mathbf{i})})$ is less than or equal to $\alpha$.

  $$\mathbb{E}_{\mathbf{i} \sim \mathsf{Uniform}(\mathsf{Domain}(\tau))}[\mathsf{Distance}_{\mathcal{H}}( \tau(Solve(\mathcal{L}_{\mathbf{i}})), Solve(\mathcal{H}_{\tau(\mathbf{i})}) )] \leq \alpha$$

• Interchange Intervention Accuracy (IIA): The proportion of interchange interventions for which the output of $\mathcal{N}$ (after abstraction by $\tau$) matches the output of $\mathcal{A}$.

  $\mathsf{IIA}(\mathcal{N}, \mathcal{A}, \tau) = \mathbb{E}_{\mathbf{i} \sim \mathsf{Uniform}(Domain(\tau))}[ \mathbbm{1}[\tau(\mathsf{Proj}(\mathcal{N}_{\mathbf{i}}, \mathbf{X}_{\mathcal{L}}^{\text{Out}})) = \mathsf{Proj}(\mathcal{A}_{\tau(\mathbf{i})} , \mathbf{X}_{\mathcal{H}}^{\text{Out}})] ]$

  If $\mathsf{Distance}_\mathcal{H}$ is a 0-1 loss (0 if equal, 1 if not), then $\mathsf{IIA} = 1-\alpha$. Thus, IIA is directly related to $\alpha$-on-average constructive abstraction.

Practically, IIA can be estimated by sampling a large number of base and source inputs, performing the corresponding interchange interventions, and checking output consistency.

Practical Implications and Connections to XAI

• LIME: Interpreted as an approximate abstraction where both the black-box model $\mathcal{N}$ and the interpretable LIME model $\mathcal{A}$ are simplified to two-variable (Input $\rightarrow$ Output) causal chains. LIME fidelity measures the $\alpha$ in an $\alpha$-on-average abstraction over a local neighborhood of inputs.
• Causal Effect Estimation (e.g., CEBaB): Estimating the effect of a real-world concept (e.g., food quality $C_{\text{food}}$) on model output $X_{\text{Out}}$ can be seen as marginalizing a larger causal graph (including data generation) to a two-variable chain $C_{\text{food}} \rightarrow X_{\text{Out}}$.
• Causal Mediation Analysis: Analyzing how an input $X$'s effect on output $Y$ is mediated by an intermediate $Z$ (e.g., a set of neurons) corresponds to a three-variable chain $X \rightarrow Z \rightarrow Y$. Complete mediation occurs if there's no direct $X \rightarrow Y$ link after abstraction. Partial mediation relates to approximate abstraction.
• Iterative Nullspace Projection (INP): Removing information about a concept $C$ from a hidden representation $\mathbf{H}$ by projecting $\mathbf{H}$ onto the nullspace of probes for $C$. This can be framed as an abstraction to a three-variable model $$(X_{\text{In}} \rightarrow L \rightarrow X_{\text{Out}})$$, where $L$ is a binary variable indicating if the information removal intervention occurred. The abstraction holds if the network's behavior matches the expected degraded performance when $L=0$ (intervention applied).
• Circuit-Based Explanations: Hypotheses about neurons representing concepts and connections implementing algorithms can be formalized as an alignment between a neural network and a high-level algorithmic causal model. For instance, a neuron that activates for "dogs or cars" (if they don't co-occur in training) could be a high-level variable in an abstract model.
• Integrated Gradients (IG): The completeness axiom of IG can be used to compute the outcome of an interchange intervention:

  $$\mathsf{Proj}(\mathcal{N}_{\mathbf{b} \cup \mathsf{IntInv}(\mathbf{b}, \langle \mathbf{s} \rangle, \langle \mathbf{Y}\rangle)}, \mathbf{X}^{\text{Out}}) = \mathsf{Proj}(\mathcal{N}_{\mathbf{b}}, \mathbf{X}^{\text{Out}}) - \sum_{i=1}^{|\mathbf{Y}|} \mathsf{IG}_i(\mathsf{Proj}(\mathcal{N}_{\mathbf{b}}, \mathbf{Y}), \mathsf{Proj}(Solve(\mathcal{N}_{\mathbf{s}}), \mathbf{Y}))$$

  where the baseline for IG is set to the "source" activation $$\mathsf{Proj}(Solve(\mathcal{N}_{\mathbf{s}}), \mathbf{Y})$$, and the input to IG is the "base" activation $$\mathsf{Proj}(\mathcal{N}_{\mathbf{b}}, \mathbf{Y})$$.

Future Applications and Extensions

The framework is extended to:

• Typed Variables: High-level variables can have types (e.g., Boolean, integer), and type consistency can be enforced in the abstraction. This was crucial in prior work by the authors for vision-based generalization tasks.
• Infinite Variables and Cycles: The paper sketches an example of modeling a bubble sort algorithm with a countably infinite number of variables (to handle sequences of arbitrary length and arbitrary sorting iterations) and how this can be abstracted to cyclic models representing equilibrium states.
• Probabilistic Models: The paper concludes by discussing how causal abstraction can be extended to probabilistic causal models. This involves aligning distributions (observational, interventional, and importantly, counterfactual) rather than deterministic states. A constructive probabilistic abstraction requires that the counterfactual distributions align, which is a stronger condition than just aligning interventional distributions.

Overall, the paper provides a rigorous, intervention-based framework for mechanistic interpretability, aiming to bridge the gap between human-understandable concepts and the low-level workings of complex AI models. Its strength lies in grounding interpretability in causality and providing tools (like interchange interventions and their accuracy metric) for empirical validation. The decomposition theorem offers a conceptual toolkit for understanding how abstractions are formed.