Latent Causal Modeling for 3D Brain MRI Counterfactuals

The number of samples in structural brain MRI studies is often too small to properly train deep learning models. Generative models show promise in addressing this issue by effectively learning the data distribution and generating high-fidelity MRI. However, they struggle to produce diverse, high-quality data outside the distribution defined by the training data. One way to address this issue is to use causal models developed for 3D volume counterfactuals. However, accurately modeling causality in high-dimensional spaces is challenging, so these models generally generate 3D brain MRIs of lower quality. To address these challenges, we propose a two-stage method that constructs a Structural Causal Model (SCM) within the latent space. In the first stage, we employ a VQ-VAE to learn a compact embedding of the MRI volume. Subsequently, we integrate our causal model into this latent space and execute a three-step counterfactual procedure using a closed-form Generalized Linear Model (GLM). Our experiments conducted on real-world high-resolution MRI data (1 mm) provided by the Alzheimer’s Disease Neuroimaging Initiative (ADNI) and the National Consortium on Alcohol and Neurodevelopment in Adolescence (NCANDA) demonstrate that our method can generate high-quality 3D MRI counterfactuals.

💡 Research Summary

The paper tackles two intertwined challenges in neuroimaging: the scarcity of large MRI datasets for training deep generative models, and the difficulty of producing realistic out‑of‑distribution samples that reflect causal manipulations such as aging or disease progression. To address these issues, the authors propose a two‑stage pipeline that embeds high‑resolution 3D T1‑weighted brain MRIs into a compact latent space using a Vector‑Quantized Variational Auto‑Encoder (VQ‑VAE) and then builds a Structural Causal Model (SCM) directly in that latent space.

In the first stage, the VQ‑VAE learns an encoder‑decoder pair together with a codebook. An MRI volume (x) is mapped by the encoder (E) to a continuous latent tensor (z), which is then quantized by nearest‑neighbor lookup in the codebook, yielding a discrete representation (Q(z)). This quantization serves both as a regularizer and as a means to dramatically reduce dimensionality while preserving essential anatomical detail.

The second stage introduces a Latent SCM (LSCM). Clinical and demographic variables—age, sex, diagnosis (alcohol use disorder vs. control), and region‑of‑interest (ROI) volumes for frontal, insular, and parietal cortices—are treated as endogenous nodes together with the latent MRI features (z). A causal graph, inspired by prior neurobiological findings, encodes directed edges such as “age → ventricular enlargement” and “diagnosis → frontal atrophy”.

Causal inference follows Pearl’s three‑step ladder (abduction, action, prediction) but is implemented with a Generalized Linear Model (GLM) that admits a closed‑form solution. All latent vectors from the training set are flattened into a matrix (Z) (size (N \times K)). The parent variables for each sample form a matrix (P) ((N \times m)). Solving the normal equations yields the coefficient matrix (B = (P^{\top}P)^{-1}P^{\top}Z). The residual (U_Z = Z - PB) acts as the exogenous noise for each latent feature.

During the “action” phase, a user‑specified intervention (e.g., set age to 50 years) is applied by modifying the parent values in the graph, producing counterfactual parent vectors (c_pa_z). The “prediction” step recombines the unchanged exogenous component with the new parents: (\tilde{Z} = U_Z + c_pa_z B). The resulting counterfactual latent tensor (\tilde{z}) is fed to the frozen VQ‑VAE decoder, yielding a high‑fidelity 3D MRI that reflects the imposed causal change.

The authors evaluate the method on two large public cohorts: ADNI (380 subjects, ages 59‑91) and NCANDA (808 subjects, ages 12‑27). After standard preprocessing (denoising, bias correction, skull stripping, affine registration, intensity normalization to