A highly generalized federated learning algorithm for brain tumor segmentation.

Area of Science:

• Medical Imaging
• Artificial Intelligence in Healthcare
• Machine Learning

Background:

• Brain image segmentation is crucial for diagnosis and treatment planning.
• Federated Learning (FL) allows collaborative AI model training while protecting patient data.
• Data imbalance and scarcity in medical datasets hinder AI model performance.

Purpose of the Study:

• To propose a client-side brain tumor image segmentation model using Virtual Adversarial Training (VAT) within a 3D U-Net.
• To address data scarcity and imbalance in federated learning environments for medical imaging.
• To enhance the generalizability of federated models through an effective aggregation strategy.

Main Methods:

• Integration of Virtual Adversarial Training (VAT) into a 3D U-Net architecture for client-side segmentation.
• Development of FedHG, a federated model aggregation strategy using public validation dataset weights.
• Incorporation of instance normalization parameters into client models during training.

Main Results:

• The proposed algorithm achieved a 2.2% improvement in Dice Similarity Coefficient (DSC) for brain tumor segmentation compared to baseline FL.
• The model demonstrated a marginal 3% performance reduction relative to centralized training.
• The method effectively improved segmentation accuracy and model generalizability under data scarcity and imbalance.

Conclusions:

• The proposed federated learning approach with VAT and FedHG effectively enhances brain tumor image segmentation accuracy and generalizability.
• This method offers a practical solution for AI in healthcare, particularly in scenarios with limited or imbalanced medical data.
• The study highlights the potential of federated learning combined with advanced training techniques for robust medical AI applications.