Latent Causal Modeling for 3D Brain MRI Counterfactuals

TL;DR

This paper introduces a novel two-stage latent space causal modeling approach for generating high-quality 3D brain MRI counterfactuals, addressing limitations of existing generative models in diversity and quality.

Contribution

It develops a structural causal model within a learned latent space using VQ-VAE and a closed-form GLM, enabling realistic counterfactual MRI generation.

Findings

Generated high-quality 3D MRI counterfactuals from real data

Demonstrated effectiveness on ADNI and NCANDA datasets

Improved diversity and realism over traditional generative models

Abstract

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.