MACAW: A Causal Generative Model for Medical Imaging

TL;DR

MACAW introduces a causal generative framework that embeds a structural causal graph into a single invertible normalizing flow via masked causal autoencoders (C-MADE). The model supports interventional sampling, counterfactual inference through abduction-action-prediction, and Bayesian classification, demonstrated on synthetic data and 23,692 UK Biobank brain MRI slices projected to 60 latent components. Results show accurate encoding of causal structure, realistic counterfactual generation, and meaningful brain aging–related changes, with age prediction evaluated using counterfactuals and Bayesian posterior ages. limitations include dependency on a predefined causal graph, scalability to 3D data, and higher computational cost for class-conditioned inference, pointing to future work in end-to-end dimensionality reduction and 3D extensions for broader clinical applicability.

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 a priori causal knowledge into deep learning models is a potential avenue to identify these problems. 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. The code for a toy dataset is available at the following link: https://github.com/vibujithan/macaw-2D.git.

Figures (8)

• Figure 1: A causal DAG (left) and its respective C-MADE network (center). The MACAW architecture (right) consists of multiple C-MADE networks connected in a series, thereby forming a normalizing flow.
• Figure 2: (a) The sum of the absolute difference between each variable in the test set after the counterfactual query for the action $do(x_2=2)$, (b) The expected values of the predicted counterfactual values and respective ground truth values.
• Figure 3: (a) Predefined causal graph for the UKBB dataset, incorporating age, sex, and BMI values. (b) Corresponding adjacency matrix, where filled cells indicate 1 and unfilled cells indicate 0.
• Figure 4: Results from two types of generative sampling: unconditional (top) and interventional (bottom). For interventional sampling, we set the age variable to 75 years.
• Figure 5: Illustration of the outcomes of counterfactual queries, where causal factors of randomly selected raw images were altered to generate new counterfactual images. Every column represents a distinct counterfactual query, identified by "$do(.)$" commands.