Optimization Stability in Federated Learning under Non-IID Data: A Comparative Study of FedAvg and FedProx with Deep Convolutional Neural Network Architectures

Gitanjali Bhimrao Yadav, Jayashree Bagade

Federated learning (FL) has emerged as a prospective model training paradigm for distributed clients, which provides data privacy while maintaining the ability to train a collaborative model.In healthcare, finance and edge intelligence applications, federated learning (FL) is a promising model training paradigm for distributed clients, which preserve data privacy while retaining the capacity of collaboratively training a model. But, this optimization when data is not independent and identically distributed (non-IID), is still a great challenge because of the heterogeneity of clients and unstable model convergence. In this paper, we provide a thorough comparison between two well-known federated optimization algorithms, FedAvg and FedProx, and three popular deep convolutional neural network (DCNN) architectures such as ResNet18, VGG16 and VGG19. This study comprehensively assesses these models in a 10-client federated setting where data is partitioned into class-skewed non-IID data sets, by examining optimized stability, convergence and classification performance across 20 communication rounds. The experimental setup investigates how the proximal regularization can reduce the client drift and enhance the learning uniformity of the client datasets. The experimental results show that FedProx consistently shows better convergence stability and better predictive performance than FedAvg as the data becomes more and more heterogeneous, whereas ResNet18 provides a good trade-off between the model complexity, communication efficiency and classification performance. The results demonstrate that enforcing strong federated optimization coupled with fitting the appropriate deep convolutional architectures could provide a more reliable way of learning in decentralized settings. This work offers practical learnings for designing scalable, privacy-preserving and stable FL systems in real-world non-IID applications.