Abstract
In this paper, we consider the optimization of federated learning (FL) over a realistic wireless multiple-input multiple-output (MIMO) communication system with digital modulation and over-the-air computation (AirComp). In such a system, MIMO devices transmit their locally trained FL models to a parameter server (PS) using beamforming to maximize the number of devices scheduled for transmission. AirComp enables efficient wireless model aggregation by the PS in bandwidth-limited settings. However, wireless channel fading can produce distortions in AirComp-based FL. To tackle this challenge, we develop a novel aggregation scheme that combines digital modulation with AirComp to mitigate wireless fading while ensuring communication efficiency. We formulate this as a joint transmit-receive beamforming design optimization problem which dynamically adjusts the beamforming matrices to minimize the FL training loss with transmission errors. To solve this problem based on limited information at the PS, we employ an artificial neural network (ANN) to estimate the local FL models of all devices. Then, we derive a closed-form optimal design of the transmit and receive beamforming matrices based on predicted FL models. Numerical evaluations validate the advantages of the proposed methodology in terms of model training performance compared with baselines.