FedUP: Uncertainty-Aware Personalized Federated Learning via Probabilistic Prototypes
TL;DR
FedUP is an uncertainty-aware personalized federated learning framework that models prototypes as probability distributions and aggregates prototype distributions through reliability-informed barycenters on the statistical manifold, suppressing unreliable contributions while preserving geometric structure.
Abstract
Prototype-based federated learning enables efficient knowledge sharing by exchanging class prototypes rather than full model parameters. However, heterogeneous client data and limited local samples increase prototype estimation variance, making many client prototypes unreliable. Existing methods usually treat prototypes as deterministic point estimates and cannot quantify their reliability, which may contaminate global prototypes and cause negative transfer. To address these challenges, we propose FedUP, an uncertainty-aware personalized federated learning framework that models prototypes as probability distributions. FedUP captures both aleatoric and epistemic uncertainty on a probabilistic simplex to guide local training and global aggregation. It further uses global probabilistic prototypes as class-conditional priors for feature augmentation, alleviating client-side data sparsity. Finally, FedUP aggregates prototype distributions through reliability-informed barycenters on the statistical manifold, suppressing unreliable contributions while preserving geometric structure. Experiments on natural and medical benchmarks demonstrate that FedUP consistently outperforms state-of-the-art methods.