Towards Dynamic Resource Allocation and Client Scheduling in Hierarchical Federated Learning: A Two-Phase Deep Reinforcement Learning Approach
Xiaojing Chen, Zhenyuan Li, Wei Ni, Xin Wang, Shunqing Zhang, Yanzan Sun, Shugong Xu, Qingqi Pei
TL;DR
This work tackles dynamic resource allocation and client scheduling in energy-harvesting hierarchical federated learning (HFL). It proposes a two-phase deep deterministic policy gradient framework (TP-DDPG) that splits decisions into two groups: phase 1 uses a DDPG actor-critic to select participating clients and configure their CPU frequencies and transmit powers, while phase 2 uses the Straggler-Aware Client Association and Bandwidth Allocation (SCABA) to optimize client associations and bandwidth, providing rewards to train the DDPG. The objective is $U=\sum_{t=1}^{RR_1} O_t - RT_g$ with $O_t=\lambda|\Omega^t|-T^t$, balancing learning delay and accuracy; the approach reduces the learnable action space and improves convergence. Experimental results on MNIST and CIFAR-10 show TP-DDPG achieves up to 39.4% reduction in training time to reach 0.9 accuracy on MNIST and attains 0.93 accuracy on CIFAR-10, outperforming baselines. This enables scalable, energy-aware HFL at the network edge with real-time adaptation to wireless channels and energy arrivals.
Abstract
Federated learning (FL) is a viable technique to train a shared machine learning model without sharing data. Hierarchical FL (HFL) system has yet to be studied regrading its multiple levels of energy, computation, communication, and client scheduling, especially when it comes to clients relying on energy harvesting to power their operations. This paper presents a new two-phase deep deterministic policy gradient (DDPG) framework, referred to as ``TP-DDPG'', to balance online the learning delay and model accuracy of an FL process in an energy harvesting-powered HFL system. The key idea is that we divide optimization decisions into two groups, and employ DDPG to learn one group in the first phase, while interpreting the other group as part of the environment to provide rewards for training the DDPG in the second phase. Specifically, the DDPG learns the selection of participating clients, and their CPU configurations and the transmission powers. A new straggler-aware client association and bandwidth allocation (SCABA) algorithm efficiently optimizes the other decisions and evaluates the reward for the DDPG. Experiments demonstrate that with substantially reduced number of learnable parameters, the TP-DDPG can quickly converge to effective polices that can shorten the training time of HFL by 39.4% compared to its benchmarks, when the required test accuracy of HFL is 0.9.
