MACAW: A Causal Generative Model for Medical Imaging

Vibujithan Vigneshwaran, Erik Yuiti Ohara, Matthias Wilms, Nils Forkert

Abstract: Although deep learning techniques show promising results for many neuroimaging tasks in research settings, they have not yet found widespread use in clinical scenarios. One of the reasons for this problem is that many machine learning models only identify correlations between the input images and the outputs of interest, which can lead to many practical problems, such as encoding of uninformative biases and reduced explainability. Thus, recent research is exploring if integrating \textit{a priori} causal knowledge into deep learning models is a potential avenue to identify these problems. However, encoding causal reasoning and generating genuine counterfactuals necessitates computationally expensive invertible processes, thus restricting analyses to a small number of causal variables and rendering them infeasible for generating even 2D images. To overcome these limitations, this work introduces a new causal generative architecture named Masked Causal Flow (MACAW) for neuroimaging applications. Within this context, three main contributions are described. First, a novel approach that integrates complex causal structures into normalizing flows is proposed. Second, counterfactual prediction is performed to identify the changes in effect variables associated with a cause variable. Finally, an explicit Bayesian inference for classification is derived and implemented, providing an inherent uncertainty estimation. The feasibility of the proposed method was first evaluated using synthetic data and then using MRI brain data from more than 23000 participants of the UK biobank study. The evaluation results show that the proposed method can (1) accurately encode causal reasoning and generate counterfactuals highlighting the structural changes in the brain known to be associated with aging, (2) accurately predict a subject's age from a single 2D MRI slice, and (3) generate new samples assuming other values for subject-specific indicators such as age, sex, and body mass index.

Submission Length: Long submission (more than 12 pages of main content)

Changes Since Last Submission:

1. A new section (Sec. 4.1.5) to the synthetic data experiments that demonstrates the benefits of explicit causal modeling.
2. New benchmarking results have now been added to the additional experiments section in the appendix (A.2), and the MorphoMNIST section has been expanded accordingly.
3. We removed the previous ad-hoc baseline comparison with HVAE.
4. Minor typo fixes

Code: https://github.com/vibujithan/macaw-2D

Assigned Action Editor: ~Krikamol_Muandet1

Submission Number: 5831