Probabilistic Federated Learning on Uncertain and Heterogeneous Data with Model Personalization

Ratun Rahman; Dinh C. Nguyen

cs.LGMar 18, 2026

Probabilistic Federated Learning on Uncertain and Heterogeneous Data with Model Personalization

Authors: Ratun Rahman, Dinh C. Nguyen

Organizations: Department of Electrical and Computer Engineering, University of Alabama in Huntsville, Huntsville, AL 35899, USA

Abstract

Conventional federated learning (FL) frameworks often suffer from training degradation due to data uncertainty and heterogeneity across local clients. Probabilistic approaches such as Bayesian neural networks (BNNs) can mitigate this issue by explicitly modeling uncertainty, but they introduce additional runtime, latency, and bandwidth overhead that has rarely been studied in federated settings. To address these challenges, we propose Meta-BayFL, a personalized probabilistic FL method that combines meta-learning with BNNs to improve training under uncertain and heterogeneous data. The framework is characterized by three main features: (1) BNN-based client models incorporate uncertainty across hidden layers to stabilize training on small and noisy datasets, (2) meta-learning with adaptive learning rates enables personalized updates that enhance local training under non-IID conditions, and (3) a unified probabilistic and personalized design improves the robustness of global model aggregation. We provide a theoretical convergence analysis and characterize the upper bound of the global model over communication rounds. In addition, we evaluate computational costs (runtime, latency, and communication) and discuss the feasibility of deployment on resource-constrained devices such as edge nodes and IoT systems. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny-ImageNet show that Meta-BayFL consistently outperforms state-of-the-art methods, including both standard and personalized FL approaches (e.g., pFedMe, Ditto, FedFomo), with up to 7.42% higher test accuracy.

Figures & tables

Fig. 1 : A simple probabilistic model structure in which probabilistic properties link every layer to every layer after it. The likelihood graph shows the likelihood of moving from the current layer to the next. It shows that probabilistic (distributional) parameters spread uncertainty through depth.

Fig. 2 : Our proposed Meta-BayFL FL algorithm where there are $K$ clients connected to the global server with adaptive model aggregation using a personalized learning approach. Each client creates local models using BNN and their local data and different temporary learning rates, and chooses the best rates for local model training. The server aggregates the local models and distributes the updated global model to all the clients.

Non-i.i.d. Type	$\mathcal{U}$	$\|\mathcal{U}\|$	ConvNet	ResNet20	ResNet32	ResNet44	ResNet56
Step	CIFAR-10	10K	65.5 $\pm$ 0.52	67.5 $\pm$ 0.52	62.7 $\pm$ 0.47	55.5 $\pm$ 0.42	57.0 $\pm$ 0.55
	CIFAR-100	50K	65.4 $\pm$ 0.48	68.2 $\pm$ 0.48	62.3 $\pm$ 0.33	55.1 $\pm$ 0.41	51.0 $\pm$ 0.59
	Tiny-ImageNet	100K	65.5 $\pm$ 0.60	67.1 $\pm$ 0.61	62.3 $\pm$ 0.53	54.5 $\pm$ 0.41	59.0 $\pm$ 0.42
Dirichlet	CIFAR-10	10K	63.9 $\pm$ 0.55	68.2 $\pm$ 0.46	67.7 $\pm$ 0.55	57.0 $\pm$ 0.48	57.0 $\pm$ 0.42
	CIFAR-100	50K	63.5 $\pm$ 0.51	68.6 $\pm$ 0.53	66.5 $\pm$ 0.41	62.0 $\pm$ 0.27	56.9 $\pm$ 0.47
	Tiny-ImageNet	100K	64.0 $\pm$ 0.55	68.2 $\pm$ 0.52	56.0 $\pm$ 0.62	61.6 $\pm$ 0.76	57.3 $\pm$ 0.42

Table I : Meta-BayFL on non-i.i.d CIFAR-10 with different unlabeled data

U

