3D Brain MRI Generation with a Clinically-Conditioned VAE-GAN and Diffusion-Driven Feature Sampling
Abstract
We introduce a 3D VAE-GAN framework that synthesizes brain MRI volumes conditioned on seven clinical attributes, such as Alzheimer’s disease (AD) diagnostic labels and key volumetric measures, including the hippocampus, amygdala, and lateral ventricle, which are known to correlate with AD. Leveraging a 3D encoder-decoder with depthwise-separable convolutions and a style-based modulation, our model efficiently captures critical biomarkers and injects clinical information directly into the generation process. During the training, two pre trained auxiliary heads, Alzheimer’s Disease and Cognitively Normal (AD/CN) classification and brain volume vector regression, provide additional cross entropy and regression losses, ensuring that generated scans remain anatomically plausible and clinically consistent. To sample realistic clinical vectors during inference, we additionally train a diffusion model on clinical vectors, enabling flexible sampling of disease states without the need for manual feature engineering. Experimental results demonstrate high-quality 3D MRI generations. Additionally, adjusting disease labels or specific brain volumes demonstrates a feasible level of conditional control, suggesting that this approach could benefit data augmentation and support clinically relevant neuroimaging tasks.