Efficient Discovery of Approximate Causal Abstractions via Neural Mechanism Sparsification

Paper: 2602.24266 Authors: Amir Asiaee Categories: cs.LG, cs.AI

The Gap

The mechanistic interpretability community wants to verify that neural networks implement causal mechanisms—that specific circuits compute specific functions in a way that holds up under intervention. The gold standard for this is finding a causal abstraction: a high-level Structural Causal Model (SCM) that faithfully represents the network’s behavior when you surgically modify its internals. But current methods hit a wall. Interchange intervention approaches (like those from Geiger et al.) require exhaustive brute-force search over all possible alignments between high-level variables and network components. Alignment-based methods demand retraining networks from scratch with alignment losses baked in. Both are computationally prohibitive for modern networks, and neither scales to the messy reality of pretrained models where you can’t just start over.

The Increment

Before: Discovering causal abstractions meant either brute-force intervention testing or retraining networks with alignment objectives. After: Pruning scores derived from interventional risk directly identify which neurons implement abstract causal variables—no search, no retraining.

The core insight is elegant: treat structured pruning not as compression, but as archaeological excavation of causal structure. Think of a trained network as a city built over ancient ruins. The pruning score is your ground-penetrating radar—it tells you which modern buildings (neurons) sit atop foundational structures (causal mechanisms) versus which are just decorative additions. The method works by deriving an Interventional Risk objective that measures how much a network’s behavior changes under hypothetical interventions. Taking the second-order Taylor expansion of this risk around the current network state yields closed-form scores for two operations: (1) replacing a unit with a constant (like demolishing a building and leaving empty ground), and (2) folding a unit into its neighbors (like merging adjacent structures). Units with low scores can be removed without disrupting the causal flow—they’re not load-bearing. High-scoring units are the pillars: remove them and the causal story collapses.

What’s particularly clever is that under uniform curvature assumptions, this interventional score reduces to simple activation variance—which explains why variance-based pruning sometimes works (when curvature is uniform) and sometimes catastrophically fails (when it’s not). The method outputs a sparse subnetwork that preserves intervention-faithful behavior, validated by actually performing interchange interventions on the pruned result.

Key Concepts

Causal Abstraction: Imagine you’re trying to understand how a car works, but you’re staring at the molecular structure of the metal. A causal abstraction is like having a schematic that shows “engine,” “transmission,” “wheels”—higher-level components with clear causal relationships. Formally, it’s a mapping from a complex low-level SCM (the neural network, where each neuron is a variable) to a simpler high-level SCM (the interpretable model, where variables represent concepts like “detects edges” or “encodes sentiment”). The key property: if you intervene on the high-level model (set “edge detector” to OFF), and then perform the corresponding intervention on the low-level model (ablate specific neurons), the outputs should match. This is intervention faithfulness—the abstraction isn’t just a loose analogy, it’s causally accurate.

Interventional Risk: Standard neural network training minimizes prediction error on observed data. But if you care about causal structure, you need to minimize error under *all possible interventions—not just what naturally occurs. Interventional Risk is the expected loss when you randomly intervene on the network’s internals, forcing neurons to take arbitrary values. Mathematically: $\mathcal{R}_{\text{int}} = \mathbb{E}_{x, \text{do}(z)} [L(f(x; \text{do}(z)), y)]$, where $\text{do}(z)$ represents surgical modifications to intermediate activations $z$. A network with low interventional risk has robust causal structure—its behavior degrades gracefully under perturbation because it’s organized around stable mechanisms, not brittle correlations.

Second-Order Pruning Criterion: When you remove a neuron, how much does interventional risk increase? Computing this exactly requires testing all possible interventions—intractable. Instead, approximate the risk change using a second-order Taylor expansion around the current network state. This gives you a closed-form score: $\Delta \mathcal{R} \approx \nabla \mathcal{R} \cdot \Delta \theta + \frac{1}{2} \Delta \theta^T H \Delta \theta$, where $H$ is the Hessian (curvature) of the risk. For pruning, $\Delta \theta$ represents setting weights to zero or merging units. The score decomposes into (1) first-order term: how much does the gradient “want” this change? and (2) second-order term: how curved is the loss landscape in this direction? Units in flat regions (low curvature, low variance) can be removed cheaply; units in steep valleys are structurally essential.

Expert Assessment

Problem significance: This tackles a real bottleneck in mechanistic interpretability. The community has compelling hypotheses about what circuits exist in transformers and CNNs, but verification is stuck in manual, small-scale analysis. If you could automatically extract validated causal abstractions from pretrained models, it would accelerate interpretability research by orders of magnitude. The affected community is currently small (mechanistic interpretability researchers, maybe a few hundred people), but the implications are broad—this is foundational work for AI safety and scientific understanding of deep learning.

Method maturity: This is proof-of-concept territory. The theoretical contribution is solid—the connection between pruning and causal abstraction is genuinely novel. But the method assumes you already know the high-level SCM structure you’re searching for, which is a significant limitation. In practice, you don’t know a priori that a network implements “edge detector → shape recognizer → object classifier”; discovering that structure is the hard part. The paper also leans heavily on uniform curvature assumptions to make computation tractable, but real networks have wildly non-uniform curvature. The authors acknowledge this implicitly by showing when variance-based pruning fails, but don’t provide a scalable solution for the general case.

Experimental rigor: The validation via interchange interventions is the right move—it’s the gold standard for causal abstraction. However, the experiments appear limited in scope (the paper abstract doesn’t specify datasets or network architectures tested). A red flag: no comparison to recent alignment-based methods like Distributed Alignment Search (DAS) or causal scrubbing variants. The claim that this is more efficient than brute-force search is likely true, but without head-to-head comparisons on the same tasks, it’s hard to assess the practical speedup. The reduction to variance-based pruning under uniform curvature is intellectually satisfying but raises the question: how often does uniform curvature actually hold in practice?

Verdict: weak accept — Solid theoretical contribution connecting two previously separate areas (pruning and causal abstraction), but experimental validation needs to be more comprehensive and the method’s practical applicability remains unclear without knowing the target abstraction structure upfront.

Takeaways

Reframe compression as structure discovery: Pruning is usually framed as “make the model smaller.” This paper shows that with the right objective (interventional risk instead of prediction loss), pruning becomes a tool for uncovering latent structure. This framing is portable—anywhere you’re compressing a complex system, ask whether the compression criterion could be designed to reveal rather than just reduce.

Second-order approximations buy you closed-form solutions: When exact computation is intractable (testing all interventions), Taylor expansion around the current state often gives you actionable scores. The trick is choosing what to expand around. Here, expanding interventional risk around the trained network turns an exponential search into a one-pass scoring procedure. This pattern—approximate an intractable objective via local curvature—shows up everywhere from neural architecture search to Bayesian optimization.

Variance-based heuristics have hidden assumptions: Practitioners often use activation variance or gradient magnitude as pruning scores because they’re cheap to compute. This paper reveals the implicit assumption: uniform curvature in the loss landscape. When that holds, variance is a perfect proxy for interventional importance. When it doesn’t, variance-based pruning can catastrophically remove causally critical units. The lesson: if you’re using a simple heuristic, understand what geometric or distributional assumption makes it valid, and test whether that assumption holds in your domain.