FedUP: Uncertainty-Aware Personalized Federated Learning via Probabilistic Prototypes

1 Introduction

Personalized Federated Learning (PFL) addresses statistical heterogeneity (Non-IID) by balancing global knowledge sharing with local adaptation, thereby mitigating client drift and performance degradation [Smith et al., 2017; Kairouz and McMahan, 2021]. However, many PFL methods require sharing full model parameters, causing high communication overhead. Prototype-based methods such as FedProto [Tan et al., 2022a] reduce this cost by exchanging class prototypes, which regularize local training through featureprototype alignment and support heterogeneous architectures.

* Corresponding authors: Weishan Zhang and Lingzhao Meng.

Despite their efficiency, existing prototype-based methods represent prototypes as deterministic point estimates and implicitly assume uniform reliability across clients. However, under data heterogeneity, prototype quality varies significantly. Clients with scarce or biased data produce unreliable prototypes that poorly represent true class distributions, while data-rich clients generate more accurate ones [Fang et al., 2025; Chen et al., 2025]. Such reliability variation is particularly prominent in medical image analysis, where class imbalance, image noise, and sample-level difficulty often lead to biased or unreliable supervision signals [Meng et al., 2026; Guo et al., 2026]. This heterogeneity in prototype quality, combined with the lack of reliability quantification, leads to two critical problems. (1) Intra-client: Unreliable prototypes mislead local training. Existing methods apply fixedstrength alignment for all clients. However, clients with unreliable prototypes should rely more on global guidance, while those with reliable prototypes should preserve local knowledge. Without reliability quantification, existing methods cannot adapt alignment strength accordingly. How to quantify prototype reliability and adaptively balance global alignment with local personalization (Challenge #1)? (2) Inter-client: Unreliable prototypes degrade global aggregation. Existing methods aggregate client prototypes without considering reliability differences. This allows unreliable prototypes from data-scarce clients to dominate the global prototype, causing negative transfer [Wang et al., 2025]. How to robustly aggregate prototypes with varying reliability while suppressing negative transfer (Challenge #2)?

To address both challenges, we propose FedUP, an uncertainty-aware personalized federated learning framework. Built upon the Dirichlet-based evidential deep learning, FedUP models the categorical prediction of a sample as following a Dirichlet distribution, thus allowing multiple potential predictions for a sample and enabling principled decomposition of predictive uncertainty into epistemic and aleatoric. On the client side, FedUP introduces an adaptive training strategy that leverages decomposed uncertainty for personalization. Epistemic uncertainty is utilized to quantify the model's lack of knowledge for each sample and serves as a dynamic gating mechanism to modulate the alignment intensity between local features and global prototypes. To mitigate the statistical vulnerability caused by the scarcity of local data, FedUP introduces a feature enhancement mechanism. The global probability prototype is regarded as a classconditional prior, encoding the semantic knowledge of the categories in the federation. While enriching the sparse local representation, it also standardizes the decision boundary of the classifier.

On the server side, as illustrated in Figure 1, unlike traditional methods that aggregate prototypes via naive averaging without considering reliability, FedUP models each client prototype as a Gaussian distribution parameterized by uncertainty-aware statistics, where uncertainty inflates the covariance to reflect prototype reliability. Prototype aggregation is then formulated as computing the barycenter on the statistical manifold, weighted by client reliability derived from both uncertainty types. This geometry-preserving aggregation avoids the variance collapse inherent in naive Euclidean averaging while effectively filtering negative transfer from low-quality prototypes contributed by data-scarce or noisy clients. We summarize our contributions as follows:

• We propose a framework that regards the prototype as a probability distribution, and quantify the reliability of the prototype.
• We design an uncertainty-driven adaptive training mechanism, which dynamically adjusts the global alignment strength and enriches the sparse feature space, balancing personalization and generalization.
• We propose a robust aggregation method based on the optimal transport theory, which can suppress the negative transfer caused by low-quality prototypes while maintaining the prototype geometry.
• Experiments on multiple datasets show that FedUP outperforms state-of-the-art methods under different degrees of data heterogeneity.

2 RelatedWork

2.1 Heterogeneous and Personalized FL

Statistical heterogeneity (Non-IID) poses a fundamental challenge to federated optimization [Makhija et al., 2024]. Regularization-based methods mitigate client drift by constraining local updates via proximal terms, such as FedProx [Li et al., 2020], or joint optimization objectives like Ditto [Li et al., 2021]. Architecture decoupling approaches support system heterogeneity by separating models into shared and private components. For instance, FedRep [Collins et al., 2021] and FedPer [Arivazhagan et al., 2019] maintain private classifiers, while FedBABU [Oh et al., 2021] freezes the classifier to update only the backbone. Although flexible, these methods are prone to overfitting private modules under data scarcity. Alternatively, Knowledge Distillation (KD) methods like FedMD [Li and Wang, 2019] and FedGen [Zhu et al., 2021] achieve model-agnostic transfer by aligning output distributions or synthesizing pseudo-data. However, KD-based approaches often incur high communication overhead and implicitly assume uniform reliability across client knowledge, failing to account for the varying quality of heterogeneous local data. Although recent generative approaches [Zhang et al., 2024b; Apell´aniz et al., 2025; Xiao et al., 2025] utilize synthetic data to mitigate data scarcity, they often require training heavy generator models, whereas our approach efficiently leverages probabilistic prototypes as lightweight generators.

2.2 Prototype-based Federated Learning

Prototype-based FL, represented by FedProto [Tan et al., 2022a], has emerged as a communication-efficient paradigm by exchanging class-level feature centroids. Subsequent works have enhanced prototype discriminability through contrastive learning (e.g., in FedProc [Mu et al., 2023] and FedPCL [Tan et al., 2022b]), gradient-based alignment such as FedPAC [Xu et al., 2023], or geometric adjustments like FedTGP [Zhang et al., 2024a] and FedGH [Yi et al., 2023]. These methods successfully reduce communication costs and support heterogeneous architectures by regularizing local features toward global representations [Guo et al., 2024].

Despite this progress, most prototype-based methods still implicitly assume equal reliability across clients. Under data heterogeneity, low-quality prototypes can mislead local alignment and contaminate global aggregation, resulting in negative transfer. Therefore, it remains important to develop a prototype-based FL framework that can quantify prototype uncertainty and perform reliability-aware aggregation.

3 Preliminaries

Consider a central server and K clients K = {1, . . . ,K}. Client i holds a private dataset Di = {(xij , yij )}ni j=1 drawn from a distinct distribution Pi(X, Y ). Due to privacy constraints, only model updates or knowledge summaries can be exchanged. The standard federated optimization objective minimizes the weighted empirical risk:

where n = ∑Ki=1 ni, fi(·; θi) denotes the client model and ℓ(·) is the classification loss.

To achieve architecture-agnostic knowledge sharing, existing methods adopt prototype-based representation. Let hi(·) be the feature extractor of client i. The deterministic prototype for class c is defined as µci = 1 |Dc i | ∑(x,y)∈Dc i hi(x), where Dc i = {(x, y) ∈ Di | y = c}. The server aggregates client prototypes via weighted averaging to obtain global prototypes, which are then broadcast to regularize local training.

However, under data heterogeneity, prototypes from datascarce clients may be highly unreliable yet appear equally credible during aggregation, leading to global knowledge drift. To address this limitation, we extend prototypes from points to probability distributions in Section 4.

4 Methodology

We present FedUP to address the two challenges by extending prototypes from deterministic points to probabilistic distributions. The core idea is to quantify prototype reliability via uncertainty. As shown in Figure 2, we employ Evidential Deep Learning to decompose predictive uncertainty into epistemic and aleatoric uncertainty. On the client side, epistemic uncertainty dynamically balances global alignment and local personalization, while global priors enrich sparse feature spaces to address Challenge #1. On the server side, prototypes are modeled as Gaussians and fused via uncertaintyweighted Wasserstein barycenters [Agueh and Carlier, 2011], preserving geometric structure while filtering unreliable contributions to address Challenge #2.

4.1 Local Training with UncertaintyModeling Evidential Uncertainty Quantification

Evidential Deep Learning (EDL) models the categorical prediction ρ as a random variable following a Dirichlet distribution Dir(ρ|α) [Malinin and Gales, 2018; Sensoy et al., 2018]. Unlike conventional networks, EDL learns concentration parameters α that enable explicit uncertainty decomposition through a single forward pass. For a given sample x and model parameters θ, the probability density function of ρ is defined as:

where α is the concentration parameter vector, Γ(·) is the Gamma function, and ∆C = {ρ ∈ RC | ∑Cc=1 ρc = 1, ρc > 0} denotes the C-dimensional unit simplex. The model prediction for class c, i.e., the expected categorical probability P(y = c|x, θ), is obtained by marginalizing over ρ:

where S = ∑Cc=1 αc is the total Dirichlet strength.

In practice, we append an evidential head after the feature extractor, which outputs non-negative evidence e(x) ∈ RC+ via softplus activation. The Dirichlet parameters are then computed as:

These Dirichlet parameters enable us to explicitly quantify the predictive uncertainty through a single forward pass

[Chen and Zhang, 2024]. The total predictive uncertainty can be decomposed into two components: aleatoric uncertainty and epistemic uncertainty. Aleatoric uncertainty Uale measures the inherent randomness in the data, which cannot be reduced by collecting more samples, reflecting factors such as class overlap or inherent noise. Epistemic uncertainty Uepi quantifies the model's lack of knowledge when facing unfamiliar data, which can be reduced as more training data is accumulated [Chen et al., 2024a; Chen et al., 2024b].

We first compute the total uncertainty as the entropy of the expected categorical distribution:

The aleatoric uncertainty captures the expected entropy under the Dirichlet distribution:

where ψ(·) is the digamma function. Finally, the epistemic uncertainty is obtained by subtracting aleatoric uncertainty from total uncertainty:

We train with the following EDL loss:

where LDir CE is the expected cross-entropy under the Dirichlet distribution, and the KL term regularizes the evidence of incorrect classes with ˜α = α - ey + 1.

Probabilistic Prototype Construction

We construct probabilistic prototypes by leveraging aleatoric uncertainty to down-weight noisy samples and to encode class-wise reliability. For client i and class c, let zj = fθ(xj) ∈ Rd denote the feature embedding for sample j ∈ Dc i . We assign a reliability weight wj = exp(-U (j) ale ) and compute the weighted prototype mean:

To model uncertainty in the prototype spread, we first compute the per-dimension empirical standard deviation σcfeat = Std({zj }j∈Dc i ) and then inflate it by class average aleatoric uncertainty:

Finally, we represent the class prototype as a diagonal Gaussian distribution:

where larger ¯U (i,c) ale yields a wider prototype distribution, suppressing unreliable or noisy classes in subsequent aggregation and alignment.

Dynamic Personalization via Epistemic Uncertainty

A fundamental tension in personalized FL is balancing global knowledge absorption against local pattern preservation. We resolve this by using epistemic uncertainty as an adaptive gate: samples where the model lacks knowledge should leverage global guidance, while confident predictions should retain local specificity. For each sample (x, y), we define a dynamic alignment coefficient:

where σ(·) is the sigmoid function, β > 0 controls sensitivity, and ¯Uepi is a running average providing normalization. This design ensures: high Uepi leads to γ → 1 (strong alignment with global prototype), while low Uepi leads to γ → 0 (preserve local representation).

Given the adaptive coefficient, we define the alignment loss as follows. Let z = fθ(x) ∈ Rd denote the feature embedding and P c g = N(µcg , Σcg ) the global probabilistic prototype for class c. To respect the distributional structure, we define the alignment loss using the Mahalanobis distance:

where δ = z - µyg denotes the residual vector. Under our diagonal covariance assumption, the inverse (Σyg )-1 = diag(1/(σy g,1 )2 , . . . , 1/(σy g,d )2 ) can be efficiently computed. Unlike naive Euclidean distance ∥z-µyg ∥2 , Mahalanobis distance automatically down-weights high global variance dimensions while emphasizing stable features.

Global Prior-Guided Feature Augmentation

The adaptive alignment mechanism is effective when local data are sufficient. In data-scarce clients, some classes contain too few samples to form reliable prototypes. Therefore, we use global probabilistic prototypes as generative priors. For each class c with insufficient local samples (|Dc i | < nmin), we draw virtual features from the global Gaussian:

We feed ˜zc into the local classifier head ϕi and optimize an auxiliary cross-entropy loss

This term encourages stable decision boundaries for sparse classes and promotes alignment with the global semantic.

The client objective is

Here, LEDL trains the evidential head, γ(x) scales alignment by epistemic uncertainty, and Laug provides priorguided augmentation. The hyperparameter λalign controls the strength of global prototype alignment.

4.2 Server-Side Prototype Aggregation

After local training, the server collects probabilistic prototypes from participating clients and aggregates them into global prototypes. The key challenge is to handle prototypes with varying reliability while preserving their distributional geometry. We address this through two steps: first, we assess prototype reliability using both epistemic and aleatoric uncertainty to compute aggregation weights; second, we aggregate prototypes via Wasserstein barycenters to avoid variance collapse inherent in Euclidean averaging.

ReliabilityAssessment viaUncertainty. Client prototypes vary in reliability depending on local data conditions. Clients with limited samples produce prototypes with high epistemic uncertainty, while clients with noisy or ambiguous data exhibit high aleatoric uncertainty. Traditional samplecount weighting cannot distinguish these cases. We leverage both uncertainty components for reliability assessment. Each client i reports class-wise uncertainties ¯U (i,c) epi and ¯U (i,c) ale alongside its local prototype P c i . The aggregation weight is then computed as:

where Sc denotes clients possessing class c, τ is a temperature parameter, and λunc balances the relative contributions of epistemic and aleatoric uncertainty in reliability estimation.

Wasserstein Barycenter Aggregation. Standard Euclidean averaging of Gaussian distributions suffers from variance collapse, where the aggregated covariance underestimates true class spread. To preserve distributional geometry, we aggregate on the 2-Wasserstein manifold, weighted by our uncertainty metrics. To aggregate the probabilistic prototypes {P c i }i∈Sc into a global prototype P c g , we minimize the weighted accumulation of squared 2-Wasserstein distances:

where the distance between two Gaussians is defined as:

W 2 2 (P1,P2) = ∥µ1 - µ2∥22 + tr(Σ1) + tr(Σ2) - 2tr((Σ1/2 1 Σ2Σ1/2 1 )1/2 ) [Li et al., 2024]. The solution to this barycenter problem yields the global mean and covariance updates. The barycentric mean is simply the weighted average µcg = ∑i∈Sc ωc i µci . The barycentric covariance Σcg is the fixed-point of the iteration Σt+1 = ∑i∈Sc ωc i (Σ1/2 t Σci Σ1/2 t )1/2 , initialized with the arithmetic mean.

5 Experiments

5.1 Experimental Setup

Datasets. We evaluate FedUP on three benchmarks covering one natural image dataset and two medical image datasets: CIFAR-10 (10 classes, 60K natural images), PathMNIST (9 tissue types, 107K medical histology images), and Kvasir-Capsule (14 gastrointestinal findings, 47K capsule endoscopy images).

Non-IID Data Partition. To simulate statistical heterogeneity in federated learning, we partition client data using a Dirichlet distribution Dir(α). For each class c, we sample a proportion vector qc ∼ Dir(α) and allocate samples to K clients accordingly. The hyperparameter α controls the degree of data heterogeneity: a smaller α implies stronger data skew and heterogeneity. In this experiment, we set α ∈ {0.1, 0.3, 0.5}, where α = 0.1 represents an extremely skewed heterogeneous scenario.

ComparedMethods. To comprehensively evaluate performance, we compare FedUP with 10 state-of-the-art (SOTA) methods categorized into four paradigms. Classical FL includes standard FedAvg [McMahan et al., 2017] and FedProx [Li et al., 2020], which introduces proximal regularization to handle heterogeneity. Personalized FL employs architecture decoupling or multi-task strategies, including FedRep [Collins et al., 2021], Ditto [Li et al., 2021], and FedALA

[Zhang et al., 2023]. Prototype-based FL utilizes feature centroids for knowledge transfer, including FedProto [Tan et al., 2022a], FedPAC [Xu et al., 2023], FedGH [Yi et al., 2023], and FedTGP [Zhang et al., 2024a]. Geometric & Bayesian FL is represented by FedWBA [Wei et al., 2025], which utilizes Wasserstein Barycenter for robust aggregation.

Implementation Details. Regarding model architecture, to ensure fair comparison, we employ ResNet-18 as the backbone for all datasets. For the training configuration, we simulate a federated system with K = 50 clients. We set the participation rate to 10% per round and train for a total of T = 300 communication rounds. For hyperparameters, local training uses an SGD optimizer with a momentum of 0.9, weight decay of 1e - 4, batch size B = 10, and local epochs E = 5. The learning rate is initialized at 0.01 and decays using a cosine annealing strategy. Regarding FedUP-specific settings, the evidential Dirichlet prior strength S is initialized to 10. The uncertainty-guided personalization coefficient β is set to 5.0, and the temperature coefficient τ for calculating aggregation weights is set to 0.5.

5.2 Main Results

Model Accuracy. Table 1 reports test accuracy across three datasets under varying non-IID heterogeneity (α ∈ {0.1, 0.3, 0.5}). FedUP consistently achieves the best performance across all settings. Notably, on PathMNIST, FedUP significantly outperforms all baselines, achieving 93.29%, 87.04%, and 85.38% under α = 0.1, 0.3, and 0.5 respectively. On CIFAR-10, FedUP achieves 90.12% at α = 0.1, comparable to FedALA while outperforming prototype-based methods such as FedProto and FedGH by over 3%. On Kvasir-Capsule, FedUP reaches 95.94% at α = 0.1 and 88.38% at α = 0.5, where the latter exceeds the second-best method by 5.52%, demonstrating stronger robustness under milder heterogeneity. These results validate that uncertaintyguided local adaptation and geometry-preserving global aggregation can effectively address the challenges posed by data heterogeneity in federated learning.

Convergence Performance. Figure 3 illustrates FedUP's superior convergence on PathMNIST. FedUP reaches 80% accuracy within just 20 rounds, approximately 2.5 times as fast as FedProto and FedPAC which require over 50 rounds. More importantly, baseline methods such as FedAvg, FedProto, and FedPAC exhibit significant oscillations under heterogeneous settings, with accuracy fluctuating by up to 1015% between rounds. In contrast, FedUP maintains a smooth and stable trajectory throughout training. This stability stems from uncertainty-aware adaptation that suppresses unreliable updates and Wasserstein aggregation that prevents global prototype drift.

5.3 Robustness Analysis

Impact of Model Heterogeneity. To evaluate robustness against model heterogeneity, we test FedUP under four scenarios where clients randomly adopt a fixed architecture from a candidate pool: HtFE2 (4-layer CNN, ResNet18); HtFE3 (ResNet10/18/34); HtFE4 (CNN, GoogleNet, MobileNetV2, ResNet18); and HtFE9 (9 variants from ResNet4 to ResNet152). As shown in Table 2 (PathMNIST, α = 0.1), baseline performance deteriorates significantly with increased heterogeneity; notably, FedProto suffers a drastic 21% drop from HtFE2 to HtFE9. In contrast, FedUP maintains remarkable stability and achieves 73.21% accuracy in the extreme HtFE9 setting, outperforming the best baseline FedTGP by 1.66%. This validates that our Wasserstein-based probabilistic aggregation aligns heterogeneous feature spaces more effectively than traditional point-based estimation.

Impact of Client Training Epochs. We investigate the robustness of FedUP against client drift by varying local epochs E ∈ {5, 10, 20} on PathMNIST (α = 0.1). As shown in Table 3, baselines like FedGH and FedWBA suffer performance degradation as E increases (e.g., FedWBA drops by 1.61%) due to overfitting local heterogeneous data. In contrast, FedUP benefits from prolonged local training, achieving peak accuracy of 91.40% at E = 20. This confirms that our uncertainty-guided regularization effectively constrains local updates within a reliable range, preventing severe drift even with extensive local computation.

5.4 Ablation Study

Table 4 presents the ablation study on PathMNIST (α = 0.5) with three modules: Feature Augmentation (FA), Uncertainty Weighting (UW), and Wasserstein Aggregation (WA). Among individual components, UW contributes most significantly (74.12% at R=300), validating that uncertainty-guided alignment effectively prevents unreliable updates from corrupting local models, while WA alone (71.45%) outperforms FA alone (69.23%), indicating that geometry-preserving aggregation provides more immediate benefits when prototype quality varies across clients. Combining any two modules yields substantial improvements, where UW+WA achieves 81.85%, demonstrating that client-side uncertainty filtering and server-side geometric aggregation are complementary. The complete FedUP framework achieves 85.38%, outperforming the best two-module variant by 3.53%, confirming that FA mitigates data sparsity, UW enables adaptive personalization, and WA ensures robust knowledge fusion.

5.5 Hyperparameter Sensitivity

Uncertainty-Guided Personalization Coefficient β. β balances local feature preservation and global alignment. As shown in Figure 4(a), performance improves as β increases until β = 5, beyond which excessive regularization suppresses valid local patterns. This confirms β = 5 as the optimal trade-off across all heterogeneity levels.

Temperature Coefficient τ. We then analyze the temperature coefficient τ for server-side aggregation weighting. As shown in Figure 4(b), small τ over-concentrates weights on a few clients, reducing diversity, while large τ yields nearuniform weights that fail to filter unreliable prototypes. Setting τ = 0.5 achieves the best balance between reliability filtering and knowledge diversity.

5.6 Visualization of Uncertainty Evolution

Figure 5 visualizes the Dirichlet distribution evolution on the simplex for hard examples on PathMNIST. In early training, the probability mass is dispersed with low evidence strength, indicating high epistemic uncertainty. As training progresses, FedUP strengthens global constraints on uncertain samples and enriches sparse representations through prior-guided augmentation. The distribution gradually concentrates toward the correct class vertex with increasing evidence, suggesting better-calibrated predictions rather than overconfident errors.

6 Conclusion

We propose FedUP, an uncertainty-aware personalized federated learning framework that models class prototypes as probability distributions. By decomposing predictive uncertainty, FedUP adaptively balances local personalization and global alignment, while suppressing unreliable prototype contributions during aggregation. Experiments on natural and medical benchmarks demonstrate the effectiveness and robustness of FedUP under data heterogeneity.

Acknowledgments

This work is supported, in part, by the National Natural Science Foundation of China under Grant Nos. 92567204 and 62072469; the Shandong Provincial Natural Science Foundation Major Basic Research Program under Grant No. ZR2024ZD20; the Shandong Data Open Innovative Application Laboratory.