Joint Local Training and Device Scheduling for Heterogeneous Federated Learning

Federated learning (FL) is a distributed machine learning paradigm that allows training a global model based on data distributed across multiple devices. However, local training on each device incurs a high energy cost. Moreover, due to the heterogeneity of computation capability, optimizing the energy efficiency of FL requires careful consideration of device selection and the number of local epochs for each device in each training round. Existing works often optimize a common number of local epochs for all devices, which may not be suitable for heterogeneous systems. In this paper, we propose a joint optimization of local training and device scheduling to balance global model accuracy and overall energy consumption given a specific training time. We formulate this optimization as a Markov decision process and propose a deep reinforcement learning (DRL)-based approach to select participating devices and determine the number of training epochs for each device in each training round. Experimental results show that the proposed approach improves model accuracy and reduces energy consumption compared to baseline schemes.