Joint Trajectory and Communication Optimization for UAV-Assisted Over-The-Air Federated Learning

This work examines the use of an unmanned aerial vehicle (UAV) to assist with the aggregation of multiple machine learning models under the federated learning framework. Different from most works that consider static parameter servers, the use of a mobile parameter server improves the efficiency of model aggregation especially when the devices are spread over a wide area. Moreover, the need to train multiple machine learning models also makes the problem more challenging due to the need to further schedule their transmissions based on the locations of the participating devices and their contributions to the training process. We also adopt over-the-air (OTA) aggregation to enable simultaneous transmission of multiple devices associated with the same model and, thus, reduce the communication overhead. Based on the above setting, we propose a joint transceiver and UAV trajectory design that aims to optimize the per-round training progress according to the weighted sum of the derived training error bound for all models. The optimization problem is first solved under a fixed transmission schedule using a block coordinate descent (BCD) method where the receiver coefficients, the transmitter coefficients, and the UAV trajectory are optimized in turn until convergence. Then, we propose a transmission scheduling design on top of the above optimization using a greedy deflation approach. The effectiveness of the proposed designs are demonstrated through numerical simulations.