UniFed: A Universal Federation of a Mixture of Highly Heterogeneous Medical Image Classification Tasks
Atefe Hassani, Islem Rekik
TL;DR
The paper addresses federated learning across highly heterogeneous medical image tasks and data distributions, where communication costs hinder practical deployment. It introduces UniFed, a universal federated framework that uses the slope of task-specific loss curves to sequence hospitals by task complexity, couples curriculum learning with dynamic local/global updates, and employs sequential model transfer between server and clients along with a small server-side mixed data regularization. Key contributions include a loss-slope based ordering, dynamic convergence-driven training, sequential client scheduling, and a mixing-based server update, which collectively improve accuracy and reduce computation and communication overhead across OCTMNIST, OrganAMNIST, and TissueMNIST under strongly and moderately Non-IID settings. This approach advances precision medicine FL by enabling a single model that generalizes across imaging modalities and diseases while mitigating resource use; future work explores knowledge distillation and evaluation on foundational medical models.
Abstract
A fundamental challenge in federated learning lies in mixing heterogeneous datasets and classification tasks while minimizing the high communication cost caused by clients as well as the exchange of weight updates with the server over a fixed number of rounds. This results in divergent model convergence rates and performance, which may hinder their deployment in precision medicine. In real-world scenarios, client data is collected from different hospitals with extremely varying components (e.g., imaging modality, organ type, etc). Previous studies often overlooked the convoluted heterogeneity during the training stage where the target learning tasks vary across clients as well as the dataset type and their distributions. To address such limitations, we unprecedentedly introduce UniFed, a universal federated learning paradigm that aims to classify any disease from any imaging modality. UniFed also handles the issue of varying convergence times in the client-specific optimization based on the complexity of their learning tasks. Specifically, by dynamically adjusting both local and global models, UniFed considers the varying task complexities of clients and the server, enhancing its adaptability to real-world scenarios, thereby mitigating issues related to overtraining and excessive communication. Furthermore, our framework incorporates a sequential model transfer mechanism that takes into account the diverse tasks among hospitals and a dynamic task-complexity based ordering. We demonstrate the superiority of our framework in terms of accuracy, communication cost, and convergence time over relevant benchmarks in diagnosing retina, histopathology, and liver tumour diseases under federated learning. Our UniFed code is available at https://github.com/basiralab/UniFed.