Unifying and Personalizing Weakly-supervised Federated Medical Image Segmentation via Adaptive Representation and Aggregation

Abstract: Federated learning (FL) enables multiple sites to collaboratively train powerful deep models without compromising data privacy and security. The statistical heterogeneity (e.g., non-IID data and domain shifts) is a primary obstacle in FL, impairing the generalization performance of the global model. Weakly supervised segmentation, which uses sparsely-grained (i.e., point-, bounding box-, scribble-, block-wise) supervision, is increasingly being paid attention to due to its great potential of reducing annotation costs. However, there may exist label heterogeneity, i.e., different annotation forms across sites. In this paper, we propose a novel personalized FL framework for medical image segmentation, named FedICRA, which uniformly leverages heterogeneous weak supervision via adaptIve Contrastive Representation and Aggregation. Concretely, to facilitate personalized modeling and to avoid confusion, a channel selection based site contrastive representation module is employed to adaptively cluster intra-site embeddings and separate inter-site ones. To effectively integrate the common knowledge from the global model with the unique knowledge from each local model, an adaptive aggregation module is applied for updating and initializing local models at the element level. Additionally, a weakly supervised objective function that leverages a multiscale tree energy loss and a gated CRF loss is employed to generate more precise pseudo-labels and further boost the segmentation performance. Through extensive experiments on two distinct medical image segmentation tasks of different modalities, the proposed FedICRA demonstrates overwhelming performance over other state-of-the-art personalized FL methods. Its performance even approaches that of fully supervised training on centralized data. Our code and data are available at https://github.com/llmir/FedICRA.

• The paper introduces FedICRA, a federated learning framework that unifies adaptive representation and personalized aggregation for weakly-supervised medical image segmentation.
• It demonstrates significantly improved segmentation accuracy with lower Hausdorff Distance compared to state-of-the-art methods across diverse imaging modalities.
• The approach enables personalized model adaptation to heterogeneous datasets, paving the way for enhanced clinical applicability and improved data privacy.

Unifying and Personalizing Weakly-supervised Federated Medical Image Segmentation

Introduction

The paper "Unifying and Personalizing Weakly-supervised Federated Medical Image Segmentation via Adaptive Representation and Aggregation" (2304.05635) tackles the challenge of medical image segmentation, a fundamental task in computer vision with significant implications in healthcare. The research introduces an innovative federated learning framework, termed FedICRA, aimed at enhancing weakly-supervised segmentation across diverse medical datasets. This methodology integrates adaptive representation learning and personalized aggregation strategies to address the heterogeneity in data distribution and annotation styles inherent in multi-center medical datasets.

Methodology

FedICRA leverages two key innovations: adaptive representation learning and a personalized aggregation method. The adaptive representation focuses on creating a generalizable feature space capable of handling variations across different datasets and annotation types. This approach is crucial in weakly-supervised settings where annotations are sparse or incomplete. The personalized aggregation method, on the other hand, allows for individual adaptation to client-specific data characteristics, thereby improving overall model performance across heterogeneous datasets.

The paper demonstrates the implementation of the proposed approach on two modalities: fundus imaging and Optical Coherence Tomography Angiography (OCTA), utilizing several datasets with varying annotation methods, such as scribbles, bounding boxes, and points. This setup highlights the flexibility and robustness of FedICRA in handling diverse medical imaging scenarios.

Figure 1: Examples of the image and the corresponding sparse annotation from each site. UA, OC, OD, FAZ, and BG respectively represent unlabeled area, optic cup, optic disc, foveal avascular zone, and background.

Results

FedICRA's performance was evaluated on the ODOC segmentation task, where it demonstrated superior results compared to several state-of-the-art (SOTA) federated learning methods such as FedAvg, FedProx, and MetaFed. The results indicated that FedICRA achieved the lowest Hausdorff Distance (HD95) across multiple test sites, signifying enhanced precision in optic disc and optic cup segmentation tasks, even with varying degrees of annotation quality.

Figure 2: Visualization results from FedICRA and other SOTA methods (CT indicates centralized training, with Weak and Full respectively denoting utilizing sparse annotations and full masks).

The visualization results corroborate these findings, where FedICRA's segmentation outputs are consistently more accurate than those of competing methods, particularly in dealing with incomplete and sparse annotations. The paper emphasizes that FedICRA's personalized models result in more consistent and generalizable segmentation outcomes across different datasets.

Discussion

The primary contribution of this research is the development of a federated learning approach tailored to the unique challenges presented by medical image segmentation tasks with weak supervision. By effectively unifying adaptation and personalization, FedICRA offers a robust solution that significantly narrows the performance gap between weakly and fully supervised learning paradigms. This research highlights the potential of federated learning to democratize access to cutting-edge medical AI technologies while respecting patient data privacy, a critical consideration in healthcare.

Moreover, the proposed methodology has profound implications for future research in federated learning and weakly-supervised segmentation. It paves the way for further explorations into adaptive learning techniques that can better handle the distributional shifts and irregularities in medical data. The adaptability and generalization capabilities demonstrated by FedICRA may extend beyond medical imaging to other domains where annotated data is scarce or costly to acquire.

Conclusion

The paper effectively presents FedICRA, a first-of-its-kind federated learning framework that empowers weakly-supervised medical image segmentation by introducing adaptive representation and personalized aggregation. The findings illustrate significant advancements in segmentation accuracy and reliability, establishing a new benchmark for federated learning applications in medical imaging. Future work will likely explore extending these techniques to broader clinical applications and other imaging modalities, reinforcing the critical role of federated learning in advancing healthcare AI.