TreeFedDG: Alleviating Global Drift in Federated Domain Generalization for Medical Image Segmentation

TL;DR

This work defines and tackles the global drift issue in federated domain generalization (FedDG) for medical imaging. It introduces TreeFedDG, a tree-structured, hierarchical aggregation framework that suppresses drift by clustering client models by parameter similarity, while maintaining multi-granularity knowledge from leaf to root. It further contributes FedStyle for targeted style mixing, a progressive personalized fusion mechanism for balanced global-local knowledge transfer, and a feature similarity-guided inference strategy that ensembles model chains along the tree. Across Fundus and Prostate MRI segmentation tasks, TreeFedDG achieves state-of-the-art performance and substantially improved cross-domain consistency, demonstrating its practical potential for privacy-preserving, multi-site medical imaging deployments.

Abstract

In medical image segmentation tasks, Domain Generalization (DG) under the Federated Learning (FL) framework is crucial for addressing challenges related to privacy protection and data heterogeneity. However, traditional federated learning methods fail to account for the imbalance in information aggregation across clients in cross-domain scenarios, leading to the Global Drift (GD) problem and a consequent decline in model generalization performance. This motivates us to delve deeper and define a new critical issue: global drift in federated domain generalization for medical imaging (FedDG-GD). In this paper, we propose a novel tree topology framework called TreeFedDG. First, starting from the distributed characteristics of medical images, we design a hierarchical parameter aggregation method based on a tree-structured topology to suppress deviations in the global model direction. Second, we introduce a parameter difference-based style mixing method (FedStyle), which enforces mixing among clients with maximum parameter differences to enhance robustness against drift. Third, we develop a a progressive personalized fusion strategy during model distribution, ensuring a balance between knowledge transfer and personalized features. Finally, during the inference phase, we use feature similarity to guide the retrieval of the most relevant model chain from the tree structure for ensemble decision-making, thereby fully leveraging the advantages of hierarchical knowledge. We conducted extensive experiments on two publicly available datasets. The results demonstrate that our method outperforms other state-of-the-art domain generalization approaches in these challenging tasks and achieves better balance in cross-domain performance.

Figures (6)

• Figure 1: Our Motivation. (a) The basic setup of Federated Domain Generalization (FedDG). (b) When domain heterogeneity is large, the FedDG model suffers from a global drift problem (FedDG-GD), resulting in significant performance differences across different sites.
• Figure 2: Overall framework workflow. (a) Hierarchical parameter aggregation and progressive personalization fusion of the tree-structured model. (b) Style mixing during client training. (c) Similarity-guided model selection during the inference phase.
• Figure 3: Qualitative comparison of the generalization results of different methods for fundus image segmentation (the top two rows) and prostate MRI segmentation (the bottom two rows)
• Figure 4: Qualitative comparison of Dice performance consistency across different methods for retinal image segmentation and prostate MRI segmentation. (Lower STD values indicate smaller performance variations and stronger model generalization capabilities.)
• Figure 5: Ablation study results on (a) the effectiveness of the tree-structured architecture, (b) the impact of the FedStyle method, and (c) the impact of the progressive fusion mechanism.