Abstract
In this paper, the performance optimization of federated learning (FL), when
deployed over a realistic wireless multiple-input multiple-output (MIMO)
communication system with digital modulation and over-the-air computation
(AirComp) is studied. In particular, an MIMO system is considered in which edge
devices transmit their local FL models (trained using their locally collected
data) to a parameter server (PS) using beamforming to maximize the number of
devices scheduled for transmission. The PS, acting as a central controller,
generates a global FL model using the received local FL models and broadcasts
it back to all devices. Due to the limited bandwidth in a wireless network,
AirComp is adopted to enable efficient wireless data aggregation. However,
fading of wireless channels can produce aggregate distortions in an
AirComp-based FL scheme. To tackle this challenge, we propose a modified
federated averaging (FedAvg) algorithm that combines digital modulation with
AirComp to mitigate wireless fading while ensuring the communication
efficiency. This is achieved by a joint transmit and receive beamforming
design, which is formulated as a optimization problem to dynamically adjust the
beamforming matrices based on current FL model parameters so as to minimize the
transmitting error and ensure the FL performance. To achieve this goal, we
first analytically characterize how the beamforming matrices affect the
performance of the FedAvg in different iterations. Based on this relationship,
an artificial neural network (ANN) is used to estimate the local FL models of
all devices and adjust the beamforming matrices at the PS for future model
transmission. The algorithmic advantages and improved performance of the
proposed methodologies are demonstrated through extensive numerical
experiments.