DRL-Enabled Hierarchical Federated Learning Optimization for Data Heterogeneity Management in Multi-Access Edge Computing

Abstract

This paper designs a novel Hierarchical Federated Learning (HFL) management scheme, enabled by deep reinforcement learning (DRL), for multi-access edge computing (MEC) environments to accelerate convergence. To do this, the proposed scheme controls the number of local updates performed by mobile devices (MDs) and the number of intermediate aggregations at the MEC server before data is transmitted to the cloud for global aggregation. This optimization aims to i) mitigate the straggler effect by balancing training times between MEC and cloud servers, and ii) reduce the risk of overfitting by avoiding excessive reliance on faster MDs. Additionally, to improve the efficiency of DRL, Bayesian optimization is employed to initialize action values, thereby avoiding inefficient exploration of actions. Extensive simulations demonstrate that our proposed scheme outperforms various benchmarks in terms of test accuracy and training time.