ConDistFL: Conditional Distillation for Federated Learning from Partially Annotated Data
This work addresses the problem of limited fully annotated data for medical image segmentation in federated learning, offering a domain-specific solution that is incremental but impactful for healthcare applications.
The paper tackles the challenge of training generalizable segmentation models in federated learning with partially annotated data by proposing ConDistFL, which combines FL with knowledge distillation to extract knowledge from unlabeled organs and tumors, resulting in significant performance improvements over FedAvg and FedOpt baselines and superior generalizability on external test datasets.
Developing a generalized segmentation model capable of simultaneously delineating multiple organs and diseases is highly desirable. Federated learning (FL) is a key technology enabling the collaborative development of a model without exchanging training data. However, the limited access to fully annotated training data poses a major challenge to training generalizable models. We propose "ConDistFL", a framework to solve this problem by combining FL with knowledge distillation. Local models can extract the knowledge of unlabeled organs and tumors from partially annotated data from the global model with an adequately designed conditional probability representation. We validate our framework on four distinct partially annotated abdominal CT datasets from the MSD and KiTS19 challenges. The experimental results show that the proposed framework significantly outperforms FedAvg and FedOpt baselines. Moreover, the performance on an external test dataset demonstrates superior generalizability compared to models trained on each dataset separately. Our ablation study suggests that ConDistFL can perform well without frequent aggregation, reducing the communication cost of FL. Our implementation will be available at https://github.com/NVIDIA/NVFlare/tree/dev/research/condist-fl.