Uncertainty-Aware Personalized Federated Learning for Realistic Healthcare Applications

Yuwei Zhang, Tong Xia, Abhirup Ghosh, Cecilia Mascolo
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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v259-zhang25b, title = {Uncertainty-Aware Personalized Federated Learning for Realistic Healthcare Applications}, author = {Zhang, Yuwei and Xia, Tong and Ghosh, Abhirup and Mascolo, Cecilia}, booktitle = {Proceedings of the 4th Machine Learning for Health Symposium}, pages = {1067--1086}, year = {2025}, editor = {Hegselmann, Stefan and Zhou, Helen and Healey, Elizabeth and Chang, Trenton and Ellington, Caleb and Mhasawade, Vishwali and Tonekaboni, Sana and Argaw, Peniel and Zhang, Haoran}, volume = {259}, series = {Proceedings of Machine Learning Research}, month = {15--16 Dec}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v259/main/assets/zhang25b/zhang25b.pdf}, url = {https://proceedings.mlr.press/v259/zhang25b.html}, 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.} }
Endnote
%0 Conference Paper %T Uncertainty-Aware Personalized Federated Learning for Realistic Healthcare Applications %A Yuwei Zhang %A Tong Xia %A Abhirup Ghosh %A Cecilia Mascolo %B Proceedings of the 4th Machine Learning for Health Symposium %C Proceedings of Machine Learning Research %D 2025 %E Stefan Hegselmann %E Helen Zhou %E Elizabeth Healey %E Trenton Chang %E Caleb Ellington %E Vishwali Mhasawade %E Sana Tonekaboni %E Peniel Argaw %E Haoran Zhang %F pmlr-v259-zhang25b %I PMLR %P 1067--1086 %U https://proceedings.mlr.press/v259/zhang25b.html %V 259 %X 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.
APA
Zhang, Y., Xia, T., Ghosh, A. & Mascolo, C.. (2025). Uncertainty-Aware Personalized Federated Learning for Realistic Healthcare Applications. Proceedings of the 4th Machine Learning for Health Symposium, in Proceedings of Machine Learning Research 259:1067-1086 Available from https://proceedings.mlr.press/v259/zhang25b.html.

Related Material