[Paper Review] High Fidelity Image Counterfactuals with Probabilistic Causal Models

Fabio De Sousa Ribeiro, Tian Xia+3 more|arXiv (Cornell University)|Jun 27, 2023

Machine Learning in Healthcare4 citations

TL;DR

This paper proposes a deep generative causal framework using probabilistic structural causal models to generate high-fidelity image counterfactuals, leveraging causal mediation analysis and variational inference. It enables accurate estimation of direct, indirect, and total effects in image data, demonstrating axiomatic soundness through effectiveness and composition properties in medical imaging benchmarks.

ABSTRACT

We present a general causal generative modelling framework for accurate estimation of high fidelity image counterfactuals with deep structural causal models. Estimation of interventional and counterfactual queries for high-dimensional structured variables, such as images, remains a challenging task. We leverage ideas from causal mediation analysis and advances in generative modelling to design new deep causal mechanisms for structured variables in causal models. Our experiments demonstrate that our proposed mechanisms are capable of accurate abduction and estimation of direct, indirect and total effects as measured by axiomatic soundness of counterfactuals.

Motivation & Objective

• To address the challenge of estimating accurate, high-fidelity counterfactual images for high-dimensional structured data like medical scans.
• To enable estimation of direct, indirect, and total causal effects on image variables using causal mediation analysis within a deep generative framework.
• To ensure counterfactuals satisfy axiomatic properties—effectiveness and composition—thereby validating their causal soundness.
• To support applications in explainability, data augmentation, fairness, and robustness in safety-critical domains such as medical imaging.
• To bridge the gap in deep learning by fulfilling all three rungs of Pearl’s ladder of causation: association, intervention, and counterfactual reasoning.

Proposed method

• The framework employs a Markovian probabilistic causal model with endogenous variables (e.g., images) and exogenous noise variables, structured via directed acyclic graphs (DAGs).
• It introduces deep causal mechanisms based on variational autoencoders (HVAEs) for image generation, trained end-to-end with structural assignments.
• Parent predictors are trained in the anticausal direction (from children to parents) using MLPs or CNNs to infer latent attributes from image observations.
• Counterfactual generation is achieved via constrained optimization using a Lagrangian formulation to enforce intervention constraints while preserving image fidelity.
• The method uses exponential moving averages of model parameters for inference and applies data augmentation (random crops, flips) during training.
• A key innovation is the use of separate, independently trained anticausal predictors for counterfactual training and evaluation to prevent shortcut learning.

Experimental results

Research questions

• RQ1 Can deep generative models produce high-fidelity image counterfactuals that satisfy axiomatic properties of causality?
• RQ2 How accurately can direct, indirect, and total causal effects be estimated on image-level variables using causal mediation analysis?
• RQ3 To what extent does the proposed framework improve counterfactual effectiveness and composition compared to prior methods?
• RQ4 Can the model generate plausible counterfactuals in data-scarce medical imaging settings with minimal spurious correlations?
• RQ5 Does the use of anticausal parent predictors mitigate shortcut learning during counterfactual training?

Key findings

• The proposed framework achieves high-fidelity counterfactual generation on brain MRI and chest X-ray datasets, with qualitative results indicating plausible image edits.
• On UK Biobank data, the sex predictor achieved ROCAUC of 0.9764 ± 2e-3, and the MRI predictor achieved perfect ROCAUC of 1.0, indicating strong parent attribute inference.
• The age predictor on UK Biobank achieved a mean absolute error (MAE) of 4.847 ± 7e-4 years, demonstrating precise regression performance.
• The brain volume and ventricle volume predictors achieved MAEs of 26.77 ± 0.39 ml and 1.958 ± 3e-2 ml, respectively, showing high accuracy in volumetric estimation.
• On MIMIC-CXR, the sex, race, and disease predictors achieved ROCAUCs of 0.9950, 0.7496, and 0.9419, respectively, indicating strong performance on diverse clinical attributes.
• Counterfactual training converged in approximately 10,000 iterations on brain MRI and 6,500 on chest X-ray, with stable optimization using the Lagrangian formulation.