Efficient Discovery of Approximate Causal Abstractions via Neural Mechanism Sparsification
| Authors | Amir Asiaee |
| Year | 2026 |
| Field | Machine Learning |
| arXiv | 2602.24266 |
| Download | |
| Categories | cs.LG, cs.AI |
Abstract
Neural networks are hypothesized to implement interpretable causal mechanisms, yet verifying this requires finding a causal abstraction -- a simpler, high-level Structural Causal Model (SCM) faithful to the network under interventions. Discovering such abstractions is hard: it typically demands brute-force interchange interventions or retraining. We reframe the problem by viewing structured pruning as a search over approximate abstractions. Treating a trained network as a deterministic SCM, we derive an Interventional Risk objective whose second-order expansion yields closed-form criteria for replacing units with constants or folding them into neighbors. Under uniform curvature, our score reduces to activation variance, recovering variance-based pruning as a special case while clarifying when it fails. The resulting procedure efficiently extracts sparse, intervention-faithful abstractions from pretrained networks, which we validate via interchange interventions.
Engineering Breakdown
Plain English
This paper tackles the problem of discovering simpler, interpretable causal models hidden within trained neural networks. The authors reframe neural network pruning as a search for 'causal abstractions'—simplified structural causal models that preserve the network's behavior under interventions. Instead of expensive brute-force methods, they derive an Interventional Risk objective with closed-form pruning criteria based on second-order analysis. Their approach shows that activation variance emerges naturally from their theory, providing theoretical grounding for why variance-based pruning works and offering a principled way to find which network units can be replaced with constants or folded into neighbors.
Core Technical Contribution
The key innovation is reframing structured pruning as an approximate causal abstraction discovery problem, moving from engineering heuristics to principled causal theory. The authors derive an Interventional Risk objective that captures when a pruned network remains faithful to the original network's causal behavior under interventions. They show that second-order Taylor expansion of this objective yields closed-form sparsification criteria—explicit formulas for deciding which units to prune without retraining. This bridges causal inference and neural network efficiency, recovering variance-based pruning as a special case while providing theoretical justification and extensions beyond it.
How It Works
The method treats a trained neural network as a deterministic Structural Causal Model (SCM) with latent variables representing hidden units. The authors formulate an Interventional Risk objective that measures how well a candidate abstraction (pruned network) matches the original network's causal behavior when you intervene on variables. They expand this objective using second-order Taylor approximation around the full network, deriving tractable criteria for deciding: (1) whether a unit can be replaced with a constant, or (2) whether a unit can be folded into its neighbors without harming causal fidelity. Under simplifying assumptions like uniform curvature, these criteria reduce to measuring activation variance across training data. The algorithm iteratively applies these criteria to identify pruning candidates, producing a smaller network that preserves causal relationships.
Production Impact
For engineers deploying large models, this approach could significantly reduce model size and inference latency by providing theoretically justified pruning guidance without expensive retraining or intervention experiments. Rather than applying generic pruning heuristics, practitioners can compute activation variance statistics on a calibration set and use closed-form formulas to identify which units are safe to remove, making the process faster and more interpretable. The method scales better than brute-force causal discovery approaches, which require many interventions on the original network. However, the production trade-off is that the method assumes networks encode causal structure—if your network is purely interpolating patterns without meaningful causal mechanisms, the abstraction may not be faithful. Integration would require storing causal graphs alongside models and validating that pruned variants still satisfy your downstream task requirements.
Limitations and When Not to Use This
The approach assumes neural networks actually implement meaningful causal mechanisms, which is unproven and may not hold for all architectures or tasks—particularly in purely empirical benchmarks like standard image classification. The theory relies on assumptions like uniform curvature and deterministic networks, simplifications that may break down in practice with ReLU networks exhibiting varying curvature and stochasticity from dropout or batch normalization. The method works on already-trained networks and doesn't address the orthogonal problem of training sparse models from scratch, potentially missing better lottery ticket solutions. Evaluating whether discovered abstractions are 'correct' requires expensive intervention experiments on the original model, making it hard to validate in production without ground truth causal graphs—the paper doesn't provide practical validation metrics for practitioners.
Research Context
This work bridges causal inference and neural network interpretability, building on the framework of causal abstractions from prior work while introducing tractable approximation methods. It relates to the 'lottery ticket hypothesis' literature on neural network sparsity but grounds the search in causal theory rather than pure performance metrics. The paper extends variance-based pruning heuristics used in model compression by deriving them from principled causal objectives, providing theoretical justification for why these methods work. It opens research directions in verifying whether trained networks encode meaningful causal structure and in designing training procedures that encourage causal mechanisms—potentially impacting both interpretability and efficiency research communities.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::