Feasibility of Federated Learning from Client Databases with Different Brain Diseases and MRI Modalities

TL;DR

Evaluation on 7 brain MRI databases with 5 different diseases shows that this FL framework can train a single model that achieves very promising results in segmenting all disease types seen during training, and can segment these diseases in new databases that contain sets of modalities different from those in training clients.

Abstract

Segmentation models for brain lesions in MRI are typically developed for a specific disease and trained on data with a predefined set of MRI modalities. Such models cannot segment the disease using data with a different set of MRI modalities, nor can they segment other types of diseases. Moreover, this training paradigm prevents a model from using the advantages of learning from heterogeneous databases that may contain scans and segmentation labels for different brain pathologies and diverse sets of MRI modalities. Additionally, the confidentiality of patient data often prevents central data aggregation, necessitating a decentralized approach. Is it feasible to use Federated Learning (FL) to train a single model on client databases that contain scans and labels of different brain pathologies and diverse sets of MRI modalities? We demonstrate promising results by combining appropriate, simple, and practical modifications to the model and training strategy: Designing a model with input channels that cover the whole set of modalities available across clients, training with random modality drop, and exploring the effects of feature normalization methods. Evaluation on 7 brain MRI databases with 5 different diseases shows that this FL framework can train a single model achieving very promising results in segmenting all disease types seen during training. Importantly, it can segment these diseases in new databases that contain sets of modalities different from those in training clients. These results demonstrate, for the first time, the feasibility and effectiveness of using FL to train a single 3D segmentation model on decentralised data with diverse brain diseases and MRI modalities, a necessary step towards leveraging heterogeneous real-world databases. Code: https://github.com/FelixWag/FedUniBrain

Figures (5)

• Figure 1: FedUniBrain (yellow box) enables the collaborative training of a single brain MRI segmentation model capable of segmenting multiple brain pathologies across multiple institutions without data sharing. FedUniBrain can train a model on data from institutions with varying sets of input MRI modalities and different brain diseases. In contrast, standard brain MRI segmentation models (grey box), can only process MRI scans with a specific set of MRI input modalities and are disease-specific.
• Figure 2: Method overview. (1) Initialization of model with distinct set of modalities across clients. (2) Federated training of the model with random modality drop. (3) Generalization to unseen clients with different input MRI modality combination.
• Figure 3: Example segmentation results for databases with different pathologies. (a) Segmentation outputs for five different databases, each having a different pathology to segment and different modality combination. Shown are 2D slices of the 3D input, the ground truth, predictions from individually trained U-Nets (on each database separately), and our FedUniBrain (trained on ATLAS, BRATS, MSSEG, TBI and WMH) model predictions. For each database we show a different modality. (b) Segmentation outputs from a FedUniBrain model trained on ATLAS, BRATS, MSSEG, TBI, and WMH, tested on excluded ISLES and Tumor2 databases with unseen modality combinations. Also, results from disease-specific models are shown: one trained on the BRATS (Tumor) and evaluated on Tumor2, and another trained on ATLAS and evaluated on ISLES for stroke lesions.
• Figure 4: Impact of a new client joining the federation. Performance (Dice score %) for both the initially trained clients and a client that joins at a later stage of training is presented. Blue bars represent the results before the new client joins (if the new client has an unseen modality, it is excluded during evaluation for the blue bar results). (a) Model initially trained on ATLAS, BRATS, MSSEG, TBI, and WMH (blue), then two scenarios are presented: after integrating ISLES client and learning its data (orange), and after integrating Tumor2 client (green). (b) Initial model trained without TBI (blue) and after integration and learning the TBI client (orange). (c) Initially trained without MSSEG (blue) and after integrating MSSEG (orange). Single-center training results on the new client are shown for comparison (violet).
• Figure C5: MRI Databases with their different modalities.