Model-Heterogeneous Semi-Supervised Federated Learning for Medical Image Segmentation

Background Introduction

Medical image segmentation plays a crucial role in clinical diagnosis as it helps doctors identify and analyze diseases. However, this task typically faces challenges such as sensitive data, privacy issues, and expensive annotation costs. While current research mainly focuses on personalized collaborative training for medical segmentation systems, it overlooks the fact that obtaining segmentation annotations is time-consuming and labor-intensive. Balancing annotation costs and segmentation performance while maintaining personalized local models has become an important research direction. Therefore, this study introduces a novel model-heterogeneous semi-supervised federated learning framework.

Paper Source

This paper, titled “Model-Heterogeneous Semi-Supervised Federated Learning for Medical Image Segmentation,” is co-authored by Yuxi Ma, Jiacheng Wang, Jing Yang, and Liansheng Wang. The authors are affiliated with the National Institute of Data Science and Health Engineering at Xiamen University and the Department of Computer Science at Xiamen University. The paper is published in the May 2024 issue of the IEEE Transactions on Medical Imaging.

Research Methods and Process

Workflow

1. Research Objective: The primary objective of this study is to propose a model-heterogeneous semi-supervised federated learning (HSSF) framework that reduces annotation burdens by leveraging unlabeled data while maintaining model personalization.

2. Key Modules: The HSSF framework introduces regularity condensation and regularity fusion to ensure personalization across different sites. Additionally, the self-assessment (SA) module and reliable pseudo-label generation (RPG) module are proposed to fully utilize unlabeled data.

3. Regularity Condensation and Regularity Fusion:

   • Regularity Condensation: The server selects the rules closest to the ground truth based on the performance of models uploaded by various sites and broadcasts these rules to all sites.
   • Regularity Fusion: Each site integrates selected beneficial rules according to its own needs.

4. Semi-Supervised Learning:

   • SA module generates real-time self-assessment confidence to detect the accuracy of model predictions.
   • RPG module generates reliable pseudo-labels based on SA confidence. Combined, these modules improve the training efficiency of unlabeled data.

Experimental Design

We evaluated the models on skin lesion and polyp lesion datasets. The experimental results demonstrated that the HSSF model outperformed other methods in heterogeneous designs and also performed well in homogeneous designs, especially in regional evaluation metrics.

Algorithm Details

The workflow of HSSF can be summarized in the following steps:

1. Data Preparation: Prepare a public dataset and multiple local datasets. The public dataset is used for shared knowledge, while local datasets are used for local operations and remain private.
2. Regularity Condensation: In every communication cycle, the server first selects the optimal rules from all local models for condensation.
3. Local Model Updates: Each site selectively integrates these rules according to its own needs and uses local unlabeled data for self-training to update its model.

Specific Methods of Semi-Supervised Learning

1. Labeled Training: First, perform supervised training on the public dataset to train the model and the SA module to effectively generate self-assessment confidence.
2. Unlabeled Self-Training: Use the confidence scores generated by the SA module to correct pseudo-labels, thereby improving their reliability. Train the local unlabeled data using consistency regularization and self-training.

Research Results

1. Evaluation Metrics: Dice coefficient and Hausdorff distance were used as the primary evaluation metrics to quantify the segmentation performance of the model. Results showed that HSSF excelled in all considered cases.

2. Comparative Experiments: HSSF was compared with other state-of-the-art federated learning methods using homogeneous and heterogeneous models. Experimental results validated the effectiveness and superiority of HSSF.

3. Ablation Experiments: Ablation experiments confirmed the effectiveness of each module, demonstrating that the SA module and RPG module significantly enhanced model performance in various scenarios.

Research Conclusion

The HSSF framework proposed in this study exhibits superior performance in medical image segmentation tasks. By introducing the SA and RPG modules, the model can self-assess and correct pseudo-labels, thereby improving the efficiency of utilizing unlabeled data. Additionally, the unique regularity condensation and regularity fusion mechanisms in the HSSF framework ensure model personalization across different sites, achieving efficient knowledge sharing in federated learning. HSSF not only marks significant progress in the medical image segmentation task but also provides an effective solution for future cross-site collaborative learning. The framework has high potential and value in practical applications by reducing annotation burdens, protecting data privacy, and enhancing model personalization.