. rows correspond to non-IID types (Step/Dirichlet) and choices of unlabeled data source and size; entries report accuracy (mean

\pm

std) for each backbone. Out-of-domain unlabeled data (e.g., CIFAR-100/Tiny-ImageNet) is sufficient to maintain or improve performance, supporting the practicality of the approach.

Data		CIFAR-10		CIFAR-100		Tiny-ImageNet
Data		FedAVG	BayFL	FedAVG	BayFL	FedAVG	BayFL
$\|D\|$ = 100%	$\epsilon$ = 0.0	66.86%	67.20%	71.25%	71.23%	73.11%	73.59%
	$\epsilon$ = 0.0001	66.71%	66.81%	71.22%	71.20%	73.05%	73.50%
	$\epsilon$ = 0.001	65.69%	66.69%	69.50%	71.05%	72.20%	73.44%
	$\epsilon$ = 0.01	63.10%	66.52%	67.98%	70.60%	70.66%	73.20%
	$\epsilon$ = 0.1	59.23%	65.65%	64.40%	70.00%	67.84%	72.19%

Table II : Difference in performance with different data size

|D|

and noise level

\epsilon

across various datasets for 5 non-IID Dirichlet clients. Each block fixes

|D|

and sweeps

\epsilon

; values compare FedAVG vs. BayFL on each dataset. The BayFL advantage widens as data becomes smaller and noisier, showing stronger robustness in challenging federated environments.

Fig. 3 : Performance comparison on the CIFAR-10 non-IID dataset between basic CNN, FedAVG (FL with CNN), and BayFL (FL with BNN) approaches. A total of 5 clients with various data distributions and batch sizes are used for FL approaches with 10 local epochs in every global round.

Variables		CIFAR-10	CIFAR-100	Tiny-ImageNet
$\epsilon$ 0.001	BayFL, $lr$ = 0.0001	66.69%	71.05%	73.44%
	BayFL, $lr$ = 0.001	66.90%	71.66%	73.75%
	BayFL, $lr$ = 0.01	66.02%	70.76%	73.41%
	Meta-BayFL	72.25 %	73.06 %	77.00 %
$\epsilon$ 0.01	BayFL, $lr$ = 0.0001	65.40%	70.05%	71.34%
	BayFL, $lr$ = 0.001	65.60%	70.33%	71.30%

Table III : Difference in performance with different noise level

\epsilon

and BayFL with different learning rates

\mathrm{lr}

across various datasets for 5 non-IID Dirichlet clients. For each

\epsilon

, we compare fixed learning rates against Meta-BayFL. Meta-learning is consistently best and degrades more gracefully under increasing noise.

Environment	Method	Acc (%)	Time/Round (s)
Simulated	FedAvg / FedProx	66.85	8.60
	Meta-BayFL (MC Dropout)	67.10	8.72
	Meta-BayFL (BNN)	69.48	11.00
Raspberry Pi	FedAvg / FedProx	66.10	27.8
	Meta-BayFL (MC Dropout)	66.40	29.4
	Meta-BayFL (BNN)	68.90	37.6

Table IV : Ablation study comparing uncertainty approximation methods under non-IID CIFAR-10 (Dirichlet

\alpha=0.5

, 5 clients, ResNet-56 backbone , same training budget). Results are reported for both simulated execution and real Raspberry Pi (CPU-only) deployment. Runtime denotes average client-side training time per communication round.

Clients	CIFAR-10	CIFAR-100	Tiny-ImageNet
5	72.25%	73.06%	77.00%
10	73.10%	74.25%	78.02%
15	73.55%	74.80%	78.40%
20	74.02%	75.10%	78.85%

Table V : Performance of Meta-BayFL under different numbers of non-IID Dirichlet clients with noise level

\epsilon=0.001

. Each row changes client count while keeping the noise fixed; entries are final test accuracy. Meta-BayFL scales positively with more clients (more diversity), with the expected trade-off of higher system overhead discussed in Section V.

