FLWGAN: Federated Learning with Wasserstein Generative Adversarial Network for Brain Tumor Segmentation

Peketi Divya, Vishnu Chalavadi, C. Krishna Mohan, Yen-Wei Chen

Abstract: Recently, the potential of deep learning in identifying complex patterns is gaining research interest in medical applications specifically for brain tumor diagnosis. To segment tumors accurately in brain MRIs, there is a need for a large amount of data for training deep learning models. Also, hospitals cannot share patient data for centralization on the server since health records are prone to privacy and ownership challenges. To deal with these challenges, we set up an efficient federated learning (FL) pipeline with Wasserstein generative adversarial networks (FLWGAN) to ensure data privacy and data sufficiency. FL preserves the data privacy of clients by sharing only the trained model parameters to a centralized server instead of raw data. A modified 3D Wasserstein generative adversarial network with gradient penalty (WGAN-GP) and is incorporated at the client side to generate image-segmentation pairs for efficient training segmentation models. Here, 3D-UNet with an attention module is used for the brain MRI segmentation. The attention module is integrated into a 3D-UNet encoder network for effective brain tumor segmentation. Our approach aims to allow each client to benefit from locally available real data and synthetic data. This process enhances the learning performance while respecting data privacy. The efficacy of our proposed pipeline is demonstrated on the brain tumor task of the medical segmentation decathlon (MSD) dataset. We designed FLWGAN frameworks for predicting four segmentation tasks, i.e., whole tumor (WT), enhanced tumor (ET), tumor core (TC), and multiclass. Our proposed approach achieves state of the art performance in terms of various segmentation metrics.