Relay-Assisted Over-the-Air Federated Learning

Zehong Lin, Hang Liu, Ying-Jun Angela Zhang

Abstract: Federated learning (FL) has recently emerged as a promising technology to enable artificial intelligence (AI) at the network edge. To improve the communication efficiency of FL, over-the-air computation allows a large number of mobile devices to concurrently upload their local models. Due to wireless channel fading, the model aggregation error at the edge server is dominated by the weakest channel among all devices, causing severe straggler issues. In this paper, we propose a relay-assisted over-the-air FL scheme to address the straggler issue. In particular, we adopt a half-duplex relay to assist the devices in uploading the local model updates to the edge server. Our scheme exploits the direct transmissions by the devices and the device-relay cooperative diversity for over-the-air model aggregation. Then, we study the transceiver design in the relay-assisted FL system. The strong coupling between the design variables renders the optimization of such a system challenging. To tackle this issue, we propose an alternating-optimization-based algorithm to optimize the transceiver and relay operation. Numerical results show that our design achieves faster convergence compared with state-of-the-art schemes.