Feasibility and benefits of joint learning from MRI databases with different brain diseases and modalities for segmentation (2405.18511v1)

Abstract: Models for segmentation of brain lesions in multi-modal MRI are commonly trained for a specific pathology using a single database with a predefined set of MRI modalities, determined by a protocol for the specific disease. This work explores the following open questions: Is it feasible to train a model using multiple databases that contain varying sets of MRI modalities and annotations for different brain pathologies? Will this joint learning benefit performance on the sets of modalities and pathologies available during training? Will it enable analysis of new databases with different sets of modalities and pathologies? We develop and compare different methods and show that promising results can be achieved with appropriate, simple and practical alterations to the model and training framework. We experiment with 7 databases containing 5 types of brain pathologies and different sets of MRI modalities. Results demonstrate, for the first time, that joint training on multi-modal MRI databases with different brain pathologies and sets of modalities is feasible and offers practical benefits. It enables a single model to segment pathologies encountered during training in diverse sets of modalities, while facilitating segmentation of new types of pathologies such as via follow-up fine-tuning. The insights this study provides into the potential and limitations of this paradigm should prove useful for guiding future advances in the direction. Code and pretrained models: https://github.com/WenTXuL/MultiUnet

• The paper shows that joint learning using Multi-Unet with modality dropout significantly improves segmentation robustness across diverse MRI datasets.
• It introduces innovative architectures like LFUnet and MAFUnet that effectively fuse multiple imaging modalities to handle missing inputs.
• The study proves that pre-training on mixed databases followed by fine-tuning on specific datasets enhances model transferability to unseen pathologies.

Feasibility and Benefits of Joint Learning from MRI Databases with Different Brain Diseases and Modalities for Segmentation

The paper discusses novel methodologies for the segmentation of brain lesions in MRI images by leveraging multiple databases that encompass a range of brain diseases and imaging modalities. Traditional segmentation models have commonly been trained on singular databases tailored for specific pathologies using predefined imaging protocols. This paper challenges the conventional approach by interrogating whether training with diverse databases and heterogeneous modalities can improve segmentation performance through joint learning strategies.

Research Questions

The paper posits several key questions aimed at exploring the potential of joint learning:

1. Can a model be trained effectively on multiple MRI databases with different modalities and pathologies?
2. Will databases with varying sets of modalities mutually benefit when used in a joint training scheme?
3. Can such a training paradigm enable the segmentation of pathologies seen during training but with different modalities at test time?
4. How would a jointly trained model perform when encountering pathologies not presented during its initial training, either directly or through subsequent fine-tuning?

Methodologies

Three primary methods are proposed and evaluated:

1. Multi-Unet:
   • A variant of the established Unet model, Multi-Unet, is designed to handle various input modalities by adapting the input layer channels to cover all unique modalities present across the training databases.
   • A modality dropout technique is introduced to enhance generalization, where random modalities are omitted during training to encourage the model to rely on any possible combination of modalities.
2. LFUnet (Late-Fusion Unet):
   • LFUnet incorporates a separate Unet for each modality, fusing the output embeddings through weighted averaging with attention mechanisms, aiming to handle missing modalities more effectively.
3. MAFUnet (Multi-scale Attention Fusion Unet):
   • MAFUnet employs separate encoders for each modality and fusions at multiple scales within the architecture.

Experimental Settings

Seven MRI databases were utilized, covering a range of brain pathologies including tumors, multiple sclerosis, traumatic brain injury, and stroke lesions. The databases contain different combinations of modalities such as T1, T2, FLAIR, and DWI.

Results

Multiple experimental results are presented to delineate the capabilities and advantages of the proposed models:

1. Robustness to Missing Modalities:
   • Performance degradation was assessed when testing without certain modalities. Training with modality dropout proved essential to preserving performance. Multi-Unet demonstrated enhanced robustness over more complex variants such as LFUnet and MAFUnet.
2. Joint Training on Multiple Databases:
   • Multi-Unet jointly trained on all databases provided comparable or superior performance to models trained individually on single databases. This showcases the feasibility of the joint training paradigm.
3. Generalization to New Databases:
   • Models jointly trained using Multi-Unet were able to segment pathologies in new databases that included different sets of modalities. For example, stroke segmentation in the ISLES dataset was improved by leveraging additional modalities seen during the joint training.
4. Segmenting Unseen Pathologies:
   • The efficacy of Multi-Unet in directly segmenting pathologies not seen during training was tested against state-of-the-art unsupervised models. While Multi-Unet outperformed these unsupervised methods, fine-tuning remains necessary to achieve closer to state-of-the-art results on target databases.
5. Fine-tuning Benefits:
   • The paper highlights the gains from pre-training on multiple databases followed by fine-tuning on specific new databases. This approach consistently outperformed training models from scratch, demonstrating the transferability of learned knowledge.

Implications and Future Work

The paper brings forward several implications:

• Practical Benefits: A single jointly trained model can simplify operational workflows by eliminating the need for multiple models tailored for specific pathologies and modality sets.
• Generalization and Transfer Learning: Improved generalization to unseen data configurations and enhanced fine-tuning capabilities signify practical advantages in diverse clinical applications.
• Simplicity vs. Complexity: Simple architectures like Multi-Unet with modality dropout can be highly effective, potentially limiting the need for more complex solutions.

Future work may focus on:

• Advanced Frameworks: Investigating advanced architectures that can better handle heterogeneous training data without significant performance trade-offs.
• Unsupervised Learning: Integrating self-supervised or unsupervised learning techniques to mitigate the need for comprehensive labeled datasets.

In conclusion, this paper not only substantiates the feasibility of joint learning from diverse MRI databases but also underscores its practical and theoretical benefits, thereby paving the way for further innovations in medical image segmentation.