Method	CIFAR-10		CIFAR-100		Tiny-ImageNet
Method	Loss Value	Accuracy	Loss Value	Accuracy	Loss Value	Accuracy
FedAVG [ 29 ]	4.35	45.50%	4.42	37.95%	4.43	36.55%
BNN [ 12 ]	2.53	52.22%	2.89	48.48%	2.75	50.18%
pFedBayes [ 21 ]	1.94	64.27%	2.33	65.31%	2.10	63.51%
pFedBe [ 22 ]	1.92	64.86%	1.99	64.03%	2.13	63.11%
Fedmask [ 20 ]	1.67	66.25%	1.88	65.50%	1.85	66.33%

Table VI : Comparison with state-of-the-art approaches.

Federated Learning

Federated learning trains shared models across data holders without requiring them to pool their raw data.

See all 943 papers

FedPBS: Proximal-Balanced Scaling Federated Learning Model for Robust Personalized Training for Non-IID Data
Mar 14, 2026Eman M. AbouNassar, Amr Elshall, Sameh AbdulahFederated Learning Non-IID Federated Learning
Mathematics and Computer Science Department, Faculty of Science, Menoufia University, Shebin El Koom, Egypt. · Computer Science Department, Faculty of Computers and Information, Menoufia University, Shebin El Koom, Egypt.
Personalized Additive Modeling for Multi-level Federated Learning
May 26, 2024Shutong Chen, Guodong Long, Tianyi Zhou +3Non-IID Federated Learning Federated Learning
University of Technology Sydney, Australia · Mohamed bin Zayed University of Artificial Intelligence, UAE · The Hong Kong Polytechnic University, China
FedPLT: Scalable, Resource-Efficient, and Heterogeneity-Aware Federated Learning via Partial Layer Training
May 4, 2026Ahmad Dabaja, Rachid El-AzouziCommunication-Efficient Distributed Training Heterogeneous Federated Learning
Laboratoire Informatique d’Avignon (LIA), Avignon University, France

Mar 14, 2026J/K move · Enter open · S save

FedPBS: Proximal-Balanced Scaling Federated Learning Model for Robust Personalized Training for Non-IID Data

Eman M. AbouNassar, Amr Elshall, Sameh Abdulah

Mathematics and Computer Science Department, Faculty of Science, Menoufia University, Shebin El Koom, Egypt. · Computer Science Department, Faculty of Computers and Information, Menoufia University, Shebin El Koom, Egypt.

Federated Learning Non-IID Federated Learning Heterogeneous Federated Learning

Open paper AskPDF arXiv

Federated learning (FL) enables a set of distributed clients to jointly train machine learning models while preserving their local data privacy, making it attractive for applications in healthcare, finance, mobility, and smart-city systems. However, FL faces several challenges, including statistical heterogeneity and uneven client participation, which can degrade convergence and model quality. In this work, we propose FedPBS, an FL algorithm that couples complementary ideas from FedBS and FedProx to address these challenges. FedPBS dynamically adapts batch sizes to client resources to support balanced and scalable participation, and selectively applies a proximal correction to small-batch clients to stabilize local updates and reduce divergence from the global model. Experiments on benchmarking datasets such as CIFAR-10 and UCI-HAR under highly non-IID settings demonstrate that FedPBS consistently outperforms state-of-the-art methods, including FedBS, FedGA, MOON, and FedProx. The results demonstrate robust performance gains under extreme data heterogeneity, with smooth loss curves indicating stable convergence across diverse federated environments. FedPBS consistently outperforms state-of-the-art federated learning baselines on UCI-HAR and CIFAR-10 under severe non-IID conditions while maintaining stable and reliable convergence.

Abstract

Federated Learning

Explore similar work

FedPBS: Proximal-Balanced Scaling Federated Learning Model for Robust Personalized Training for Non-IID Data

Personalized Additive Modeling for Multi-level Federated Learning

FedPLT: Scalable, Resource-Efficient, and Heterogeneity-Aware Federated Learning via Partial Layer Training