ConDistFL: Conditional Distillation for Federated Learning from Partially Annotated Data

TL;DR

ConDistFL introduces a federated learning framework that leverages knowledge distillation with conditional probability representations to train generalized segmentation models from partially annotated data, reducing communication costs and improving performance.

Contribution

The paper presents a novel federated learning approach combining knowledge distillation and conditional probability to handle partially annotated data for segmentation tasks.

Findings

Outperforms FedAvg and FedOpt baselines.

Demonstrates superior generalizability on external datasets.

Reduces communication costs by less frequent aggregation.

Abstract

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.