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.
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.
St
Selected subset of clients in communication round t .
w(t)
Global model parameters at round t .
wq(t)
Local model of client q before update in round t .
wq(t+1)
Locally updated model of client q sent to the server.
Dq
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 ( η )
1e-3
Batch size (B)
64
Dirichlet concentration ( α )
0.2
Local epochs (E)
20
Table 2: Default experimental configuration of our study based on UCI-HAR.
Figure 3: Client class-distribution matrices for different Dirichlet concentration parameters α . Smaller α 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
α =0.2
α =0.5
α =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 ( α=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
α =0.2
α =0.5
α =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 ( α=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 (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.
Prakash Kumbhakar, Shrey Srivastava, Haroon R Lone
Class imbalance is a common problem in deep learning that severely degrades performance. In federated learning (FL), it is a critical factor contributing to non-identically distributed data (non-IID). Building on several previous attempts, we define and analyze imbalance issues in FL at three levels: inter-case, inter-class, and inter-client. Inter-case imbalance addresses the imbalance in every single class; inter-class imbalance compares the number of data between different classes. Inter-client imbalance represents different skewness of local data between clients. Based on these concepts, we propose FedBB, which consists of two main components: (1) Positive Negative Balanced (PNB) loss function addresses the inter-case and inter-class imbalances in local training, enhancing generalization on highly skewed local client datasets. It optimizes both multi-label and multi-class classifications by assigning higher weights to minority cases or classes. (2) Client Balanced Reweighting (CBR) reweights clients based on inter-client imbalance during model aggregation, giving greater weight to models trained on less skewed datasets. Various experiments on X-ray and natural image datasets demonstrate that FedBB outperforms other algorithms in both performance and efficiency. Additionally, it requires limited statistical information, which is beneficial for privacy protection. Through ablation studies, we proved that PNB loss and CBR independently contribute to performance. As FedBB aims to build a global model that accurately classifies all classes, it can serve as a baseline for the generic and personalized FL.
Haengbok Chung, Jae Sung Lee
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
Federated learning (FL) is a promising paradigm of machine learning, which preserves user privacy by enabling learning without sharing raw data with a cloud server. Straggling clients have been a problem for FL as they introduce delays in aggregating the local models and hence, the convergence of the global model. Therefore, it is important to have a mechanism that ensures fast convergence of the global model as well as good FL participation rate. Another issue for the convergence of a model in FL is the non-independent and identically distributed (non-iid) data across the clients. Prior approaches based on probabilistic client selection do not work well under non-iid data especially when the number of clients is small. We show scenarios where such approaches fail and propose a joint client-training data selection algorithm for fast convergence of FL models. Our experiments on CIFAR-100 dataset show that convergence of the FL model can be significantly improved over prior works that can consider non-iid data and heterogeneous computation and higher model accuracy.
Muzaffer Citir, Hiroki Nishikawa, Sangyoung Park
Smart Mobility Systems, Technical University of Berlin, Germany · Graduate School of Information Science and Technology, The University of Osaka, Japan