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

TL;DR

TreeFedDG introduces a hierarchical federated learning framework with style mixing and personalized fusion to improve medical image segmentation across diverse domains, effectively addressing global drift and enhancing model generalization.

Contribution

The paper proposes a novel tree topology framework with style mixing and personalized fusion strategies to mitigate global drift in federated domain generalization for medical imaging.

Findings

Outperforms state-of-the-art domain generalization methods.

Achieves better cross-domain performance balance.

Demonstrates robustness against data heterogeneity.

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.