Joint Antenna Selection and Beamforming for Massive MIMO-Enabled Over-the-Air Federated Learning

Saba Asaad, Hina Tabassum, Chongjun Ouyang, Ping Wang

Abstract: Over-the-air federated learning (OTA-FL) is an emerging technique to reduce the computation and communication overload caused by the orthogonal transmissions of the model updates in conventional federated learning (FL). This reduction is achieved at the expense of introducing aggregation error that can be efficiently suppressed by means of receive beamforming via large array-antennas. This paper studies OTA-FL in massive multiple-input multiple-output (MIMO) systems with limited number of radio frequency (RF)-chains. For this setting, the beamforming for over-the-air model aggregation needs to be addressed jointly with antenna selection. This leads to an NP-hard problem due to its combinatorial nature. We develop three different algorithms to solve the problem. First, we use the penalty dual decomposition (PDD) technique and propose a two-tier algorithm for joint antenna selection and beamforming. The second algorithm interprets the antenna selection task as a sparse recovery problem and invokes the least absolute shrinkage and selection operator (Lasso) algorithm to approximate the sparse solution. The third algorithm invokes the same sparse recovery based interpretation, but employs the low-complexity method of fast iterative soft-thresholding to find a sparse solution. Convergence and complexity analysis is presented for all the algorithms. The numerical investigations depict that the two algorithms based on the sparse recovery interpretation outperform the PDD-based algorithm, when the number of RF-chains at the edge server is much smaller than its array size. However, as the number of RF-chains increases, the PDD-based algorithm outperforms. Our simulations further depict that learning performance with all the antennas being active at the parameter server (PS) can be closely tracked by selecting less than 20% of the antennas at the PS.