Uncertainty and Scale-Calibrated Contrastive Federated Segmentation under Client Heterogeneity

TMLR Paper7517 Authors

Abstract: Federated learning presents a promising approach for medical image segmentation, particularly in addressing data privacy concerns. However, it faces significant challenges due to data heterogeneity across participating clients. This heterogeneity introduces variations in data scales and distributions, making it difficult to balance spatial accuracy and feature similarity when managing multidimensional heterogeneous data. To address these challenges, we propose a novel \textbf{Uncertainty- and Scale-Calibrated Contrastive Federated Segmentation under Client Heterogeneity (SAFCF)} with two key approaches: (i) an \textbf{uncertainty-driven dynamic scale-adaptive weighted aggregation (DSWA)} method, which balances the influence of local client data scales and reduces model drift caused by data heterogeneity through the use of epistemic uncertainty in weighted aggregation, and (ii) a \textbf{contrastive federated segmentation loss (CFSL)}, a local loss function that effectively balances spatial accuracy and feature similarity at the pixel level of an image by combining modified Dice loss with improved contrastive loss. Additionally, epistemic uncertainty layer learns weight distributions to introduce uncertainty, further improving model robustness and enabling adaptive learning from diverse data during training. Our framework demonstrates substantial improvements on standard benchmark medical image segmentation datasets, especially under highly non-IID conditions, when compared to traditional algorithms.

Submission Type: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Konstantin_Mishchenko1

Submission Number: 7517