Divide, Conquer and Unite: Hierarchical Style-Recalibrated Prototype Alignment for Federated Medical Image Segmentation

TL;DR

Federated learning for medical image segmentation suffers from cross-site heterogeneity due to different scanners and protocols. The authors propose FedBCS, a two-fold approach combining frequency-domain style recalibration to construct domain-invariant prototypes and context-aware dual-level prototype alignment to fuse multi-level encoder–decoder information. They provide convergence guarantees and demonstrate improved Dice scores on histology nuclei and prostate MRI datasets, with favorable communication efficiency. The work advances robust cross-domain segmentation without sharing raw data, offering practical benefits for multi-institution clinical deployment.

Abstract

Federated learning enables multiple medical institutions to train a global model without sharing data, yet feature heterogeneity from diverse scanners or protocols remains a major challenge. Many existing works attempt to address this issue by leveraging model representations (e.g., mean feature vectors) to correct local training; however, they often face two key limitations: 1) Incomplete Contextual Representation Learning: Current approaches primarily focus on final-layer features, overlooking critical multi-level cues and thus diluting essential context for accurate segmentation. 2) Layerwise Style Bias Accumulation: Although utilizing representations can partially align global features, these methods neglect domain-specific biases within intermediate layers, allowing style discrepancies to build up and reduce model robustness. To address these challenges, we propose FedBCS to bridge feature representation gaps via domain-invariant contextual prototypes alignment. Specifically, we introduce a frequency-domain adaptive style recalibration into prototype construction that not only decouples content-style representations but also learns optimal style parameters, enabling more robust domain-invariant prototypes. Furthermore, we design a context-aware dual-level prototype alignment method that extracts domain-invariant prototypes from different layers of both encoder and decoder and fuses them with contextual information for finer-grained representation alignment. Extensive experiments on two public datasets demonstrate that our method exhibits remarkable performance.

Figures (5)

• Figure 1: Problem illustration of existing methods under domain skew setting. The top row highlights two major challenges—incomplete contextual representation learning and layerwise style bias accumulation.
• Figure 2: Architecture illustration of FedBCS. To address layerwise style bias accumulation, we implement frequency-domain style recalibration during local prototype construction. To capture complete contextual representations, we extract and align multi-level prototypes from both encoder and decoder pathways. These prototypes are aggregated at the server to derive global prototypes for guiding local training. Zoom in for details.
• Figure 3: Analysis of FedBCS with different temperature (\ref{['tau']}). "Base" denotes FedAvg. See details in \ref{['sec:ablation']}.
• Figure 4: Performance analysis. (a) Convergence comparison of differentmethods across communication rounds. (b) Effectiveness of FSR in addressing layerwise style bias compared to input-level normalization. See details in \ref{['sec:ablation']}.
• Figure 5: Qualitative comparison of segmentation results between our method and other state-of-the-art approaches. The top two rows present results for the prostate MRI segmentation task, while the bottom two rows correspond to the histology nuclei segmentation task.