Integrating Anatomical Priors into a Causal Diffusion Model

2025-09-10

AI Summary

Existing counterfactual brain MRI generation models lack anatomical inductive bias and struggle to capture disease-associated subtle morphological changes. Method: We propose the first framework integrating probabilistic causal graphs with diffusion models, featuring a voxel-wise anatomical constraint mechanism: a 3D ControlNet encodes anatomical masks and embeds them into a counterfactual denoising UNet, coupled with a 3D diffusion decoder. This explicitly unifies causal intervention modeling with structural priors to enable controllable, interpretable synthesis of localized regional alterations. Contribution/Results: Evaluated on multiple public datasets, our method significantly outperforms baselines in both image fidelity (FID, LPIPS) and consistency with downstream structural measurements (cortical thickness, volume), achieving state-of-the-art performance. It successfully reproduces sub-millimeter pathological differences reported in the literature for Alzheimer’s disease and schizophrenia.

Technology Category

• Machine Learning: Causal Learning
• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Causality

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

3D brain MRI studies often examine subtle morphometric differences between cohorts that are hard to detect visually. Given the high cost of MRI acquisition, these studies could greatly benefit from image syntheses, particularly counterfactual image generation, as seen in other domains, such as computer vision. However, counterfactual models struggle to produce anatomically plausible MRIs due to the lack of explicit inductive biases to preserve fine-grained anatomical details. This shortcoming arises from the training of the models aiming to optimize for the overall appearance of the images (e.g., via cross-entropy) rather than preserving subtle, yet medically relevant, local variations across subjects. To preserve subtle variations, we propose to explicitly integrate anatomical constraints on a voxel-level as prior into a generative diffusion framework. Called Probabilistic Causal Graph Model (PCGM), the approach captures anatomical constraints via a probabilistic graph module and translates those constraints into spatial binary masks of regions where subtle variations occur. The masks (encoded by a 3D extension of ControlNet) constrain a novel counterfactual denoising UNet, whose encodings are then transferred into high-quality brain MRIs via our 3D diffusion decoder. Extensive experiments on multiple datasets demonstrate that PCGM generates structural brain MRIs of higher quality than several baseline approaches. Furthermore, we show for the first time that brain measurements extracted from counterfactuals (generated by PCGM) replicate the subtle effects of a disease on cortical brain regions previously reported in the neuroscience literature. This achievement is an important milestone in the use of synthetic MRIs in studies investigating subtle morphological differences.

Problem

• Generating anatomically plausible counterfactual brain MRIs
• Preserving fine-grained anatomical details in synthetic images
• Replicating subtle disease effects on cortical brain regions

Innovation

• Integrates anatomical constraints into diffusion model
• Uses probabilistic graph for voxel-level anatomical priors
• Employs ControlNet-constrained counterfactual denoising UNet