Start of funding 01.07.2023

BRAIN-Gen: Brain Representations Learning with Generative Models

Prof. Dr. Nassir Navab
Technische Universität München
Fakultät für Informatik

Prof. Dr. Ehsan Adeli
Stanford University
Department of Psychiatry and Behavioral Science



Our project aims to develop a generative model to unravel the intricate interactions of counterfactual factors — such as age, sex, and education — and their influence on brain structure, function, and cognition. Utilizing vast brain datasets, we will develop a Disentangling Autoencoder (DAE) model, leveraging recent advancements in generative and disentanglement techniques. The DAE will reconstruct and analyze brain imaging datasets, extracting latent representations of complex factors. This effort will enhance our understanding of these factors' impacts on brain health, potentially guiding the creation of targeted interventions and therapies for neurological conditions.

Final report:
Introduction

This report highlights collaborative research between the Computer-Aided Medical Procedures (CAMP) lab at the Technical University of Munich (TUM) and the Stanford Translational AI Research (STAI) lab at Stanford University, focusing on cutting-edge developments in medical image analysis. The collaboration has yielded significant contributions in areas such as generative modeling for multimodal MRI synthesis and the application of latent diffusion models for counterfactual image generation. This report synthesizes key findings and methodologies from the publications [1,2,3] resulting from the joint efforts of these research groups.

Physics-Informed Latent Diffusion for Multimodal MRI Synthesis

In [1], we addressed the challenge of limited and variable modality availability across different datasets, which hinders the development of universally applicable generative models. We developed a physics-informed generative model to synthesize a variable number of brain MRI modalities, including those absent in the training data. Our model employs a two-step generative process, leveraging latent diffusion models (LDMs):

* Physical Tissue Property Maps: We utilize LDMs to generate unobserved physical tissue property maps: proton density (PD), longitudinal relaxation time (T1), and transverse relaxation time (T2). To combine information from different modalities, we use product-of-experts (PoE) to aggregate unimodal inference distributions into a joint multimodal distribution. We also condition the encoder on acquisition parameters using adaptive group normalization (AdaGN).

* Physical Signal Model: We then apply a physical signal model to the generated tissue property maps to synthesize the final MRI scans. The signal model consists of equations that approximate the relationship between tissue properties (PD, T1, T2) and signal intensity, depending on acquisition parameters like echo time (TE), repetition time (TR), and inversion time (TI).

This physics-informed approach allows us to generate more realistic and physically plausible MRI images. Our experiments show that this model can generate MRI contrasts not present in the training data and that the generated tissue properties have realistic distributions.

Hybrid Convolutional-Transformer Architecture for Segmentation

We also explored hybrid architectures for medical image segmentation to combine the benefits of convolutional neural networks (CNNs) and transformer networks. This research [2] stems from the observation that self-supervised pre-trained vision transformers, such as DINO, learn attention maps that capture valuable spatial information. To harness this, we developed DAINet (Dual Attention-Image Network), a specialized CNN architecture. DAINet incorporates attention map visualizations from a pre-trained DINO model as an additional input stream alongside the original image. The architecture employs a dual-encoder design: (1) the first encoder processes the original image, and (2) The second encoder processes the attention map visualizations. The outputs of these encoders are concatenated and fed into a shared decoder. We also introduced a switching mechanism to control the skip connections between the encoders and the decoder. DAINet outperforms both purely CNN-based and transformer-based architectures in medical image segmentation tasks, achieving state-of-the-art performance with reduced computational cost.

Latent Drifting in Diffusion Models for Counterfactual Medical Image Synthesis

We also addressed the challenge of generating counterfactual medical images, which are crucial for exploring hypothetical scenarios in medical contexts. We developed Latent Drifting (LD) [3], a novel method to condition diffusion models for this task. LD introduces a hyperparameter σ into the diffusion process to adapt a pre-trained model's learned distribution to a new, target medical image distribution. This is important because fine-tuning diffusion models on limited medical datasets is often difficult due to the distribution shift between general image datasets (used for pre-training) and specialized medical image datasets. Diffusion models are trained to denoise data iteratively, and LD modifies this process by adding the σ hyperparameter to the target of the forward process and the reverse diffusion process. We define LD as a min-max optimization problem, where we aim to minimize the distance between the generated data distribution and the target data distribution while maximizing the counterfactual fidelity. The counterfactual fidelity is measured using a loss function that ensures similarity to the original image while also achieving the desired counterfactual outcome. The efficacy of LD is demonstrated through counterfactual image generation tasks, showing its ability to generate realistic and condition-faithful images.

Key Findings and Implications

Our key findings in this collaboration are summarized below:

* [1] demonstrates that physics-informed generative models can synthesize realistic and diverse MRI contrasts, even those not present in the training data, by disentangling tissue properties and acquisition parameters.

* In [2], we found that incorporating attention map visualizations from pre-trained transformers enhances CNN-based segmentation models, achieving state-of-the-art performance with reduced computational cost.

* [3] shows that Latent Drifting effectively conditions diffusion models for counterfactual medical image generation, improving realism and fidelity compared to standard fine-tuning methods.

Conclusion

Our collaborative research has yielded substantial contributions to medical image analysis, pushing the boundaries of what's possible with deep learning. By combining our expertise, we have developed innovative solutions for MRI synthesis, counterfactual image generation, and image segmentation. These advancements hold significant promise for improving diagnostic accuracy, treatment planning, and our overall understanding of medical images.

References

[1] S. Lüpke, Y. Yeganeh, E. Adeli, N. Navab, and A. Farshad. “Physics-Informed Latent Dif- fusion for Multimodal Brain MRI Synthesis”. In: Medical Image Computing and Computer Assisted Intervention – MICCAI 2024 Workshops. Cham: Springer Nature Switzerland, 2025, pp. 198–207.

[2] Y. Yeganeh, I. Charisiadis, M. Hasny, M. Hartenberger, B. Ommer, N. Navab, A. Farshad, and E. Adeli. “Latent Drifting in Diffusion Models for Counterfactual Medical Image Syn- thesis”. In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (2025).

[3] Y. Yeganeh, A. Farshad, P. Weinberger, S.-A. Ahmadi, E. Adeli, and N. Navab. “Transformers pay attention to convolutions leveraging emerging properties of ViTs by dual attention-image network”. In: Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023, pp. 2304–2315.