[Paper Review] Multiscale Metamorphic VAE for 3D Brain MRI Synthesis

Jaivardhan Kapoor, Jakob H. Macke+1 more|arXiv (Cornell University)|Jan 9, 2023

Generative Adversarial Networks and Image Synthesis4 citations

TL;DR

This paper proposes a multiscale metamorphic variational autoencoder (M³AE) for 3D brain MRI synthesis that generates high-fidelity MRI volumes by applying composable, multiscale diffeomorphic and intensity transformations to a fixed anatomical template. The method achieves state-of-the-art FID scores while preserving a disentangled, interpretable latent space, significantly outperforming prior VAE and GAN-based approaches in distribution coverage and anatomical fidelity.

ABSTRACT

Generative modeling of 3D brain MRIs presents difficulties in achieving high visual fidelity while ensuring sufficient coverage of the data distribution. In this work, we propose to address this challenge with composable, multiscale morphological transformations in a variational autoencoder (VAE) framework. These transformations are applied to a chosen reference brain image to generate MRI volumes, equipping the model with strong anatomical inductive biases. We structure the VAE latent space in a way such that the model covers the data distribution sufficiently well. We show substantial performance improvements in FID while retaining comparable, or superior, reconstruction quality compared to prior work based on VAEs and generative adversarial networks (GANs).

Motivation & Objective

• To address the challenge of generating high-fidelity 3D brain MRIs with sufficient coverage of the data distribution.
• To incorporate strong anatomical inductive biases into a variational autoencoder framework for improved generative modeling.
• To enable a disentangled, manipulable latent space for downstream analysis while maintaining reconstruction quality.
• To extend morphological transformation-based synthesis from 2D to 3D MRI volumes with multiscale composition.
• To overcome limitations of GANs (e.g., mode collapse, lack of latent space) and VAEs (e.g., blurry outputs) in 3D brain MRI generation.

Proposed method

• The model uses a variational autoencoder with a fixed reference brain template as the base image for generation.
• It decomposes the generative process into two decoders: one for diffeomorphic deformation fields (ϕ) and one for additive intensity transformations (A).
• Transformations are applied in a coarse-to-fine, multiscale cascade, with each level doubling in resolution, enabling hierarchical refinement.
• The final output is formed by composing transformations as M(T|ϕ,A) = (T + A) ∘ ϕ, with the last stage being purely intensity-based.
• The loss function combines ELBO with scale-specific regularization terms for deformation fields and intensity maps, including smoothness and divergence constraints.
• The model is trained using AdamW with a learning rate of 0.0003, and inference uses a latent vector sampled from a standard normal distribution.

Figure 1 : Description of the proposed M 3 AE model. The decoder consists of two separate backbones Decoder ϕ and Decoder A that output diffeomorphic deformation fields and additive intensity transformations, respectively. These transformations are composed from coarse-to-fine scales that increase b

Experimental results

Research questions

• RQ1 Can multiscale, composable morphological transformations improve 3D brain MRI synthesis fidelity and distribution coverage compared to standard VAEs and GANs?
• RQ2 Does using a fixed anatomical template with deformation and intensity transformations enable better anatomical consistency in generated MRI volumes?
• RQ3 Can a disentangled latent space be preserved in a VAE-based 3D MRI generator while achieving high-fidelity outputs?
• RQ4 How does the proposed method compare to SOTA baselines in terms of FID, SSIM, PSNR, and MSE metrics?
• RQ5 What are the limitations in topological diversity and visual quality introduced by multiscale deformation fields?

Key findings

• The M³AE model achieves a FID score of 11.2 on 2D slices, substantially outperforming the α WGAN baseline and β VAE, indicating superior distribution coverage and sample quality.
• The model maintains reconstruction quality comparable to β VAE, with SSIM, PSNR, and MSE values within a narrow margin, despite the anatomical constraints.
• Qualitative analysis shows that M³AE samples are the most anatomically correct, with cortical folds and structural features well-preserved, unlike artifacts seen in β VAE and α WGAN.
• The model exhibits occasional wavy visual artifacts due to poorly generated fine-scale deformation fields, indicating a need for improved regularization at high resolutions.
• The method demonstrates limited topological diversity in cortical folds, suggesting a need for enhanced modeling of sulcal and gyral variability in future work.
• The disentangled latent space enables meaningful downstream analysis, unlike GAN-based methods that lack such a structured latent representation.

Figure 2 : Preprocessed volumes from the ADNI dataset