CoBAAF: Controlled Bayesian Air Aggregation Federated Learning from Heterogeneous Data

Federated learning (FL) is an emerging machine learning paradigm for training models across multiple edge devices holding local data sets, without explicitly exchanging the data. A major challenge in FL is to reduce the bandwidth and energy consumption due to the repeated transmissions of large volumes of data by a large number of users over the wireless channel, and to handle statistical heterogeneity of users data. In this paper we present a novel algorithm, dubbed Controlled Bayesian Air Aggregation Federated-learning (CoBAAF), that handles statistical heterogeneity in noisy networks using a joint design of three main steps in FL: Model distribution, local training, and global aggregation. Specifically, CoBAAF controls the drift in local updates using a correction term, and allows users to transmit their data signal simultaneously over MAC. Second, it adopts a Bayesian approach to average properly the channel output, thus mitigating the effect of the noise and fading induced by the channel. We analyze the convergence of CoBAAF to the loss minimizing model theoretically, showing its ability to achieve a convergence rate similar to that achieved over error-free channels. Extensive simulation results demonstrate the improved convergence of CoBAAF for training in machine learning problems.