[Paper Review] Multi-site Organ Segmentation with Federated Partial Supervision and Site Adaptation

Pengbo Liu, Mengke Sun+1 more|arXiv (Cornell University)|Feb 8, 2023

Medical Imaging and Analysis4 citations

TL;DR

This paper proposes a two-stage federated learning framework for multi-site organ segmentation that first trains a unified multi-organ model under partial supervision across sites using a novel loss design to resolve 'knowledge conflict', then adapts this model per site via fine-tuning. The method achieves performance comparable to centrally trained models while preserving data privacy.

ABSTRACT

Objective and Impact Statement: Accurate organ segmentation is critical for many clinical applications at different clinical sites, which may have their specific application requirements that concern different organs. Introduction: However, learning high-quality, site-specific organ segmentation models is challenging as it often needs on-site curation of a large number of annotated images. Security concerns further complicate the matter. Methods: The paper aims to tackle these challenges via a two-phase aggregation-then-adaptation approach. The first phase of federated aggregation learns a single multi-organ segmentation model by leveraging the strength of 'bigger data', which are formed by (i) aggregating together datasets from multiple sites that with different organ labels to provide partial supervision, and (ii) conducting partially supervised learning without data breach. The second phase of site adaptation is to transfer the federated multi-organ segmentation model to site-specific organ segmentation models, one model per site, in order to further improve the performance of each site's organ segmentation task. Furthermore, improved marginal loss and exclusion loss functions are used to avoid 'knowledge conflict' problem in a partially supervision mechanism. Results and Conclusion: Extensive experiments on five organ segmentation datasets demonstrate the effectiveness of our multi-site approach, significantly outperforming the site-per-se learned models and achieving the performance comparable to the centrally learned models.

Motivation & Objective

• To address the challenge of training high-quality, site-specific organ segmentation models when large-scale annotated data are unavailable due to privacy and security constraints.
• To enable collaborative model training across multiple clinical sites without sharing raw medical images, preserving data privacy.
• To resolve the 'knowledge conflict' problem in partially supervised multi-organ segmentation where organs are labeled inconsistently across sites.
• To improve site-specific segmentation performance by adapting a globally trained federated model to local data distributions.
• To develop a scalable, privacy-preserving framework for multi-organ segmentation using federated learning with partial supervision.

Proposed method

• Proposes a two-phase approach: (1) federated aggregation of partially labeled datasets from multiple sites to train a unified multi-organ segmentation model, and (2) site-specific adaptation to improve per-site performance.
• Introduces marginal and exclusion loss functions to mitigate 'knowledge conflict' in partially supervised learning, ensuring foreground objects in one site do not interfere with background in another.
• Employs a combination of Dice, cross-entropy, focal, Top-K, and Lovasz losses—both marginal and exclusion variants—for robust training under partial supervision.
• Uses Federated Averaging (FedAvg) for model aggregation, where clients train locally and send updated weights to a central server for global model update, ensuring data never leave local sites.
• Applies three fine-tuning strategies for site adaptation: full fine-tuning, encoder freezing, and decoder freezing, to tailor the global model to local data distributions.
• Uses a client-server FL setup with model averaging, where each client selects its best local model based on validation performance.

Figure 1: CT images of different partially labeled datasets for different tasks at different clinical sites. The red, green, blue, yellow, and cyan color represent liver, spleen, pancreas, right kidney, and left kidney, respectively.

Experimental results

Research questions

• RQ1 Can a single multi-organ segmentation model be effectively trained across multiple clinical sites with partially labeled data while avoiding knowledge conflict?
• RQ2 How can data privacy be preserved when training a shared segmentation model across institutions with heterogeneous data labels and acquisition protocols?
• RQ3 To what extent can a globally trained federated model be adapted to improve performance on site-specific organ segmentation tasks?
• RQ4 Which fine-tuning strategy (full, encoder-freeze, decoder-freeze) yields the best performance for site adaptation in multi-organ segmentation?
• RQ5 Does the proposed loss design effectively reduce performance degradation due to partial supervision in multi-site settings?

Key findings

• The proposed federated multi-organ segmentation model significantly outperforms models trained independently per site on each dataset, demonstrating the benefit of collaborative learning.
• The model achieves performance comparable to a centrally trained model, even though data are never shared and are only partially labeled across sites.
• The combination of marginal and exclusion losses effectively mitigates the 'knowledge conflict' problem in partially supervised multi-organ segmentation.
• Site adaptation via fine-tuning further improves segmentation performance on individual sites, with the best results achieved using encoder-freezing or full fine-tuning strategies.
• The method maintains strong generalization across five different organ segmentation datasets, confirming its robustness and scalability.
• The ablation study confirms that both marginal and exclusion losses contribute significantly to performance gains, especially in handling inconsistent labeling across sites.

Figure 2: Visualization of the segmentation results of our proposed method and other SOTA methods.