Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Eman M. AbouNassara, Amr Elshalla, Sameh Abdulah

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.

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

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.
Paper Structure (17 sections, 18 equations, 7 figures, 4 tables, 4 algorithms)

This paper contains 17 sections, 18 equations, 7 figures, 4 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Federated Learning Under Statistical Heterogeneity
  4. Foundations of Federated Learning
  5. Statistical Heterogeneity in FL
  6. FedProx: A Proximal Approach
  7. FedProx Formulation
  8. FedBS: Balancing Subtasks for Scalability
  9. FedBS Formulation
  10. FedPBS: Proximal-Balanced FL
  11. Experimental Results
  12. Datasets and Non-IID Data Construction
  13. Client Performance Under Non-IID data
  14. Accuracy Comparison with Baseline Algorithms
  15. Convergence Behavior: Loss versus Training Rounds
  16. ...and 2 more sections

Figures (7)

  • 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.
  • 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.
  • Figure 5: Performance comparison of federated learning baselines against FedPBS under varying non-IID levels on the Cifar10 dataset.
  • ...and 2 more figures