DRL-empowered joint batch size and weighted aggregation adjustment mechanism for federated learning on non-IID data

Abstract

To address the accuracy degradation as well as prolonged convergence time due to the inherent data heterogeneity among end-devices in federated learning (FL), we introduce the joint batch size and weighted aggregation adjustment problem, which is non-convex problem. To adjust optimal hyperparameters, we develop deep reinforcement learning (DRL) to empower a mechanism known as Batch size and Weighted aggregation Adjustment (BWA). Experimental evaluation demonstrates that BWA not only outperforms methods optimized solely from either a local training or server perspective but also achieves higher accuracy, with an increase of up to 5.53% compared to FedAvg, and additionally accelerates convergence speeds.