Source-linked AI summary
Finding Alignments Between Interpretable Causal Variables and Distributed Neural Representations
Atticus Geiger, Zhengxuan Wu, Christopher Potts, Thomas Icard, Noah D. Goodman
TL;DR
Existing causal abstraction analyses rely on brute-force alignment search and localist neuron assignments, limiting how neural implementations can be identified. DAS learns alignments by gradient descent in distributed, non-standard representational bases, and experiments find perfect alignments while exposing different underlying substructures. The method therefore supports causal-abstraction analyses of distributed neural representations, subject to remaining discrete hyperparameters and alignment assumptions.
Problem
Existing causal abstraction methods require computationally intensive brute-force alignment search and presuppose disjoint neuron groups for high-level variables.
Method
DAS learns high-level-to-low-level alignments with gradient descent and distributed interchange interventions in non-standard representational bases.
Results
DAS discovers perfect alignments for distributed representations of equality and lexical-entailment relations, while lexical-entailment representations can decompose into word-identity sub-representations.
Takeaways & Limitations
DAS reveals conceptual and causal substructure in trained neural networks that localist or brute-force alignment analyses may miss.
Takeaways & Limitations
DAS still uses discrete hyperparameters, and causal-abstraction conclusions can depend on alignment choice and the assumed representational basis.
Abstract
from arXiv · showhide
Causal abstraction is a promising theoretical framework for explainable artificial intelligence that defines when an interpretable high-level causal model is a faithful simplification of a low-level deep learning system. However, existing causal abstraction methods have two major limitations: they require a brute-force search over alignments between the high-level model and the low-level one, and they presuppose that variables in the high-level model will align with disjoint sets of neurons in the low-level one. In this paper, we present distributed alignment search (DAS), which overcomes these limitations. In DAS, we find the alignment between high-level and low-level models using gradient descent rather than conducting a brute-force search, and we allow individual neurons to play multiple distinct roles by analyzing representations in non-standard bases-distributed representations. Our experiments show that DAS can discover internal structure that prior approaches miss. Overall, DAS removes previous obstacles to conducting causal abstraction analyses and allows us to find conceptual structure in trained neural nets.
1. Introduction
Causal abstraction asks whether interpretable symbolic algorithms faithfully explain neural networks, but prior methods struggle with alignment search and distributed representations. DAS addresses these obstacles with gradient-based alignment and non-standard representational bases, revealing structure in equality and entailment tasks.
- Motivation: Causal abstraction formalizes when a complex causal system, such as a neural network, implements a simpler symbolic algorithm.Interchange interventions compare aligned high-level and low-level variables by forcing neural states to match source-input values.
- Limitations: Previous methods face brute-force alignment search and assume high-level variables map to disjoint neuron groups.Exhaustive search can miss optimal alignments, while the localist assumption excludes overlapping conceptual roles for individual neurons.
- DAS: DAS learns alignments with gradient descent and analyzes neural representations in alternative bases where interpretable dimensions need not correspond to individual neurons.Distributed interchange interventions rotate representations, intervene on selected dimensions, and rotate them back to the neuron-aligned basis.
- Experiments: DAS finds perfect alignment for a binary lexical-entailment variable in a fine-tuned language model.The natural-language-inference task uses premise and hypothesis sentences differing in one word pair, with entailment versus contradiction or neutrality as labels.
- Findings: DAS also reveals that apparent lexical-entailment structure may decompose into word-identity representations, whereas hierarchical equality relations remain abstracted from their constituent entities.This contrast indicates that the same perfect abstraction score can conceal different causal substructures.
2. Related Work
Causal abstraction provides a formal account of how high-level causal variables summarize low-level mechanisms under interventions. Related work uses this framework to analyze faithful and interpretable AI explanations, while interchange-intervention training supplies an objective for learning alignments.
- Causal abstraction: Causal abstraction characterizes when a high-level causal model is an abstract description of a low-level causal model.High-level variables may correspond to overlapping low-level variables that summarize causal mechanisms under hard or soft interventions.
- Applications: The framework has unified intervention-based explanation methods, including LIME, causal effect estimation, mediation analysis, and nullspace projection.These methods are presented as related applications of causal abstraction to faithful and interpretable AI explanations.
- Training objectives: Interchange intervention training objectives are minimized when a high-level causal model abstracts a neural network under a specified alignment.This objective is used in the paper to learn alignments between interpretable causal variables and deep-learning representations.
3. Methods
The paper defines causal abstraction through aligned interventions and introduces distributed interventions and DAS to find faithful alignments beyond standard neuron-based representations.
- Constructive Causal Abstraction: Causal abstraction requires a high-level model to faithfully render the causal mechanisms of aligned low-level variables.
- Constructive Causal Abstraction: Interchange interventions replace target variables in a base-input setting with values obtained from corresponding source-input settings, producing counterfactual outputs.
- Constructive Causal Abstraction: A network can solve a task perfectly yet fail causal abstraction when its low-level and high-level interchange interventions have different effects under an alignment.
- Distributed Interventions: Rotating a hidden representation can reveal an alignment with perfect behavioral and counterfactual accuracy when the relevant structure is not localized in the standard basis.
- Distributed Interventions: Distributed interventions transform representations, intervene on orthogonal sub-spaces, and transform them back, allowing interpretable structure to be distributed across neurons.
- Distributed Alignment Search: DAS optimizes orthogonal alignment matrices with stochastic gradient descent against interchange-intervention objectives while keeping the high- and low-level models frozen.
- Approximate Causal Abstraction: IIA measures the proportion of aligned interchange interventions with equivalent high-level and low-level effects, coinciding with approximate causal abstraction when below 100%.
4. Hierarchical Equality Experiment
The hierarchical equality experiment tests whether a feed-forward network represents the two equality relations as intermediate variables. DAS compares alternative causal models and alignment strategies, finding stronger support for distributed equality representations than localist alternatives.
- Task and model: The task predicts whether two object pairs have the same equality relation, returning True when both pairs are equal or both unequal.Examples include (A,A,B,B) as True and (A,B,C,C) as False.
- Task and model: A three-layer ReLU feed-forward network is trained on randomly initialized object vectors and evaluated on held-out random vectors.The input vectors are in R^n, while biases and weights use dimensions based on 4n.
- High-level models: DAS evaluates a causal model in which left and right equality relations are computed separately before predicting the final label.Alternative models represent only the first equality relation or the identity of the first input object.
- Results: DAS learns higher-IIA alignments than brute-force localist searches, reaching perfect or near-perfect alignment where the best brute-force result is 0.60 and the best localist result is 0.73.The strongest results occur for the model representing both equality relations.
- Results: The left-equality representation cannot be decomposed into the identity of the first argument, supporting an abstract equality representation rather than stored input identities.The identity-subspace alternative has very low performance.
5. Monotonicity NLI Experiment
The Monotonicity NLI experiment tests whether fine-tuned BERT represents negation and lexical entailment as interpretable intermediate variables. DAS finds a perfect distributed alignment, while substructure analysis shows that lexical entailment can be decomposed into word identities.
- Task and model: MoNLI pairs premise and hypothesis sentences differing in one word, with labels determined by lexical entailment and reversed when negation is present.The changed hypothesis word is a hypernym or hyponym of the premise word.
- Task and model: The experiment fine-tunes uncased BERT-base and analyzes distributed representations of the [CLS] token because the final classification uses that representation.The model has 12 layers, 12 heads, and hidden dimension 768.
- High-level models: DAS tests a high-level model with two Boolean variables: whether negation is present and whether the premise word entails the hypothesis word.Alternative models represent only lexical entailment or the identity of the premise word.
- Results: 100% IIA is achieved for layer 9 with intervention size 256, aligning separate non-standard basis dimensions to lexical entailment and negation.The result indicates that 256 dimensions encode the word relation and another 256 encode negation.
- Results: The lexical-entailment representation is nearly perfectly decomposable into separate representations of the premise-word and hypothesis-word identities.Thus the apparent entailment representation functions as a data structure containing two word-identity representations.
6. Conclusion
The conclusion presents DAS as a gradient-descent method for aligning interpretable causal variables with distributed neural representations. Its experiments recover simple equality and lexical-entailment variables, while revealing different causal substructures within those representations.
- Conclusion: DAS learns distributed alignments with gradient descent, improving on brute-force search while allowing more interpretable alignments than localist methods.The method addresses both alignment-search and disjoint-neuron assumptions.
- Conclusion: The experiments discover perfect alignments between distributed representations and binary variables for equality and lexical entailment relations.The conclusion frames these as simple symbolic variables represented within trained neural networks.
- Conclusion: Lexical-entailment representations can decompose into word-identity subrepresentations, whereas equality representations cannot be decomposed into the identities of participating entities.The contrast motivates examining the causal substructure of neural representations.
A.1. Training Data for distributed alignment search (DAS)
DAS learns rotation matrices from counterfactual input-output examples generated by interventions on high-level variables. The training data differ by task-specific variables and replacements.
- Training data: Each rotation-matrix training example combines a base input, source inputs, high-level interventions, and corresponding counterfactual outputs.The counterfactual outputs are the network outputs under the high-level intervention effects.
- Training data: For hierarchical equality, source inputs provide equality relations or first-object changes that are interchanged with the base input.The sampled intervention depends on whether the high-level model represents both relations, left equality, or first-object identity.
- Training data: For Monotonicity NLI, source inputs provide negation or lexical-entailment values, while identity models interchange a replacement lexeme from another hypothesis sentence.Sampling avoids invalid entailment-label cases.
A.2. Reproducibility
The experiments train task models and alignment rotations using specified datasets, hyperparameters, repeated seeds, and baseline procedures.
- Hierarchical Equality Experiment: The hierarchical equality model uses 1.92M training pairs and reaches 100% training and testing accuracy.The model is trained for 10 epochs before reaching perfect training accuracy and evaluated on unseen input-output pairs.
- Hierarchical Equality Experiment: Rotation learning for each hierarchical equality high-level model uses 640K training pairs and 19.2K testing pairs.Both training phases use batch size 6.4K, a maximum of 10 epochs, learning rate 1e-3, and early stopping patience of 10K.
- Training Procedure: Each experiment is run three times with distinct random seeds on a single NVIDIA 2080 Ti RTX 11GB GPU.Hierarchical-equality training converges in less than ten minutes, and datasets are balanced across the two labels.
- Monotonicity NLI Experiment: The monotonicity NLI models are trained on 10K MoNLI examples and reach 100% training and testing accuracy.Models are fine-tuned for five epochs before reaching perfect task accuracy on the training and hold-out testing sets.
- Baselines: The brute-force baseline searches disjoint sliding-window neuron groups and selects the alignment with the highest interchange intervention accuracy.Its worst-case hypothesis count is estimated as Cn^m, where n is total representation dimension and m is variable dimension.
- Baselines: The localist baseline converts a learned orthogonal matrix into a localist alignment matrix before evaluating interchange intervention accuracy.Subspace alignment can also be learned sequentially by freezing an existing rotation and training another rotation on top of it.
Appendix B. Runtime Comparison: Brute-force Search Baseline vs. DAS
The runtime comparison evaluates DAS against brute-force search under matched settings, while rotation analyses assess whether learned alignments are distributed.
- Runtime Comparison: DAS runtime is invariant to the number of testing hypotheses, unlike brute-force search whose exhaustive computation can become intractable.The brute-force estimate is conservative because it considers only a limited subset of possible alignments.
- Runtime Comparison: DAS requires training, whereas the brute-force baseline searches candidate alignments without additional training.DAS runtime could be further reduced through early stopping or optimized training-data size.
- Learned Rotations: Most learned rotation eigenvectors are rotated across tasks, indicating that the discovered representations are not trivial.The results suggest high-level variables are represented in distributed, non-standard bases that direct activation probes may miss.
Appendix D. Common Questions
The common questions clarify rotation properties, approximate-abstraction interpretation, prerequisites, scalability, and practical uses of DAS.
- Learned Rotation Matrix: The learned full square rotation matrix is orthonormal because PyTorch’s parametrization guarantees orthonormality.Orthonormality is used to isolate rotation rather than scaling.
- Optimization Stability: Orthogonal-matrix optimization depends partly on initialization, so the experiments use distinct random seeds and select the best result.Different initialization schemes are identified as a possible future direction.
- Learned Rotation Matrix: Orthogonal transformations are not theoretically required, but nonlinear invertible transformations may better align variables represented on nonlinear sub-manifolds.The paper leaves such transformations for future work because they are computationally difficult to find.
- Prerequisites: DAS requires a partial or complete causal graph with interchangeable high-level variables and counterfactual data samples.The counterfactual samples provide base and source inputs differing in high-level-variable values.
- Interpreting IIA: When interchange intervention accuracy is below 100%, it is interpreted through α-on-average approximate causal abstraction and relative IIA rankings.Relative rankings indicate which high-level model is a better approximation of the low-level model.
- Practical Uses and Scaling: DAS transforms representations into an operatable state for interpretable interchange interventions, but full rotation matrices can become intractable for large models.A BERT-BASE full rotation is estimated at about 15.4B parameters; aligned-subspace or low-rank representations are proposed for tractability.