UniFed: A Universal Federation of a Mixture of Highly Heterogeneous Medical Image Classification Tasks
This addresses the problem of divergent model performance and high communication costs in federated learning for precision medicine, particularly in real-world hospital settings with diverse imaging modalities and tasks, representing a novel method for a known bottleneck.
The paper tackles the challenge of federated learning with highly heterogeneous medical image classification tasks by introducing UniFed, a universal federated learning paradigm that dynamically adjusts models to handle varying task complexities, achieving improvements in accuracy, communication cost, and convergence time for diagnosing diseases like retina, histopathology, and liver tumors.
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.