# UniFed: A Universal Federation of a Mixture of Highly Heterogeneous Medical Image Classification Tasks

arXiv: https://arxiv.org/abs/2408.07075 · HTML: https://mlsanity.com/paper/2408.07075/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 27 of its 37 references that are in the map.
- First posted on arXiv 29 Jul 2024.

## Builds on (most cited first, 27 of 27)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://mlsanity.com/paper/2102.12122/) (2021)
- [MedMNIST v2 -- A large-scale lightweight benchmark for 2D and 3D biomedical image classification](https://mlsanity.com/paper/2110.14795/) (2021)
- [Federated Multi-Task Learning](https://mlsanity.com/paper/1705.10467/) (2017)
- [Three Approaches for Personalization with Applications to Federated Learning](https://mlsanity.com/paper/2002.10619/) (2020)
- [The Non-IID Data Quagmire of Decentralized Machine Learning](https://mlsanity.com/paper/1910.00189/) (2019)
- [Curriculum Learning: A Survey](https://mlsanity.com/paper/2101.10382/) (2021)
- [Federated Multi-Task Learning under a Mixture of Distributions](https://mlsanity.com/paper/2108.10252/) (2021)
- [FedALA: Adaptive Local Aggregation for Personalized Federated Learning](https://mlsanity.com/paper/2212.01197/) (2022)
- [FairFed: Enabling Group Fairness in Federated Learning](https://mlsanity.com/paper/2110.00857/) (2021)
- [FedFed: Feature Distillation against Data Heterogeneity in Federated Learning](https://mlsanity.com/paper/2310.05077/) (2023)
- [A Survey on Heterogeneous Federated Learning](https://mlsanity.com/paper/2210.04505/) (2022)
- [Sparse Random Networks for Communication-Efficient Federated Learning](https://mlsanity.com/paper/2209.15328/) (2022)
- [Convergence Analysis of Sequential Federated Learning on Heterogeneous Data](https://mlsanity.com/paper/2311.03154/) (2023)
- [Learning Tasks for Multitask Learning: Heterogenous Patient Populations in the ICU](https://mlsanity.com/paper/1806.02878/) (2018)
- [Self-Paced Learning: an Implicit Regularization Perspective](https://mlsanity.com/paper/1606.00128/) (2016)
- [FedFMC: Sequential Efficient Federated Learning on Non-iid Data](https://mlsanity.com/paper/2006.10937/) (2020)
- [FedBone: Towards Large-Scale Federated Multi-Task Learning](https://mlsanity.com/paper/2306.17465/) (2023)
- [MAS: Towards Resource-Efficient Federated Multiple-Task Learning](https://mlsanity.com/paper/2307.11285/) (2023)
- [Asynchronous Multi-Task Learning](https://mlsanity.com/paper/1609.09563/) (2016)
- [FedLPS: Heterogeneous Federated Learning for Multiple Tasks with Local Parameter Sharing](https://mlsanity.com/paper/2402.08578/) (2024)
- [Personalized Federated Learning with Server-Side Information](https://mlsanity.com/paper/2205.11044/) (2022)
- [Federated Learning with Server Learning: Enhancing Performance for Non-IID Data](https://mlsanity.com/paper/2210.02614/) (2022)

## Abstract (arXiv)

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.