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

TL;DR

FedPBS is a novel federated learning algorithm that dynamically balances client participation and stabilizes local updates, significantly improving model performance and convergence stability under highly non-IID data distributions.

Contribution

The paper introduces FedPBS, which combines adaptive batch sizing and proximal correction to enhance robustness and scalability in federated learning with non-IID data.

Findings

Outperforms state-of-the-art methods like FedBS, FedGA, MOON, and FedProx.

Achieves stable convergence with smooth loss curves.

Demonstrates significant performance gains under extreme data heterogeneity.

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.