Beamforming Vector Design and Device Selection in Over-the-Air Federated Learning

In this paper, we consider a beamforming vector design and device selection problem in over-the-air computation (AirComp) for federated learning. Since the learning performance improves as more devices participate in the federated learning aggregation, we formulate a beamforming vector optimization problem that maximizes the number of selected devices under a given target aggregation mean-squared error. This AirComp uplink beamforming problem with device selection is shown to have the same form as the downlink multicast beamforming problem with user selection, which establishes the AirComp-multicasting duality. We design a low-complexity algorithm based on the projected subgradient method that is orders of magnitude faster than conventional semidefinite relaxation-based algorithms and faster than local model training on the devices, which makes it possible to implement the proposed wireless federated learning in real time. Numerical results show that the proposed algorithm provides significant multiple antenna beamforming gains and achieves the performance of the ideal federated learning system with no aggregation errors.