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

Eman M. AbouNassar; Amr Elshall; Sameh Abdulah

cs.LGMar 14, 2026

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

Authors: Eman M. AbouNassar, Amr Elshall, Sameh Abdulah

Organizations: 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.

Abstract

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.

Figures & tables

Figure 1: The Federated Learning process: 1) the server distributes the global model to all clients; 2) each client updates the model using its local data; 3) clients send their local updates back to the server; and 4) the server aggregates these updates to form the new global model.

Figure 2: The FedPBS proposed Model: 1) the server distributes the aggregation model to all clients; 2) each client estimates the variance and batch value then updates the model; 3) the fedprox-term added for all clients with small batches, and high variance then the model is updated; 4) clients send local updates back to the server; and 4) the server aggregates these updates for the new global model.

Notation	Description
$Q$	Total number of participating clients.
$S_{t}$	Selected subset of clients in communication round $t$ .
$w^{(t)}$	Global model parameters at round $t$ .
$w_{q}^{(t)}$	Local model of client $q$ before update in round $t$ .
$w_{q}^{(t+1)}$	Locally updated model of client $q$ sent to the server.
$D_{q}$	Local dataset of client $q$ .

Table 1: Notation summary of the proposed FedPBS framework.

Parameter (Abbreviation)	Default Value
Number of clients (Q)	10
Communication rounds (T)	100
Learning rate ( $\eta$ )	1e-3
Batch size (B)	64
Dirichlet concentration ( $\alpha$ )	0.2
Local epochs (E)	20

Table 2: Default experimental configuration of our study based on UCI-HAR.

Panel of Figure 3 — Figure 3: Client class-distribution matrices for different Dirichlet concentration parameters $\alpha$ . Smaller $\alpha$ implies stronger label skew and higher heterogeneity.

Figure 4: Performance comparison of federated learning baselines against FedPBS under varying non-IID levels on the UCI-HAR dataset.

Method	$\alpha$ =0.2	$\alpha$ =0.5	$\alpha$ =1
FedPBS	0.9403	0.9230	0.9092
FedBS	0.8700	0.8588	0.8273
FedProx	0.9175	0.8991	0.8986
FedGA	0.9077	0.9189	0.9151
Moon	0.9190	0.9161	0.9091

Table 3: Performance comparison of existing algorithms on the UCI-HAR dataset under varying Dirichlet non-IID settings (

\alpha=0.2,0.5,1.0

Figure 5: Performance comparison of federated learning baselines against FedPBS under varying non-IID levels on the Cifar10 dataset.

Method	$\alpha$ =0.2	$\alpha$ =0.5	$\alpha$ =1
FedPBS	0.6373	0.6339	0.63613
FedBS	0.6313	0.63411	0.6331
FedProx	0.63172	0.5217	0.5712
FedGA	0.5144	0.5501	0.63011
Moon	0.6324	0.6314	0.63611

Table 4: Performance comparison of existing algorithms on the Cifar10 dataset under varying Dirichlet non-IID settings (

\alpha=0.2,0.5,1.0

Figure 6: Graph with (y-axis) loss behaviour through the optimization process (UCI).

Figure 7: Graph with (y-axis) loss behaviour through the optimization process (Cifar10).

Federated Learning

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

See all 943 papers

FedPrism: Adaptive Personalized Federated Learning under Non-IID Data
Mar 9, 2026Prakash Kumbhakar, Shrey Srivastava, Haroon R LoneFederated Learning Non-IID Federated Learning
IISER Bhopal, Bhopal, India
Multi-Level Analyzation of Imbalance to Resolve Non-IID-Ness in Federated Learning
Jun 8, 2026Haengbok Chung, Jae Sung LeeClass-Imbalanced Learning Federated Learning Aggregation
Interdisciplinary Program in Artificial Intelligence, Seoul National University, Republic of Korea · Department of Nuclear Medicine, Seoul National University College of Medicine, Republic of Korea · Brightonix Imaging Inc. Seoul, Republic of Korea
Client and Training Data Selection for Computationally Efficient Synchronized Federated Learning
Sep 30, 2026Muzaffer Citir, Hiroki Nishikawa, Sangyoung ParkClient Selection in Federated Learning Non-IID Federated Learning
Smart Mobility Systems, Technical University of Berlin, Germany · Graduate School of Information Science and Technology, The University of Osaka, Japan

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

FedPrism: Adaptive Personalized Federated Learning under Non-IID Data

Prakash Kumbhakar, Shrey Srivastava, Haroon R Lone

IISER Bhopal, Bhopal, India

Federated Learning Non-IID Federated Learning Heterogeneous Federated Learning Clustered Federated Learning

Open paper AskPDF arXiv

Federated Learning (FL) suffers significant performance degradation in real-world deployments characterized by moderate to extreme statistical heterogeneity (non-IID client data). While global aggregation strategies promote broad generalization, they often fail to capture the diversity of local data distributions, leading to suboptimal personalization. We address this problem with FedPrism, a framework that uses two main strategies. First, it uses a Prism Decomposition method that builds each client's model from three parts: a global foundation, a shared group part for similar clients, and a private part for unique local data. This allows the system to group similar users together automatically and adapt if their data changes. Second, we include a Dual-Stream design that runs a general model alongside a local specialist. The system routes predictions between the general model and the local specialist based on the specialist's confidence. Through systematic experiments on non-IID data partitions, we demonstrate that FedPrism exceeds static aggregation and hard-clustering baselines, achieving significant accuracy gains under high heterogeneity. These results establish FedPrism as a robust and flexible solution for federated learning in heterogeneous environments, effectively balancing generalizable knowledge with adaptive personalization.

Abstract

Federated Learning

Explore similar work

FedPrism: Adaptive Personalized Federated Learning under Non-IID Data

Multi-Level Analyzation of Imbalance to Resolve Non-IID-Ness in Federated Learning

Client and Training Data Selection for Computationally Efficient Synchronized Federated Learning