Federated Learning with Partially Labeled Data: A Conditional Distillation Approach

2024-12-25

AI Summary

In medical federated learning, model divergence and catastrophic forgetting arise due to partial organ/lesion annotations across centers, label inconsistency, and large inter-center intensity distribution shifts (e.g., contrast-enhanced vs. non-contrast CT). To address this without requiring fully annotated data or cross-center label alignment, we propose a conditional feature distillation framework. It enforces semantic consistency locally at each client via a label-aware loss and a lightweight communication protocol. Evaluated on multi-center 3D CT and 2D chest X-ray datasets, our method achieves superior segmentation accuracy over state-of-the-art methods while reducing both computational and communication overhead. Its core innovation lies in the first conditional distillation mechanism explicitly designed for partial annotations and robustness to cross-domain intensity variations—effectively alleviating knowledge transfer bottlenecks under data scarcity and strict privacy constraints.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Computer Vision: Segmentation
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

In medical imaging, developing generalized segmentation models that can handle multiple organs and lesions is crucial. However, the scarcity of fully annotated datasets and strict privacy regulations present significant barriers to data sharing. Federated Learning (FL) allows decentralized model training, but existing FL methods often struggle with partial labeling, leading to model divergence and catastrophic forgetting. We propose ConDistFL, a novel FL framework incorporating conditional distillation to address these challenges. ConDistFL enables effective learning from partially labeled datasets, significantly improving segmentation accuracy across distributed and non-uniform datasets. In addition to its superior segmentation performance, ConDistFL maintains computational and communication efficiency, ensuring its scalability for real-world applications. Furthermore, ConDistFL demonstrates remarkable generalizability, significantly outperforming existing FL methods in out-of-federation tests, even adapting to unseen contrast phases (e.g., non-contrast CT images) in our experiments. Extensive evaluations on 3D CT and 2D chest X-ray datasets show that ConDistFL is an efficient, adaptable solution for collaborative medical image segmentation in privacy-constrained settings.

Problem

• Federated Learning
• Medical Image Analysis
• Privacy Protection

Innovation

• ConDistFL
• Conditional Distillation
• Federated Learning