On the Mechanistic Interpretability of Neural Networks for Causality in Bio-statistics

TL;DR

This paper explores how mechanistic interpretability techniques can be applied to neural networks used in causal inference within bio-statistics, aiming to enhance understanding, validation, and comparison of models in health-related research.

Contribution

It demonstrates methods to probe neural network representations, visualize computational pathways, and compare mechanisms across models for improved causal analysis in bio-statistics.

Findings

MI tools can validate internal representations of NNs.

Visualization reveals how NNs process confounders and treatments.

Comparison methods highlight differences between statistical, ML, and NN models.

Abstract

Interpretable insights from predictive models remain critical in bio-statistics, particularly when assessing causality, where classical statistical and machine learning methods often provide inherent clarity. While Neural Networks (NNs) offer powerful capabilities for modeling complex biological data, their traditional "black-box" nature presents challenges for validation and trust in high-stakes health applications. Recent advances in Mechanistic Interpretability (MI) aim to decipher the internal computations learned by these networks. This work investigates the application of MI techniques to NNs within the context of causal inference for bio-statistics. We demonstrate that MI tools can be leveraged to: (1) probe and validate the internal representations learned by NNs, such as those estimating nuisance functions in frameworks like Targeted Minimum Loss-based Estimation (TMLE); (2) discover and visualize the distinct computational pathways employed by the network to process different types of inputs, potentially revealing how confounders and treatments are handled; and (3) provide methodologies for comparing the learned mechanisms and extracted insights across statistical, machine learning, and NN models, fostering a deeper understanding of their respective strengths and weaknesses for causal bio-statistical analysis.