[edit]
Uncertainty-Aware Personalized Federated Learning for Realistic Healthcare Applications
Proceedings of the 4th Machine Learning for Health Symposium, PMLR 259:1067-1086, 2025.
Abstract
Healthcare applications require accurate and uncertainty-aware machine learning models, providing confidence rather than only blackbox predictions. However, training such deep learning models with insufficient data at individual sites (e.g., hospitals) poses a challenge. Federated learning (FL) mitigates this by allowing data holders to collaboratively train models without sharing sensitive health data. Yet, we identify two major realistic challenges when building uncertainty estimates in FL, severe data heterogeneity and high computational overhead. This paper proposes FedEE, an uncertainty-aware and efficient personalized FL framework for realistic healthcare applications. FedEE achieves an efficient way of ensembling by incorporating lightweight early exit blocks into a single backbone model. These blocks are personalized for each client to tackle data heterogeneity. Experiments with four FL strategies and three datasets demonstrate that FedEE achieves up to a 15% improvement in uncertainty estimation from vanilla softmax entropy and is competitive with expensive baselines, showcasing in the order of 5× improved efficiency with a 5-member ensemble.