Weighting Schemes for Federated Learning in Heterogeneous and Imbalanced Segmentation Datasets

Federated learning allows for training deep learning models from various sources (e.g., hospitals) without sharing patient information, but only the model weights. Two central problems arise when sending the updated weights to the central node in a federation: the imbalance of the datasets and data heterogeneity caused by differences in scanners or acquisition protocols. In this paper, we benchmark the federated average algorithm and adapt two weighting functions to counteract the effect of data imbalance. The approaches are validated on a segmentation task with synthetic data from imbalanced centers, and on two multi-centric datasets with the clinically relevant tasks of stroke infarct core prediction and brain tumor segmentation. The results show that accounting for the imbalance in the data sources improves the federated average aggregation in different perfusion CT and structural MRI images in the ISLES and BraTS19 datasets, respectively.