FedMix: Mixed Supervised Federated Learning for Medical Image Segmentation

Abstract

The purpose of federated learning is to enable multiple clients to jointly train a machine learning model without sharing data. However, the existing methods for training an image segmentation model have been based on an unrealistic assumption that the training set for each local client is annotated in a similar fashion and thus follows the same image supervision level. To relax this assumption, in this work, we propose a label-agnostic unified federated learning framework, named FedMix, for medical image segmentation based on mixed image labels. In FedMix, each client updates the federated model by integrating and effectively making use of all available labeled data ranging from strong pixel-level labels, weak bounding box labels, to weakest image-level class labels. Based on these local models, we further propose an adaptive weight assignment procedure across local clients, where each client learns an aggregation weight during the global model update. Compared to the existing methods, FedMix not only breaks through the constraint of a single level of image supervision, but also can dynamically adjust the aggregation weight of each local client, achieving rich yet discriminative feature representations. To evaluate its effectiveness, experiments have been carried out on two challenging medical image segmentation tasks, i.e., breast tumor segmentation and skin lesion segmentation. The results validate that our proposed FedMix outperforms the state-of-the-art methods by a large margin 1 .

Simon Fraser University's discourse processing lab, Burnaby, BC, Canada

“…We compare our QA-SplitFed against five baseline methods: Naive SplitFed (W = 1 [8] (SGD optimizer and momentum of 0.9), SplitAVG [6], and FedMix [7] (β = 1.5, λ = 10). All models are trained and validated over E = 12 local epochs and G = 10 global epochs.…”

“…Different types of data heterogeneity could occur in the context of collaborative learning, e.g. unbalanced data [2], different feature map distributions [6], or different levels of annotation accuracies over clients [7]. In the literature, these are commonly referred to as non-IID (Independent and Identically Distributed) scenarios.…”

“…Their proposed SplitFed algorithm was tested on CIFAR-10 and CIFAR-100 datasets. A weighted averaging strategy was proposed in [7] that took advantage of labeled data for creating pseudo labels for unlabeled data. It used clients with better training quality when updating the global model.…”

Current and Emerging Trends in Medical Image Segmentation With Deep Learning

“…The authors also suggested a new federated cross-ensemble learning technique that together trains and sets up various models. Wicaksana et al [213] proposed FedMix, an FL strategy that employed mixed image labels specifically to segment anatomical region-of-interest from medical images. These labels incorporated substantial pixel-wise annotations, weak bounding boxes, and image-wise class annotations.…”

Regularized Weight Aggregation in Networked Federated Learning for Glioblastoma Segmentation

“…The challenge of data heterogeneity and domain shifting was recently tackled in novel ways by, for example, federated disentanglement learning via disentangling the parameter space into shape and appearance [6] and automated federated averaging based on Dirichlet distribution [22]. Dynamic Re-Weighting mechanisms [12], federated cross ensemble learning [24], and label-agnostic (mixed labels) unified FL formed by a mixture of the client distributions [21] have been recently proposed to relax an unrealistic assumption that each client's training set will be annotated similarly and therefore follows the same image supervision level during the training of an image segmentation model. Although extensive research has been carried out on FL, there is still a need for methods to enable the development of more generalized FL models for clinical use which can effectively deal with statistical heterogeneity in weight aggregation, communication efficiency, and privacy with security.…”