Q-Learning-Aided Offloading Strategy in Edge-Assisted Federated Learning over Industrial IoT

Abstract

Federated learning (FL) is a key solution to realizing a cost-efficient and intelligent Industrial Internet of Things (IIoT). To improve training efficiency and mitigate the straggler effect of FL, this paper investigates an edge-assisted FL framework over an IIoT system by combining it with a mobile edge computing (MEC) technique. In the proposed edge-assisted FL framework, each IIoT device with weak computation capacity can offload partial local data to an edge server with strong computing power for edge training. In order to obtain the optimal offloading strategy, we formulate an FL loss function minimization problem under the latency constraint in the proposed edge-assisted FL framework by optimizing the offloading data size of each device. An optimal offloading strategy is first derived in a perfect channel state information (CSI) scenario. Then, we extend the strategy into an imperfect CSI scenario and accordingly propose a Q-learning-aided offloading strategy. Finally, our simulation results show that our proposed Q-learning-based offloading strategy can improve FL test accuracy by about 4.7% compared to the conventional FL scheme. Furthermore, the proposed Q-learning-based offloading strategy can achieve similar performance to the optimal offloading strategy and always outperforms the conventional FL scheme in different system parameters, which validates the effectiveness of the proposed edge-assisted framework and Q-learning-based offloading strategy.