Optimizing Federated Learning Configurations for MRI Prostate Segmentation and Cancer Detection: A Simulation Study.

Area of Science:

• Artificial Intelligence in Radiology
• Medical Imaging Analysis
• Machine Learning in Healthcare

Background:

• Federated learning (FL) offers a privacy-preserving approach for developing AI models across multiple institutions without sharing raw patient data.
• Prostate cancer detection and segmentation using biparametric MRI are crucial for accurate diagnosis and treatment planning.
• Developing robust AI models for these tasks requires diverse datasets and optimized training strategies.

Purpose of the Study:

• To develop and optimize a federated learning (FL) framework for enhanced biparametric MRI prostate segmentation and clinically significant prostate cancer (csPCa) detection.
• To evaluate the performance of the optimized FL model against local client models and a baseline FL model.
• To determine the optimal FL configurations (epochs, rounds, aggregation strategies) for both segmentation and detection tasks.

Main Methods:

• A retrospective study utilized the Flower FL framework to train a nnU-net-based architecture on biparametric MRI data from January 2010 to August 2021.
• Model development involved optimizing local epochs, federated rounds, and aggregation strategies (FedMedian, FedAdagrad) for prostate segmentation (4 clients, 1294 patients) and csPCa detection (3 clients, 1440 patients).
• Performance was assessed using Dice scores for segmentation and the Prostate Imaging: Cancer Artificial Intelligence (PI-CAI) score for csPCa detection on independent test sets, with statistical significance determined by permutation testing.

Main Results:

• Optimized FL configurations (1 epoch, 300 rounds with FedMedian for segmentation; 5 epochs, 200 rounds with FedAdagrad for detection) significantly improved performance over average client models.
• The optimized FL model demonstrated substantial improvements in prostate segmentation (Dice score: 0.73 ± 0.06 to 0.88 ± 0.03; P ≤ .01) and csPCa detection (PI-CAI score: 0.63 ± 0.07 to 0.74 ± 0.06; P ≤ .01).
• The optimized FL model showed superior lesion detection performance compared to the FL-baseline model (PI-CAI score: 0.72 ± 0.06 to 0.74 ± 0.06; P ≤ .01), with no significant difference in segmentation performance.

Conclusions:

• Federated learning enhances the performance and generalizability of AI models for MRI prostate segmentation and csPCa detection compared to local training approaches.
• Optimizing FL framework configurations, including aggregation strategies and training parameters, further boosts the model's lesion detection capabilities.
• This study highlights the potential of federated learning in advancing AI-driven diagnostic tools for prostate cancer while preserving patient data privacy.